Let image generation extension hosts control output persistence

codex-rs/app-server/src/extensions.rs

@@ -76,7 +76,9 @@ where
     codex_mcp_extension::install(&mut builder);
     codex_mcp_extension::install_executor_plugins(&mut builder, environment_manager);
     codex_web_search_extension::install(&mut builder, auth_manager.clone());
-    codex_image_generation_extension::install(&mut builder, auth_manager);
+    codex_image_generation_extension::install(&mut builder, auth_manager, |config: &Config| {
+        Some(config.codex_home.clone())
+    });
     let skill_providers = codex_skills_extension::SkillProviders::new()
         .with_executor_provider(executor_skill_provider)
         .with_orchestrator_provider(Arc::new(

codex-rs/app-server/tests/suite/v2/imagegen_extension.rs

@@ -134,7 +134,7 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul
         .context("image output should include model-visible path hint")?;
     assert!(
         output_hint.contains(&saved_path.display().to_string()),
-        "output hint should identify the path core saved"
+        "output hint should identify the path the extension saved"
     );
     assert!(
         !requests[1]

codex-rs/core/src/stream_events_utils.rs

@@ -321,7 +321,7 @@ pub(crate) struct HandleOutputCtx {
     pub cancellation_token: CancellationToken,
 }
 
-async fn apply_turn_item_contributors(
+pub(crate) async fn apply_turn_item_contributors(
     sess: &Session,
     turn_store: &ExtensionData,
     item: &mut TurnItem,

codex-rs/core/src/tools/handlers/extension_tools.rs

@@ -16,6 +16,7 @@ use crate::sandboxing::SandboxPermissions;
 use crate::session::session::Session;
 use crate::session::turn_context::TurnContext;
 use crate::stream_events_utils::TurnItemContributorPolicy;
+use crate::stream_events_utils::apply_turn_item_contributors;
 use crate::stream_events_utils::finalize_turn_item;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolPayload;
@@ -71,10 +72,7 @@ struct CoreTurnItemEmitter {
 fn extension_turn_item(item: ExtensionTurnItem) -> TurnItem {
     match item {
         ExtensionTurnItem::WebSearch(item) => TurnItem::WebSearch(item),
-        ExtensionTurnItem::ImageGeneration(mut item) => {
-            item.saved_path = None;
-            TurnItem::ImageGeneration(item)
-        }
+        ExtensionTurnItem::ImageGeneration(item) => TurnItem::ImageGeneration(item),
     }
 }
 
@@ -95,15 +93,31 @@ impl TurnItemEmitter for CoreTurnItemEmitter {
             let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else {
                 return;
             };
-            let mut item = extension_turn_item(item);
-            finalize_turn_item(
-                session.as_ref(),
-                turn.as_ref(),
-                TurnItemContributorPolicy::Run(turn.extension_data.as_ref()),
-                &mut item,
-                turn.collaboration_mode.mode == codex_protocol::config_types::ModeKind::Plan,
-            )
-            .await;
+            let item = match item {
+                ExtensionTurnItem::ImageGeneration(item) => {
+                    let mut item = TurnItem::ImageGeneration(item);
+                    apply_turn_item_contributors(
+                        session.as_ref(),
+                        turn.extension_data.as_ref(),
+                        &mut item,
+                    )
+                    .await;
+                    item
+                }
+                ExtensionTurnItem::WebSearch(item) => {
+                    let mut item = TurnItem::WebSearch(item);
+                    finalize_turn_item(
+                        session.as_ref(),
+                        turn.as_ref(),
+                        TurnItemContributorPolicy::Run(turn.extension_data.as_ref()),
+                        &mut item,
+                        turn.collaboration_mode.mode
+                            == codex_protocol::config_types::ModeKind::Plan,
+                    )
+                    .await;
+                    item
+                }
+            };
             session.emit_turn_item_completed(turn.as_ref(), item).await;
         })
     }
@@ -534,10 +548,11 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn image_generation_publication_is_finalized_by_core() {
-        let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor));
+    async fn image_generation_publication_preserves_extension_saved_path() {
         let (session, turn, rx) = crate::session::tests::make_session_and_context_with_rx().await;
-        let expected_path = crate::stream_events_utils::image_generation_artifact_path(
+        let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor));
+        let expected_path = test_path_buf("/tmp/extension-claimed.png").abs();
+        let default_path = crate::stream_events_utils::image_generation_artifact_path(
             &turn.config.codex_home,
             &session.thread_id.to_string(),
             "call-image",
@@ -598,9 +613,6 @@ mod tests {
                 saved_path: Some(expected_path.clone()),
             }
         );
-        assert_eq!(
-            std::fs::read(&expected_path).expect("generated artifact should be saved"),
-            b"png"
-        );
+        assert!(!default_path.exists());
     }
 }

codex-rs/core/tests/suite/extension_sandbox.rs

@@ -26,6 +26,7 @@ use codex_protocol::request_permissions::PermissionGrantScope;
 use codex_protocol::request_permissions::RequestPermissionProfile;
 use codex_protocol::request_permissions::RequestPermissionsResponse;
 use codex_protocol::user_input::UserInput;
+use codex_utils_absolute_path::AbsolutePathBuf;
 use core_test_support::responses;
 use core_test_support::skip_if_no_network;
 use core_test_support::skip_if_sandbox;
@@ -48,10 +49,13 @@ const TINY_PNG_BYTES: &[u8] = &[
 const TINY_PNG_BASE64: &str = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==";
 const TINY_PNG_DATA_URL: &str = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==";
 
-fn image_generation_extensions(auth: &CodexAuth) -> Arc<ExtensionRegistry<Config>> {
+fn image_generation_extensions(
+    auth: &CodexAuth,
+    resolve_save_root: impl Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync + 'static,
+) -> Arc<ExtensionRegistry<Config>> {
     let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
     let mut extension_builder = ExtensionRegistryBuilder::<Config>::new();
-    install_image_generation_extension(&mut extension_builder, auth_manager);
+    install_image_generation_extension(&mut extension_builder, auth_manager, resolve_save_root);
     Arc::new(extension_builder.build())
 }
 
@@ -61,7 +65,7 @@ async fn extension_tool_receives_turn_environment_sandbox() -> Result<()> {
 
     let server = responses::start_mock_server().await;
     let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
-    let extensions = image_generation_extensions(&auth);
+    let extensions = image_generation_extensions(&auth, |config| Some(config.codex_home.clone()));
     let mut builder = test_codex()
         .with_auth(auth)
         .with_extensions(extensions)
@@ -141,7 +145,7 @@ async fn extension_tool_receives_turn_environment_sandbox() -> Result<()> {
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
-async fn extension_tool_uses_granted_turn_permissions() -> Result<()> {
+async fn extension_tool_uses_granted_turn_permissions_without_local_persistence() -> Result<()> {
     skip_if_no_network!(Ok(()));
     skip_if_sandbox!(Ok(()));
 
@@ -157,7 +161,7 @@ async fn extension_tool_uses_granted_turn_permissions() -> Result<()> {
         .await;
 
     let auth = CodexAuth::create_dummy_chatgpt_auth_for_testing();
-    let extensions = image_generation_extensions(&auth);
+    let extensions = image_generation_extensions(&auth, |_config| None);
     let base_permission_profile = PermissionProfile::workspace_write_with(
         &[],
         NetworkSandboxPolicy::Restricted,
@@ -299,9 +303,14 @@ async fn extension_tool_uses_granted_turn_permissions() -> Result<()> {
         .last_request()
         .context("missing request containing extension output")?;
     let output = request.function_call_output(image_call_id);
-    let image = &output["output"][0];
-    assert_eq!(image["type"], "input_image");
-    assert_eq!(image["image_url"], TINY_PNG_DATA_URL);
+    assert_eq!(
+        output["output"],
+        json!([{
+            "type": "input_image",
+            "image_url": TINY_PNG_DATA_URL,
+        }])
+    );
+    assert!(!test.config.codex_home.join("generated_images").exists());
 
     Ok(())
 }

codex-rs/core/tests/suite/responses_lite.rs

@@ -28,7 +28,9 @@ fn responses_extensions(auth: &CodexAuth) -> Arc<ExtensionRegistry<Config>> {
     let auth_manager = codex_core::test_support::auth_manager_from_auth(auth.clone());
     let mut extension_builder = ExtensionRegistryBuilder::<Config>::new();
     install_web_search_extension(&mut extension_builder, Arc::clone(&auth_manager));
-    install_image_generation_extension(&mut extension_builder, auth_manager);
+    install_image_generation_extension(&mut extension_builder, auth_manager, |config| {
+        Some(config.codex_home.clone())
+    });
     Arc::new(extension_builder.build())
 }
 

codex-rs/ext/image-generation/Cargo.toml

@@ -13,8 +13,10 @@ doctest = false
 workspace = true
 
 [dependencies]
+base64 = { workspace = true }
 codex-api = { workspace = true }
 codex-core = { workspace = true }
+codex-exec-server = { workspace = true }
 codex-extension-api = { workspace = true }
 codex-login = { workspace = true }
 codex-model-provider = { workspace = true }
@@ -28,6 +30,7 @@ http = { workspace = true }
 schemars = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
+tracing = { workspace = true }
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }

codex-rs/ext/image-generation/src/extension.rs

@@ -21,23 +21,26 @@ use crate::tool::ImageGenerationTool;
 #[derive(Clone)]
 struct ImageGenerationExtension {
     auth_manager: Arc<AuthManager>,
+    resolve_save_root: Arc<SaveRootResolver>,
 }
 
+type SaveRootResolver = dyn Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync;
+
 #[derive(Clone)]
 struct ImageGenerationExtensionConfig {
     available: bool,
     provider: ModelProviderInfo,
-    codex_home: AbsolutePathBuf,
+    save_root: Option<AbsolutePathBuf>,
 }
 
-impl From<&Config> for ImageGenerationExtensionConfig {
+impl ImageGenerationExtensionConfig {
     /// Resolves whether standalone image generation should be available for a thread.
-    fn from(config: &Config) -> Self {
+    fn from_config(config: &Config, resolve_save_root: &SaveRootResolver) -> Self {
         Self {
             // Core selects this executor per turn using the feature flag or model metadata.
             available: config.model_provider.is_openai(),
             provider: config.model_provider.clone(),
-            codex_home: config.codex_home.clone(),
+            save_root: resolve_save_root(config),
         }
     }
 }
@@ -51,7 +54,10 @@ impl ThreadLifecycleContributor<Config> for ImageGenerationExtension {
         Box::pin(async move {
             input
                 .thread_store
-                .insert(ImageGenerationExtensionConfig::from(input.config));
+                .insert(ImageGenerationExtensionConfig::from_config(
+                    input.config,
+                    self.resolve_save_root.as_ref(),
+                ));
         })
     }
 }
@@ -65,7 +71,10 @@ impl ConfigContributor<Config> for ImageGenerationExtension {
         _previous_config: &Config,
         new_config: &Config,
     ) {
-        thread_store.insert(ImageGenerationExtensionConfig::from(new_config));
+        thread_store.insert(ImageGenerationExtensionConfig::from_config(
+            new_config,
+            self.resolve_save_root.as_ref(),
+        ));
     }
 }
 
@@ -88,15 +97,22 @@ impl ToolContributor for ImageGenerationExtension {
                 config.provider.clone(),
                 Some(self.auth_manager.clone()),
             )),
-            config.codex_home.clone(),
+            config.save_root.clone(),
             thread_store.level_id().to_string(),
         ))]
     }
 }
 
 /// Installs the standalone image-generation extension contributors.
-pub fn install(registry: &mut ExtensionRegistryBuilder<Config>, auth_manager: Arc<AuthManager>) {
-    let extension = Arc::new(ImageGenerationExtension { auth_manager });
+pub fn install(
+    registry: &mut ExtensionRegistryBuilder<Config>,
+    auth_manager: Arc<AuthManager>,
+    resolve_save_root: impl Fn(&Config) -> Option<AbsolutePathBuf> + Send + Sync + 'static,
+) {
+    let extension = Arc::new(ImageGenerationExtension {
+        auth_manager,
+        resolve_save_root: Arc::new(resolve_save_root),
+    });
     registry.thread_lifecycle_contributor(extension.clone());
     registry.config_contributor(extension.clone());
     registry.tool_contributor(extension);