Route extension image generation through the native image completion pipeline (#24972)

## Why

The standalone `image_gen.imagegen` extension should behave like native
image generation for artifact persistence and UI completion, while
returning its save-location guidance as part of the tool result instead
of injecting a developer message.

## What Changed

- Added an image-generation completion hook for extension tools so core
can persist generated images and emit the existing `ImageGeneration`
lifecycle events.
- Reused core image artifact persistence for extension output and
removed extension-local save-path/file-writing logic.
- Split shared image persistence from built-in finalization so native
image generation keeps its existing developer-message instruction
behavior.
- Returned the generated image save-location instruction through the
extension `FunctionCallOutput`, alongside the generated image input for
model follow-up.
- Preserved the existing image-generation event shape for current UI and
replay compatibility.
- Avoided cloning the full generated-image base64 payload when emitting
the in-progress image item.
- Removed dependencies no longer needed after moving persistence out of
the extension crate.

## Fast Follow
- Adjust the existing Extension API and add a general `TurnItem`
finalization path for re-usability of code

## Validation

- Ran `just fmt`.
- Ran `just bazel-lock-update`.
- Ran `just bazel-lock-check`.
- Ran `just test -p codex-tools -p codex-extension-api -p
codex-image-generation-extension`.
- Ran `just test -p codex-core
image_generation_publication_is_finalized_by_core`.
- Ran `just test -p codex-core
handle_output_item_done_records_image_save_history_message`.
- Ran `just fix -p codex-tools -p codex-extension-api -p codex-core -p
codex-image-generation-extension`.

codex-rs/ext/extension-api/src/lib.rs

@@ -13,6 +13,7 @@ pub use capabilities::ResponseItemInjector;
 pub use codex_tools::ConversationHistory;
 pub use codex_tools::ExtensionTurnItem;
 pub use codex_tools::FunctionCallError;
+pub use codex_tools::ImageGenerationCompletionFuture;
 pub use codex_tools::JsonToolOutput;
 pub use codex_tools::NoopTurnItemEmitter;
 pub use codex_tools::ResponsesApiTool;

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

@@ -14,7 +14,6 @@ workspace = true
 
 [dependencies]
 async-trait = { workspace = true }
-base64 = { workspace = true }
 codex-api = { workspace = true }
 codex-core = { workspace = true }
 codex-extension-api = { workspace = true }
@@ -28,10 +27,6 @@ http = { workspace = true }
 schemars = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 serde_json = { workspace = true }
-tokio = { workspace = true, features = ["fs"] }
-tracing = { workspace = true }
 
 [dev-dependencies]
 pretty_assertions = { workspace = true }
-tempfile = { workspace = true }
-tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

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

@@ -1,4 +1,3 @@
-use std::path::PathBuf;
 use std::sync::Arc;
 
 use codex_core::config::Config;
@@ -17,7 +16,6 @@ use codex_model_provider_info::ModelProviderInfo;
 
 use crate::backend::CodexImagesBackend;
 use crate::tool::ImageGenerationTool;
-use crate::tool::generated_image_output_dir;
 
 #[derive(Clone)]
 struct ImageGenerationExtension {
@@ -28,7 +26,6 @@ struct ImageGenerationExtension {
 struct ImageGenerationExtensionConfig {
     enabled: bool,
     provider: ModelProviderInfo,
-    codex_home: PathBuf,
 }
 
 impl From<&Config> for ImageGenerationExtensionConfig {
@@ -38,7 +35,6 @@ impl From<&Config> for ImageGenerationExtensionConfig {
             enabled: config.features.enabled(Feature::ImageGenExt)
                 && config.model_provider.is_openai(),
             provider: config.model_provider.clone(),
-            codex_home: config.codex_home.to_path_buf(),
         }
     }
 }
@@ -80,13 +76,9 @@ impl ToolContributor for ImageGenerationExtension {
             return Vec::new();
         }
 
-        vec![Arc::new(ImageGenerationTool::new(
-            CodexImagesBackend::new(create_model_provider(
-                config.provider.clone(),
-                Some(self.auth_manager.clone()),
-            )),
-            generated_image_output_dir(&config.codex_home, thread_store.level_id()),
-        ))]
+        vec![Arc::new(ImageGenerationTool::new(CodexImagesBackend::new(
+            create_model_provider(config.provider.clone(), Some(self.auth_manager.clone())),
+        )))]
     }
 }
 

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

@@ -20,14 +20,13 @@ use super::GeneratedImageOutput;
 use super::ImageRequest;
 use super::ImagegenAction;
 use super::ImagegenArgs;
-use super::generated_image_output_dir;
 use super::imagegen_tool_spec;
-use super::persist_generated_image;
 use super::request_for_action;
 use crate::IMAGE_GEN_NAMESPACE;
 use crate::IMAGEGEN_TOOL_NAME;
 
 const RESULT: &str = "cG5n";
+const OUTPUT_HINT: &str = "Generated images are saved to /tmp as /tmp/call-1.png by default.";
 
 #[test]
 fn uses_reserved_image_gen_namespace() {
@@ -55,15 +54,11 @@ fn generate_uses_fixed_request_defaults() {
     );
 }
 
-#[tokio::test]
-async fn generated_output_returns_image_input_and_persists_artifact() {
-    let tempdir = tempfile::tempdir().expect("tempdir");
-    let output_hint = persist_generated_image(tempdir.path(), "call-1", RESULT)
-        .await
-        .expect("generated image should persist");
+#[test]
+fn generated_output_returns_image_input_and_output_hint() {
     let output = GeneratedImageOutput {
         result: RESULT.to_string(),
-        output_hint: Some(output_hint),
+        output_hint: Some(OUTPUT_HINT.to_string()),
     };
 
     let ResponseInputItem::FunctionCallOutput {
@@ -84,19 +79,10 @@ async fn generated_output_returns_image_input_and_persists_artifact() {
                 detail: Some(DEFAULT_IMAGE_DETAIL),
             },
             FunctionCallOutputContentItem::InputText {
-                text: format!(
-                    "Generated images are saved to {} as {} by default.\n\
-                     If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it.",
-                    tempdir.path().display(),
-                    tempdir.path().join("call-1.png").display(),
-                ),
+                text: OUTPUT_HINT.to_string(),
             },
         ]
     );
-    assert_eq!(
-        std::fs::read(tempdir.path().join("call-1.png")).expect("saved generated image"),
-        b"png"
-    );
 }
 
 #[test]
@@ -265,14 +251,6 @@ fn edit_without_image_history_returns_tool_error() {
     );
 }
 
-#[test]
-fn generated_image_output_dir_is_scoped_to_sanitized_thread_id() {
-    assert_eq!(
-        generated_image_output_dir(std::path::Path::new("/tmp/codex-home"), "thread/1"),
-        std::path::PathBuf::from("/tmp/codex-home/generated_images/thread_1")
-    );
-}
-
 fn args(action: ImagegenAction, prompt: &str) -> ImagegenArgs {
     ImagegenArgs {
         prompt: prompt.to_string(),

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

@@ -1,8 +1,3 @@
-use std::path::Path;
-use std::path::PathBuf;
-
-use base64::Engine;
-use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use codex_api::ImageBackground;
 use codex_api::ImageEditRequest;
 use codex_api::ImageGenerationRequest;
@@ -41,21 +36,16 @@ use crate::backend::CodexImagesBackend;
 const IMAGE_MODEL: &str = "gpt-image-2";
 const MAX_EDIT_IMAGES: usize = 5;
 const IMAGEGEN_DESCRIPTION: &str = include_str!("../imagegen_description.md");
-const GENERATED_IMAGE_ARTIFACTS_DIR: &str = "generated_images";
 
 #[derive(Clone)]
 pub(crate) struct ImageGenerationTool {
     backend: CodexImagesBackend,
-    output_dir: PathBuf,
 }
 
 impl ImageGenerationTool {
     /// Creates an image-generation tool backed by an image API executor.
-    pub(crate) fn new(backend: CodexImagesBackend, output_dir: PathBuf) -> Self {
-        Self {
-            backend,
-            output_dir,
-        }
+    pub(crate) fn new(backend: CodexImagesBackend) -> Self {
+        Self { backend }
     }
 }
 
@@ -94,7 +84,6 @@ impl ToolExecutor<ToolCall> for ImageGenerationTool {
     async fn handle(&self, call: ToolCall) -> Result<Box<dyn ToolOutput>, FunctionCallError> {
         let args = parse_args(&call)?;
         let request = request_for_action(&args, call.conversation_history.items())?;
-
         let response = match request {
             ImageRequest::Generate(request) => self.backend.generate(request).await,
             ImageRequest::Edit(request) => self.backend.edit(request).await,
@@ -107,18 +96,10 @@ impl ToolExecutor<ToolCall> for ImageGenerationTool {
                 "image generation returned no image data".to_string(),
             ));
         };
-        let output_hint =
-            match persist_generated_image(&self.output_dir, &call.call_id, &result).await {
-                Ok(output_hint) => Some(output_hint),
-                Err(err) => {
-                    tracing::warn!(
-                        call_id = %call.call_id,
-                        output_dir = %self.output_dir.display(),
-                        "failed to save generated image: {err}"
-                    );
-                    None
-                }
-            };
+        let output_hint = call
+            .turn_item_emitter
+            .image_generation_completed(call.call_id.clone(), args.prompt, result.clone())
+            .await;
         Ok(Box::new(GeneratedImageOutput {
             result,
             output_hint,
@@ -268,58 +249,6 @@ fn parse_args(call: &ToolCall) -> Result<ImagegenArgs, FunctionCallError> {
         .map_err(|err| FunctionCallError::RespondToModel(err.to_string()))
 }
 
-/// Resolves where generated images for one thread are persisted by the extension.
-pub(crate) fn generated_image_output_dir(codex_home: &Path, thread_id: &str) -> PathBuf {
-    codex_home
-        .join(GENERATED_IMAGE_ARTIFACTS_DIR)
-        .join(sanitize_path_component(thread_id))
-}
-
-fn generated_image_output_path(output_dir: &Path, call_id: &str) -> PathBuf {
-    output_dir.join(format!("{}.png", sanitize_path_component(call_id)))
-}
-
-fn sanitize_path_component(value: &str) -> String {
-    let sanitized: String = value
-        .chars()
-        .map(|ch| {
-            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
-                ch
-            } else {
-                '_'
-            }
-        })
-        .collect();
-    if sanitized.is_empty() {
-        "generated_image".to_string()
-    } else {
-        sanitized
-    }
-}
-
-async fn persist_generated_image(
-    output_dir: &Path,
-    call_id: &str,
-    result: &str,
-) -> Result<String, String> {
-    let bytes = BASE64_STANDARD
-        .decode(result.trim().as_bytes())
-        .map_err(|err| format!("invalid image generation payload: {err}"))?;
-    tokio::fs::create_dir_all(output_dir)
-        .await
-        .map_err(|err| err.to_string())?;
-    tokio::fs::write(generated_image_output_path(output_dir, call_id), bytes)
-        .await
-        .map_err(|err| err.to_string())?;
-
-    Ok(format!(
-        "Generated images are saved to {} as {} by default.\n\
-         If you need to use a generated image at another path, copy it and leave the original in place unless the user explicitly asks you to delete it.",
-        output_dir.display(),
-        generated_image_output_path(output_dir, call_id).display(),
-    ))
-}
-
 /// Builds the namespace function schema exposed to the model.
 fn imagegen_tool_spec() -> ToolSpec {
     let mut schema_value = serde_json::to_value(
@@ -369,7 +298,7 @@ impl ToolOutput for GeneratedImageOutput {
         true
     }
 
-    /// Returns generated bytes and persisted-artifact context for the model's follow-up response.
+    /// Returns generated bytes and persisted-artifact context for model follow-up.
     fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem {
         let mut content = vec![FunctionCallOutputContentItem::InputImage {
             image_url: format!("data:image/png;base64,{}", self.result),