From 57f337a8e90360649112f8fc01d9ee2f9f2afa24 Mon Sep 17 00:00:00 2001 From: Won Park Date: Tue, 2 Jun 2026 12:00:04 -0700 Subject: [PATCH] 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` --- .../tests/suite/v2/image_generation.rs | 214 ++++++++++++++++++ codex-rs/app-server/tests/suite/v2/mod.rs | 1 + codex-rs/core/src/stream_events_utils.rs | 72 +++--- .../src/tools/handlers/extension_tools.rs | 170 +++++++------- codex-rs/ext/extension-api/src/lib.rs | 1 - codex-rs/ext/image-generation/src/tests.rs | 17 +- codex-rs/ext/image-generation/src/tool.rs | 37 +-- codex-rs/tools/src/lib.rs | 1 - codex-rs/tools/src/tool_call.rs | 28 +-- 9 files changed, 385 insertions(+), 156 deletions(-) create mode 100644 codex-rs/app-server/tests/suite/v2/image_generation.rs diff --git a/codex-rs/app-server/tests/suite/v2/image_generation.rs b/codex-rs/app-server/tests/suite/v2/image_generation.rs new file mode 100644 index 000000000..4752f1497 --- /dev/null +++ b/codex-rs/app-server/tests/suite/v2/image_generation.rs @@ -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::(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::(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 { + 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 +"# + ), + ) +} diff --git a/codex-rs/app-server/tests/suite/v2/mod.rs b/codex-rs/app-server/tests/suite/v2/mod.rs index 160598370..950cccebb 100644 --- a/codex-rs/app-server/tests/suite/v2/mod.rs +++ b/codex-rs/app-server/tests/suite/v2/mod.rs @@ -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; diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index bb1f5fba6..7914c8d48 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -131,6 +131,7 @@ pub(crate) async fn persist_image_generation_item( turn_context: &TurnContext, image_item: &mut ImageGenerationItem, ) -> Option { + 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::(); - 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::(); + 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, diff --git a/codex-rs/core/src/tools/handlers/extension_tools.rs b/codex-rs/core/src/tools/handlers/extension_tools.rs index 4ad88ac5b..0a1e2e58c 100644 --- a/codex-rs/core/src/tools/handlers/extension_tools.rs +++ b/codex-rs/core/src/tools/handlers/extension_tools.rs @@ -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>>, + 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::() + .is_some() + ); } #[async_trait::async_trait] @@ -434,15 +452,28 @@ mod tests { &self, call: codex_tools::ToolCall, ) -> Result, 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(), - )) - ); } } diff --git a/codex-rs/ext/extension-api/src/lib.rs b/codex-rs/ext/extension-api/src/lib.rs index 2f00a8dd9..4705581c8 100644 --- a/codex-rs/ext/extension-api/src/lib.rs +++ b/codex-rs/ext/extension-api/src/lib.rs @@ -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; diff --git a/codex-rs/ext/image-generation/src/tests.rs b/codex-rs/ext/image-generation/src/tests.rs index f866a9018..44f2d80eb 100644 --- a/codex-rs/ext/image-generation/src/tests.rs +++ b/codex-rs/ext/image-generation/src/tests.rs @@ -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), + }] ); } diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 75a79db2b..41d407d91 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -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 for ImageGenerationTool { async fn handle(&self, call: ToolCall) -> Result, 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 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, } 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 { diff --git a/codex-rs/tools/src/lib.rs b/codex-rs/tools/src/lib.rs index 28c4f93bb..1cee3f7d2 100644 --- a/codex-rs/tools/src/lib.rs +++ b/codex-rs/tools/src/lib.rs @@ -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; diff --git a/codex-rs/tools/src/tool_call.rs b/codex-rs/tools/src/tool_call.rs index 31143f9f1..e3fc21e86 100644 --- a/codex-rs/tools/src/tool_call.rs +++ b/codex-rs/tools/src/tool_call.rs @@ -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 + Send + 'a>>; -/// Future returned when an image-generation extension publishes completed image bytes. -pub type ImageGenerationCompletionFuture<'a> = - Pin> + 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.