Route standalone image generation through host finalization md (#25176)

## Why Standalone image-generation extensions emitted turn items through the low-level event path, bypassing host-owned finalization such as image persistence and contributor processing. At the same time, the generated-image save-path hint must remain visible to the model through the extension tool's `FunctionCallOutput`, rather than the legacy built-in developer-message path. ## What changed - Extended `ExtensionTurnItem` to support image-generation items while keeping the extension-facing emitter API limited to `emit_started` and `emit_completed`. - Routed extension completion through core `finalize_turn_item`, so standalone image-generation items receive host-owned processing and persisted `saved_path` values before publication. - Kept legacy built-in image generation on its existing developer-message hint path, while standalone image generation returns its deterministic saved-path hint in `FunctionCallOutput`. - Shared the image artifact path and output-hint formatting used by core and the image-generation extension. - Passed thread identity through extension tool calls so standalone image generation can construct the same intended artifact path as core. - Added an app-server integration test covering real standalone image generation, saved artifact publication, model-visible output hint wiring, and absence of the legacy developer-message hint. ## Validation - `just fmt` - `just test -p codex-image-generation-extension` - `just test -p codex-web-search-extension` - `just test -p codex-goal-extension` - `just test -p codex-memories-extension` - Targeted `codex-core` tests for image save history, extension completion finalization, and contributor execution - `just test -p codex-app-server standalone_image_generation_returns_saved_path_hint_to_model` - `just fix -p codex-core` - `just fix -p codex-image-generation-extension` - `just bazel-lock-update` - `just bazel-lock-check`
2026-07-01 00:31:56 +08:00 · 2026-06-02 12:00:04 -07:00
parent e29071e4c9
commit 57f337a8e9
9 changed files with 385 additions and 156 deletions
@@ -0,0 +1,214 @@
+use std::path::Path;
+use std::time::Duration;
+
+use anyhow::Context;
+use anyhow::Result;
+use app_test_support::ChatGptAuthFixture;
+use app_test_support::McpProcess;
+use app_test_support::to_response;
+use app_test_support::write_chatgpt_auth;
+use codex_app_server_protocol::ItemCompletedNotification;
+use codex_app_server_protocol::JSONRPCResponse;
+use codex_app_server_protocol::RequestId;
+use codex_app_server_protocol::ThreadItem;
+use codex_app_server_protocol::ThreadStartParams;
+use codex_app_server_protocol::ThreadStartResponse;
+use codex_app_server_protocol::TurnStartParams;
+use codex_app_server_protocol::TurnStartResponse;
+use codex_app_server_protocol::UserInput as V2UserInput;
+use codex_config::types::AuthCredentialsStoreMode;
+use core_test_support::responses;
+use pretty_assertions::assert_eq;
+use serde_json::json;
+use tempfile::TempDir;
+use tokio::time::timeout;
+use wiremock::Mock;
+use wiremock::MockServer;
+use wiremock::ResponseTemplate;
+use wiremock::matchers::method;
+use wiremock::matchers::path;
+
+const RESULT: &str = "cG5n";
+
+// macOS and Windows Bazel CI can spend tens of seconds starting app-server
+// subprocesses or processing test RPCs under load.
+#[cfg(any(target_os = "macos", windows))]
+const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(60);
+#[cfg(not(any(target_os = "macos", windows)))]
+const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(10);
+
+#[tokio::test]
+async fn standalone_image_generation_persists_image_and_returns_it_to_model() -> Result<()> {
+    let call_id = "image-run-1";
+    let server = responses::start_mock_server().await;
+    mount_image_response(&server).await;
+
+    let response_mock = responses::mount_sse_sequence(
+        &server,
+        vec![
+            responses::sse(vec![
+                responses::ev_response_created("resp-1"),
+                responses::ev_function_call_with_namespace(
+                    call_id,
+                    "image_gen",
+                    "imagegen",
+                    &json!({
+                        "action": "generate",
+                        "prompt": "paint a blue whale",
+                    })
+                    .to_string(),
+                ),
+                responses::ev_completed("resp-1"),
+            ]),
+            responses::sse(vec![
+                responses::ev_assistant_message("msg-1", "Done"),
+                responses::ev_completed("resp-2"),
+            ]),
+        ],
+    )
+    .await;
+
+    let codex_home = TempDir::new()?;
+    create_config_toml(codex_home.path(), &server.uri())?;
+    write_chatgpt_auth(
+        codex_home.path(),
+        ChatGptAuthFixture::new("access-chatgpt"),
+        AuthCredentialsStoreMode::File,
+    )?;
+
+    let mut mcp = McpProcess::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?;
+    timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??;
+
+    let thread_req = mcp
+        .send_thread_start_request(ThreadStartParams::default())
+        .await?;
+    let thread_resp: JSONRPCResponse = timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_response_message(RequestId::Integer(thread_req)),
+    )
+    .await??;
+    let ThreadStartResponse { thread, .. } = to_response::<ThreadStartResponse>(thread_resp)?;
+
+    let turn_req = mcp
+        .send_turn_start_request(TurnStartParams {
+            thread_id: thread.id,
+            client_user_message_id: None,
+            input: vec![V2UserInput::Text {
+                text: "Generate an image".to_string(),
+                text_elements: Vec::new(),
+            }],
+            ..Default::default()
+        })
+        .await?;
+    let turn_resp: JSONRPCResponse = timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_response_message(RequestId::Integer(turn_req)),
+    )
+    .await??;
+    let _turn: TurnStartResponse = to_response::<TurnStartResponse>(turn_resp)?;
+
+    let completed = timeout(
+        DEFAULT_READ_TIMEOUT,
+        wait_for_image_generation_completed(&mut mcp),
+    )
+    .await??;
+    timeout(
+        DEFAULT_READ_TIMEOUT,
+        mcp.read_stream_until_notification_message("turn/completed"),
+    )
+    .await??;
+
+    let ThreadItem::ImageGeneration {
+        status,
+        revised_prompt,
+        result,
+        saved_path: Some(saved_path),
+        ..
+    } = completed.item
+    else {
+        panic!("expected completed image generation item with saved path");
+    };
+    assert_eq!(status, "completed");
+    assert_eq!(revised_prompt.as_deref(), Some("paint a blue whale"));
+    assert_eq!(result, RESULT);
+    assert_eq!(std::fs::read(&saved_path)?, b"png");
+
+    let requests = response_mock.requests();
+    assert_eq!(requests.len(), 2);
+    let output = requests[1].function_call_output(call_id);
+    assert_eq!(
+        output["output"][0],
+        json!({
+            "type": "input_image",
+            "image_url": format!("data:image/png;base64,{RESULT}"),
+            "detail": "high",
+        })
+    );
+    assert_eq!(output["output"].as_array().map(Vec::len), Some(1));
+    assert!(
+        !requests[1]
+            .message_input_texts("developer")
+            .iter()
+            .any(|text| text.contains("Generated images are saved to")),
+        "standalone image generation should not emit the legacy developer-message hint"
+    );
+
+    Ok(())
+}
+
+async fn wait_for_image_generation_completed(
+    mcp: &mut McpProcess,
+) -> Result<ItemCompletedNotification> {
+    loop {
+        let notification = mcp
+            .read_stream_until_notification_message("item/completed")
+            .await?;
+        let completed: ItemCompletedNotification = serde_json::from_value(
+            notification
+                .params
+                .context("item/completed notification should include params")?,
+        )?;
+        if matches!(&completed.item, ThreadItem::ImageGeneration { .. }) {
+            return Ok(completed);
+        }
+    }
+}
+
+async fn mount_image_response(server: &MockServer) {
+    Mock::given(method("POST"))
+        .and(path("/api/codex/images/generations"))
+        .respond_with(ResponseTemplate::new(200).set_body_json(json!({
+            "created": 1,
+            "data": [{"b64_json": RESULT}],
+        })))
+        .expect(1)
+        .mount(server)
+        .await;
+}
+
+fn create_config_toml(codex_home: &Path, server_uri: &str) -> std::io::Result<()> {
+    std::fs::write(
+        codex_home.join("config.toml"),
+        format!(
+            r#"
+model = "mock-model"
+approval_policy = "never"
+sandbox_mode = "read-only"
+model_provider = "openai-custom"
+chatgpt_base_url = "{server_uri}"
+
+[features]
+imagegenext = true
+
+[model_providers.openai-custom]
+name = "OpenAI"
+base_url = "{server_uri}/api/codex"
+wire_api = "responses"
+request_max_retries = 0
+stream_max_retries = 0
+supports_websockets = false
+requires_openai_auth = true
+"#
+        ),
+    )
+}
@@ -17,6 +17,7 @@ mod experimental_feature_list;
 mod external_agent_config;
 mod fs;
 mod hooks_list;
+mod image_generation;
 mod initialize;
 mod marketplace_add;
 mod marketplace_remove;
@@ -131,6 +131,7 @@ pub(crate) async fn persist_image_generation_item(
    turn_context: &TurnContext,
    image_item: &mut ImageGenerationItem,
 ) -> Option<AbsolutePathBuf> {
+    image_item.saved_path = None;
    let session_id = sess.conversation_id.to_string();
    match save_image_generation_result(
        &turn_context.config.codex_home,
@@ -163,15 +164,12 @@ pub(crate) async fn persist_image_generation_item(
    }
 }

-pub(crate) async fn finalize_image_generation_item(
+async fn record_image_generation_instructions(
    sess: &Session,
    turn_context: &TurnContext,
-    image_item: &mut ImageGenerationItem,
+    image_item: &ImageGenerationItem,
 ) {
-    if persist_image_generation_item(sess, turn_context, image_item)
-        .await
-        .is_none()
-    {
+    if image_item.saved_path.is_none() {
        return;
    }
    let session_id = sess.conversation_id.to_string();
@@ -530,27 +528,16 @@ pub(crate) async fn handle_non_tool_response_item(
        | ResponseItem::WebSearchCall { .. }
        | ResponseItem::ImageGenerationCall { .. } => {
            let mut turn_item = parse_turn_item(item)?;
-            if let TurnItemContributorPolicy::Run(turn_store) = contributor_policy {
-                apply_turn_item_contributors(sess, turn_store, &mut turn_item).await;
-            }
-            if let TurnItem::AgentMessage(agent_message) = &mut turn_item {
-                let combined = agent_message
-                    .content
-                    .iter()
-                    .map(|entry| match entry {
-                        codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(),
-                    })
-                    .collect::<String>();
-                let (stripped, memory_citation) =
-                    strip_hidden_assistant_markup_and_parse_memory_citation(&combined, plan_mode);
-                agent_message.content =
-                    vec![codex_protocol::items::AgentMessageContent::Text { text: stripped }];
-                if agent_message.memory_citation.is_none() {
-                    agent_message.memory_citation = memory_citation;
-                }
-            }
-            if let TurnItem::ImageGeneration(image_item) = &mut turn_item {
-                finalize_image_generation_item(sess, turn_context, image_item).await;
+            finalize_turn_item(
+                sess,
+                turn_context,
+                contributor_policy,
+                &mut turn_item,
+                plan_mode,
+            )
+            .await;
+            if let TurnItem::ImageGeneration(image_item) = &turn_item {
+                record_image_generation_instructions(sess, turn_context, image_item).await;
            }
            Some(turn_item)
        }
@@ -564,6 +551,37 @@ pub(crate) async fn handle_non_tool_response_item(
    }
 }

+pub(crate) async fn finalize_turn_item(
+    sess: &Session,
+    turn_context: &TurnContext,
+    contributor_policy: TurnItemContributorPolicy<'_>,
+    turn_item: &mut TurnItem,
+    plan_mode: bool,
+) {
+    if let TurnItemContributorPolicy::Run(turn_store) = contributor_policy {
+        apply_turn_item_contributors(sess, turn_store, turn_item).await;
+    }
+    if let TurnItem::AgentMessage(agent_message) = &mut *turn_item {
+        let combined = agent_message
+            .content
+            .iter()
+            .map(|entry| match entry {
+                codex_protocol::items::AgentMessageContent::Text { text } => text.as_str(),
+            })
+            .collect::<String>();
+        let (stripped, memory_citation) =
+            strip_hidden_assistant_markup_and_parse_memory_citation(&combined, plan_mode);
+        agent_message.content =
+            vec![codex_protocol::items::AgentMessageContent::Text { text: stripped }];
+        if agent_message.memory_citation.is_none() {
+            agent_message.memory_citation = memory_citation;
+        }
+    }
+    if let TurnItem::ImageGeneration(image_item) = &mut *turn_item {
+        persist_image_generation_item(sess, turn_context, image_item).await;
+    }
+}
+
 pub(crate) fn last_assistant_message_from_item(
    item: &ResponseItem,
    plan_mode: bool,
@@ -4,7 +4,6 @@ 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::ToolSearchInfo;
@@ -12,12 +11,11 @@ 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::stream_events_utils::TurnItemContributorPolicy;
+use crate::stream_events_utils::finalize_turn_item;
 use crate::tools::context::ToolInvocation;
 use crate::tools::context::ToolOutput;
 use crate::tools::context::ToolPayload;
@@ -76,6 +74,10 @@ 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)
+        }
    }
 }

@@ -85,8 +87,9 @@ impl TurnItemEmitter for CoreTurnItemEmitter {
            let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else {
                return;
            };
-            let item = extension_turn_item(item);
-            session.emit_turn_item_started(turn.as_ref(), &item).await;
+            session
+                .emit_turn_item_started(turn.as_ref(), &extension_turn_item(item))
+                .await;
        })
    }

@@ -95,54 +98,18 @@ impl TurnItemEmitter for CoreTurnItemEmitter {
            let (Some(session), Some(turn)) = (self.session.upgrade(), self.turn.upgrade()) else {
                return;
            };
-            let item = extension_turn_item(item);
+            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;
            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 {
@@ -167,6 +134,8 @@ async fn to_extension_call(invocation: &ToolInvocation) -> ExtensionToolCall {
 mod tests {
    use std::sync::Arc;

+    use codex_extension_api::ExtensionData;
+    use codex_extension_api::TurnItemContributor;
    use codex_protocol::items::TurnItem;
    use codex_protocol::items::WebSearchItem;
    use codex_protocol::models::ContentItem;
@@ -174,10 +143,13 @@ mod tests {
    use codex_protocol::models::WebSearchAction;
    use codex_protocol::protocol::EventMsg;
    use codex_tools::ExtensionTurnItem;
+    use codex_utils_absolute_path::test_support::PathExt;
+    use codex_utils_absolute_path::test_support::test_path_buf;
    use pretty_assertions::assert_eq;
    use serde_json::json;
    use tokio::sync::Mutex;

+    use super::CoreTurnItemEmitter;
    use super::ExtensionToolAdapter;
    use crate::tools::context::ToolCallSource;
    use crate::tools::context::ToolInvocation;
@@ -409,8 +381,54 @@ mod tests {
        assert_eq!(end.action, expected.action);
    }

-    struct ImageGenerationExtensionExecutor {
-        output_hint: Arc<Mutex<Option<String>>>,
+    struct ImageGenerationExtensionExecutor;
+
+    #[derive(Debug)]
+    struct ExtensionTurnItemContributorRan;
+
+    struct RecordExtensionTurnItemContributor;
+
+    #[async_trait::async_trait]
+    impl TurnItemContributor for RecordExtensionTurnItemContributor {
+        async fn contribute(
+            &self,
+            _thread_store: &ExtensionData,
+            turn_store: &ExtensionData,
+            _item: &mut TurnItem,
+        ) -> Result<(), String> {
+            turn_store.insert(ExtensionTurnItemContributorRan);
+            Ok(())
+        }
+    }
+
+    #[tokio::test]
+    async fn extension_completion_runs_turn_item_contributors() {
+        let (mut session, turn) = crate::session::tests::make_session_and_context().await;
+        let mut builder = codex_extension_api::ExtensionRegistryBuilder::new();
+        builder.turn_item_contributor(Arc::new(RecordExtensionTurnItemContributor));
+        session.services.extensions = Arc::new(builder.build());
+        let session = Arc::new(session);
+        let turn = Arc::new(turn);
+        let emitter = CoreTurnItemEmitter {
+            session: Arc::downgrade(&session),
+            turn: Arc::downgrade(&turn),
+        };
+
+        codex_tools::TurnItemEmitter::emit_completed(
+            &emitter,
+            ExtensionTurnItem::WebSearch(WebSearchItem {
+                id: "search-1".to_string(),
+                query: "contributors".to_string(),
+                action: WebSearchAction::Other,
+            }),
+        )
+        .await;
+
+        assert!(
+            turn.extension_data
+                .get::<ExtensionTurnItemContributorRan>()
+                .is_some()
+        );
    }

    #[async_trait::async_trait]
@@ -434,15 +452,28 @@ mod tests {
            &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(),
-                )
+            call.turn_item_emitter
+                .emit_started(ExtensionTurnItem::ImageGeneration(
+                    codex_protocol::items::ImageGenerationItem {
+                        id: call.call_id.clone(),
+                        status: "in_progress".to_string(),
+                        revised_prompt: None,
+                        result: String::new(),
+                        saved_path: None,
+                    },
+                ))
+                .await;
+            call.turn_item_emitter
+                .emit_completed(ExtensionTurnItem::ImageGeneration(
+                    codex_protocol::items::ImageGenerationItem {
+                        id: call.call_id,
+                        status: "completed".to_string(),
+                        revised_prompt: Some("A tiny blue square".to_string()),
+                        result: "cG5n".to_string(),
+                        saved_path: Some(test_path_buf("/tmp/extension-claimed.png").abs()),
+                    },
+                ))
                .await;
-            *self.output_hint.lock().await = output_hint;
            Ok(Box::new(codex_tools::JsonToolOutput::new(
                json!({ "ok": true }),
            )))
@@ -451,10 +482,7 @@ mod tests {

    #[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 handler = ExtensionToolAdapter::new(Arc::new(ImageGenerationExtensionExecutor));
        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,
@@ -521,17 +549,5 @@ mod tests {
            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(),
-            ))
-        );
    }
 }
@@ -13,7 +13,6 @@ 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;
@@ -26,7 +26,6 @@ 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,10 +54,9 @@ fn generate_uses_fixed_request_defaults() {
 }

 #[test]
-fn generated_output_returns_image_input_and_output_hint() {
+fn generated_output_returns_image_input() {
    let output = GeneratedImageOutput {
        result: RESULT.to_string(),
-        output_hint: Some(OUTPUT_HINT.to_string()),
    };

    let ResponseInputItem::FunctionCallOutput {
@@ -73,15 +71,10 @@ fn generated_output_returns_image_input_and_output_hint() {
    };
    assert_eq!(
        content_items,
-        vec![
-            FunctionCallOutputContentItem::InputImage {
-                image_url: format!("data:image/png;base64,{RESULT}"),
-                detail: Some(DEFAULT_IMAGE_DETAIL),
-            },
-            FunctionCallOutputContentItem::InputText {
-                text: OUTPUT_HINT.to_string(),
-            },
-        ]
+        vec![FunctionCallOutputContentItem::InputImage {
+            image_url: format!("data:image/png;base64,{RESULT}"),
+            detail: Some(DEFAULT_IMAGE_DETAIL),
+        }]
    );
 }

@@ -3,6 +3,7 @@ use codex_api::ImageEditRequest;
 use codex_api::ImageGenerationRequest;
 use codex_api::ImageQuality;
 use codex_api::ImageUrl;
+use codex_extension_api::ExtensionTurnItem;
 use codex_extension_api::FunctionCallError;
 use codex_extension_api::ToolCall;
 use codex_extension_api::ToolExecutor;
@@ -11,6 +12,7 @@ use codex_extension_api::ToolOutput;
 use codex_extension_api::ToolPayload;
 use codex_extension_api::ToolSpec;
 use codex_extension_api::parse_tool_input_schema;
+use codex_protocol::items::ImageGenerationItem;
 use codex_protocol::models::ContentItem;
 use codex_protocol::models::DEFAULT_IMAGE_DETAIL;
 use codex_protocol::models::FunctionCallOutputBody;
@@ -84,6 +86,15 @@ 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())?;
+        call.turn_item_emitter
+            .emit_started(ExtensionTurnItem::ImageGeneration(ImageGenerationItem {
+                id: call.call_id.clone(),
+                status: "in_progress".to_string(),
+                revised_prompt: None,
+                result: String::new(),
+                saved_path: None,
+            }))
+            .await;
        let response = match request {
            ImageRequest::Generate(request) => self.backend.generate(request).await,
            ImageRequest::Edit(request) => self.backend.edit(request).await,
@@ -96,14 +107,16 @@ impl ToolExecutor<ToolCall> for ImageGenerationTool {
                "image generation returned no image data".to_string(),
            ));
        };
-        let output_hint = call
-            .turn_item_emitter
-            .image_generation_completed(call.call_id.clone(), args.prompt, result.clone())
+        call.turn_item_emitter
+            .emit_completed(ExtensionTurnItem::ImageGeneration(ImageGenerationItem {
+                id: call.call_id.clone(),
+                status: "completed".to_string(),
+                revised_prompt: Some(args.prompt),
+                result: result.clone(),
+                saved_path: None,
+            }))
            .await;
-        Ok(Box::new(GeneratedImageOutput {
-            result,
-            output_hint,
-        }))
+        Ok(Box::new(GeneratedImageOutput { result }))
    }
 }

@@ -284,7 +297,6 @@ fn imagegen_tool_spec() -> ToolSpec {

 struct GeneratedImageOutput {
    result: String,
-    output_hint: Option<String>,
 }

 impl ToolOutput for GeneratedImageOutput {
@@ -298,17 +310,12 @@ impl ToolOutput for GeneratedImageOutput {
        true
    }

-    /// Returns generated bytes and persisted-artifact context for model follow-up.
+    /// Returns generated bytes for model follow-up.
    fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem {
-        let mut content = vec![FunctionCallOutputContentItem::InputImage {
+        let content = vec![FunctionCallOutputContentItem::InputImage {
            image_url: format!("data:image/png;base64,{}", self.result),
            detail: Some(DEFAULT_IMAGE_DETAIL),
        }];
-        if let Some(output_hint) = &self.output_hint {
-            content.push(FunctionCallOutputContentItem::InputText {
-                text: output_hint.clone(),
-            });
-        }
        ResponseInputItem::FunctionCallOutput {
            call_id: call_id.to_string(),
            output: FunctionCallOutputPayload {
@@ -64,7 +64,6 @@ pub use responses_api::mcp_tool_to_responses_api_tool;
 pub use responses_api::tool_definition_to_responses_api_tool;
 pub use tool_call::ConversationHistory;
 pub use tool_call::ExtensionTurnItem;
-pub use tool_call::ImageGenerationCompletionFuture;
 pub use tool_call::NoopTurnItemEmitter;
 pub use tool_call::ToolCall;
 pub use tool_call::TurnItemEmissionFuture;
@@ -1,6 +1,7 @@
 use crate::FunctionCallError;
 use crate::ToolName;
 use crate::ToolPayload;
+use codex_protocol::items::ImageGenerationItem;
 use codex_protocol::items::WebSearchItem;
 use codex_protocol::models::ResponseItem;
 use codex_utils_output_truncation::TruncationPolicy;
@@ -29,20 +30,14 @@ impl ConversationHistory {
 /// Future returned when an extension tool emits a visible turn-item lifecycle event.
 pub type TurnItemEmissionFuture<'a> = Pin<Box<dyn Future<Output = ()> + Send + 'a>>;

-/// Future returned when an image-generation extension publishes completed image bytes.
-pub type ImageGenerationCompletionFuture<'a> =
-    Pin<Box<dyn Future<Output = Option<String>> + Send + 'a>>;
-
-/// Visible turn items that an extension fully owns and may emit as-is.
-///
-/// Add only item kinds that require no additional host finalization before
-/// persistence or client delivery. Richer items need a host-owned publish path.
+/// Visible turn items that an extension may publish into the host lifecycle.
 #[derive(Clone, Debug, PartialEq)]
 pub enum ExtensionTurnItem {
    WebSearch(WebSearchItem),
+    ImageGeneration(ImageGenerationItem),
 }

-/// Host-provided capability for extension tools to emit finalized visible turn items.
+/// Host-provided capability for extension tools to emit visible turn items.
 ///
 /// Implementations route lifecycle events through the host's normal item event
 /// pipeline, including any persistence and client delivery owned by the host.
@@ -50,21 +45,8 @@ pub trait TurnItemEmitter: Send + Sync {
    /// Emits the beginning of one visible turn item.
    fn emit_started<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a>;

-    /// Emits the completion of one visible turn item.
+    /// Emits one visible turn item after host-owned finalization.
    fn emit_completed<'a>(&'a self, item: ExtensionTurnItem) -> TurnItemEmissionFuture<'a>;
-
-    /// Publishes image bytes for host persistence and visible completion.
-    ///
-    /// Returns persisted-artifact context for the extension's model-facing
-    /// function output when the host saves the generated image successfully.
-    fn image_generation_completed<'a>(
-        &'a self,
-        _call_id: String,
-        _prompt: String,
-        _result: String,
-    ) -> ImageGenerationCompletionFuture<'a> {
-        Box::pin(std::future::ready(None))
-    }
 }

 /// Turn-item emitter used when a caller does not expose visible item emission.