Route extension image generation through the native image completion pipeline

codex-rs/core/src/stream_events_utils.rs

@@ -5,6 +5,7 @@ use base64::Engine;
 use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
 use codex_extension_api::ExtensionData;
 use codex_protocol::config_types::ModeKind;
+use codex_protocol::items::ImageGenerationItem;
 use codex_protocol::items::TurnItem;
 use codex_utils_stream_parser::strip_citations;
 use tokio_util::sync::CancellationToken;
@@ -125,6 +126,68 @@ async fn save_image_generation_result(
     Ok(path)
 }
 
+pub(crate) async fn persist_image_generation_item(
+    sess: &Session,
+    turn_context: &TurnContext,
+    image_item: &mut ImageGenerationItem,
+) -> Option<AbsolutePathBuf> {
+    let session_id = sess.conversation_id.to_string();
+    match save_image_generation_result(
+        &turn_context.config.codex_home,
+        &session_id,
+        &image_item.id,
+        &image_item.result,
+    )
+    .await
+    {
+        Ok(path) => {
+            image_item.saved_path = Some(path.clone());
+            Some(path)
+        }
+        Err(err) => {
+            let output_path = image_generation_artifact_path(
+                &turn_context.config.codex_home,
+                &session_id,
+                &image_item.id,
+            );
+            let output_dir = output_path
+                .parent()
+                .unwrap_or_else(|| turn_context.config.codex_home.clone());
+            tracing::warn!(
+                call_id = %image_item.id,
+                output_dir = %output_dir.display(),
+                "failed to save generated image: {err}"
+            );
+            None
+        }
+    }
+}
+
+pub(crate) async fn finalize_image_generation_item(
+    sess: &Session,
+    turn_context: &TurnContext,
+    image_item: &mut ImageGenerationItem,
+) {
+    if persist_image_generation_item(sess, turn_context, image_item)
+        .await
+        .is_none()
+    {
+        return;
+    }
+    let session_id = sess.conversation_id.to_string();
+    let image_output_path =
+        image_generation_artifact_path(&turn_context.config.codex_home, &session_id, "<image_id>");
+    let image_output_dir = image_output_path
+        .parent()
+        .unwrap_or_else(|| turn_context.config.codex_home.clone());
+    let message: ResponseItem = ContextualUserFragment::into(ImageGenerationInstructions::new(
+        image_output_dir.display(),
+        image_output_path.display(),
+    ));
+    sess.record_conversation_items(turn_context, &[message])
+        .await;
+}
+
 /// Persist a completed model response item and record any cited memory usage.
 pub(crate) async fn record_completed_response_item(
     sess: &Session,
@@ -487,49 +550,7 @@ pub(crate) async fn handle_non_tool_response_item(
                 }
             }
             if let TurnItem::ImageGeneration(image_item) = &mut turn_item {
-                let session_id = sess.conversation_id.to_string();
-                match save_image_generation_result(
-                    &turn_context.config.codex_home,
-                    &session_id,
-                    &image_item.id,
-                    &image_item.result,
-                )
-                .await
-                {
-                    Ok(path) => {
-                        image_item.saved_path = Some(path);
-                        let image_output_path = image_generation_artifact_path(
-                            &turn_context.config.codex_home,
-                            &session_id,
-                            "<image_id>",
-                        );
-                        let image_output_dir = image_output_path
-                            .parent()
-                            .unwrap_or_else(|| turn_context.config.codex_home.clone());
-                        let message: ResponseItem =
-                            ContextualUserFragment::into(ImageGenerationInstructions::new(
-                                image_output_dir.display(),
-                                image_output_path.display(),
-                            ));
-                        sess.record_conversation_items(turn_context, &[message])
-                            .await;
-                    }
-                    Err(err) => {
-                        let output_path = image_generation_artifact_path(
-                            &turn_context.config.codex_home,
-                            &session_id,
-                            &image_item.id,
-                        );
-                        let output_dir = output_path
-                            .parent()
-                            .unwrap_or_else(|| turn_context.config.codex_home.clone());
-                        tracing::warn!(
-                            call_id = %image_item.id,
-                            output_dir = %output_dir.display(),
-                            "failed to save generated image: {err}"
-                        );
-                    }
-                }
+                finalize_image_generation_item(sess, turn_context, image_item).await;
             }
             Some(turn_item)
         }

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

@@ -4,15 +4,19 @@ use std::sync::Weak;
 use codex_protocol::items::TurnItem;
 use codex_tools::ConversationHistory;
 use codex_tools::ExtensionTurnItem;
+use codex_tools::ImageGenerationCompletionFuture;
 use codex_tools::ToolCall as ExtensionToolCall;
 use codex_tools::ToolName;
 use codex_tools::ToolSpec;
 use codex_tools::TurnItemEmissionFuture;
 use codex_tools::TurnItemEmitter;
 
+use crate::context::ContextualUserFragment;
+use crate::context::ImageGenerationInstructions;
 use crate::function_tool::FunctionCallError;
 use crate::session::session::Session;
 use crate::session::turn_context::TurnContext;
+use crate::stream_events_utils::persist_image_generation_item;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolOutput;
 use crate::tools::context::ToolPayload;
@@ -90,6 +94,50 @@ impl TurnItemEmitter for CoreTurnItemEmitter {
             session.emit_turn_item_completed(turn.as_ref(), item).await;
         })
     }
+
+    fn image_generation_completed<'a>(
+        &'a self,
+        call_id: String,
+        prompt: String,
+        result: String,
+    ) -> ImageGenerationCompletionFuture<'a> {
+        Box::pin(async move {
+            let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else {
+                return None;
+            };
+            let mut item = codex_protocol::items::ImageGenerationItem {
+                id: call_id,
+                status: "completed".to_string(),
+                revised_prompt: Some(prompt),
+                result,
+                saved_path: None,
+            };
+            let output_hint =
+                persist_image_generation_item(session.as_ref(), turn.as_ref(), &mut item)
+                    .await
+                    .map(|saved_path| {
+                        let output_dir = saved_path
+                            .parent()
+                            .unwrap_or_else(|| turn.config.codex_home.clone());
+                        ImageGenerationInstructions::new(output_dir.display(), saved_path.display())
+                            .body()
+                    });
+            let started_item = codex_protocol::items::ImageGenerationItem {
+                id: item.id.clone(),
+                status: "in_progress".to_string(),
+                revised_prompt: None,
+                result: String::new(),
+                saved_path: None,
+            };
+            session
+                .emit_turn_item_started(turn.as_ref(), &TurnItem::ImageGeneration(started_item))
+                .await;
+            session
+                .emit_turn_item_completed(turn.as_ref(), TurnItem::ImageGeneration(item))
+                .await;
+            output_hint
+        })
+    }
 }
 
 async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall {
@@ -352,4 +400,130 @@ mod tests {
         assert_eq!(end.query, expected.query);
         assert_eq!(end.action, expected.action);
     }
+
+    struct ImageGenerationExtensionExecutor {
+        output_hint: Arc<Mutex<Option<String>>>,
+    }
+
+    #[async_trait::async_trait]
+    impl codex_extension_api::ToolExecutor<codex_tools::ToolCall> for ImageGenerationExtensionExecutor {
+        fn tool_name(&self) -> codex_tools::ToolName {
+            codex_tools::ToolName::namespaced("image_gen", "imagegen")
+        }
+
+        fn spec(&self) -> codex_tools::ToolSpec {
+            codex_tools::ToolSpec::Function(codex_tools::ResponsesApiTool {
+                name: "imagegen".to_string(),
+                description: "Generates an image.".to_string(),
+                strict: false,
+                parameters: codex_tools::JsonSchema::default(),
+                output_schema: None,
+                defer_loading: None,
+            })
+        }
+
+        async fn handle(
+            &self,
+            call: codex_tools::ToolCall,
+        ) -> Result<Box<dyn codex_tools::ToolOutput>, codex_tools::FunctionCallError> {
+            let output_hint = call
+                .turn_item_emitter
+                .image_generation_completed(
+                    call.call_id,
+                    "A tiny blue square".to_string(),
+                    "cG5n".to_string(),
+                )
+                .await;
+            *self.output_hint.lock().await = output_hint;
+            Ok(Box::new(codex_tools::JsonToolOutput::new(
+                json!({ "ok": true }),
+            )))
+        }
+    }
+
+    #[tokio::test]
+    async fn image_generation_publication_is_finalized_by_core() {
+        let output_hint = Arc::new(Mutex::new(None));
+        let handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor {
+            output_hint: Arc::clone(&output_hint),
+        }));
+        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(
+            &turn.config.codex_home,
+            &session.conversation_id.to_string(),
+            "call-image",
+        );
+        let invocation = ToolInvocation {
+            session,
+            turn,
+            cancellation_token: tokio_util::sync::CancellationToken::new(),
+            tracker: Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::new())),
+            call_id: "call-image".to_string(),
+            tool_name: codex_tools::ToolName::namespaced("image_gen", "imagegen"),
+            source: ToolCallSource::Direct,
+            payload: ToolPayload::Function {
+                arguments: "{}".to_string(),
+            },
+        };
+
+        crate::tools::registry::ToolExecutor::handle(&handler, invocation)
+            .await
+            .expect("extension call should succeed");
+
+        let started = rx.recv().await.expect("item started event");
+        let EventMsg::ItemStarted(started) = started.msg else {
+            panic!("expected item started event");
+        };
+        let TurnItem::ImageGeneration(started_item) = started.item else {
+            panic!("expected image generation item");
+        };
+        let begin = rx.recv().await.expect("legacy image start event");
+        assert!(matches!(begin.msg, EventMsg::ImageGenerationBegin(_)));
+        let completed = rx.recv().await.expect("item completed event");
+        let EventMsg::ItemCompleted(completed) = completed.msg else {
+            panic!("expected item completed event");
+        };
+        let TurnItem::ImageGeneration(completed_item) = completed.item else {
+            panic!("expected image generation item");
+        };
+        let end = rx.recv().await.expect("legacy image end event");
+        assert!(matches!(end.msg, EventMsg::ImageGenerationEnd(_)));
+
+        assert_eq!(
+            started_item,
+            codex_protocol::items::ImageGenerationItem {
+                id: "call-image".to_string(),
+                status: "in_progress".to_string(),
+                revised_prompt: None,
+                result: String::new(),
+                saved_path: None,
+            }
+        );
+        assert_eq!(
+            completed_item,
+            codex_protocol::items::ImageGenerationItem {
+                id: "call-image".to_string(),
+                status: "completed".to_string(),
+                revised_prompt: Some("A tiny blue square".to_string()),
+                result: "cG5n".to_string(),
+                saved_path: Some(expected_path.clone()),
+            }
+        );
+        assert_eq!(
+            std::fs::read(&expected_path).expect("generated artifact should be saved"),
+            b"png"
+        );
+        assert_eq!(
+            *output_hint.lock().await,
+            Some(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.",
+                expected_path
+                    .parent()
+                    .expect("generated image path should have a parent")
+                    .display(),
+                expected_path.display(),
+            ))
+        );
+    }
 }

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"] }