From 958f93f899c99e8954535c0a6a2e75adde8fd601 Mon Sep 17 00:00:00 2001 From: Won Park Date: Fri, 13 Mar 2026 10:29:19 -0700 Subject: [PATCH] sending back imagaegencall response back to responseapi (#14558) Sending back the ResponseItem::ImageGenerationCall as is, because it is now supported from the API-side. --- codex-rs/core/src/codex_tests.rs | 28 ++- codex-rs/core/src/context_manager/history.rs | 3 - .../core/src/context_manager/history_tests.rs | 53 ++--- .../core/src/context_manager/normalize.rs | 45 +---- codex-rs/core/src/stream_events_utils.rs | 11 +- .../core/src/stream_events_utils_tests.rs | 7 +- codex-rs/core/tests/suite/model_switching.rs | 185 +++++++++++++++++- 7 files changed, 211 insertions(+), 121 deletions(-) diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 76ee7bcb5..a1d88ff3d 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/codex-rs/core/src/codex_tests.rs @@ -59,6 +59,7 @@ use codex_app_server_protocol::AppInfo; use codex_otel::TelemetryAuthMode; use codex_protocol::models::BaseInstructions; use codex_protocol::models::ContentItem; +use codex_protocol::models::DeveloperInstructions; use codex_protocol::models::ResponseInputItem; use codex_protocol::models::ResponseItem; use codex_protocol::openai_models::ModelsResponse; @@ -154,15 +155,6 @@ fn developer_input_texts(items: &[ResponseItem]) -> Vec<&str> { .collect() } -fn default_image_save_developer_message_text() -> String { - let image_output_dir = crate::stream_events_utils::default_image_generation_output_dir(); - format!( - "Generated images are saved to {} as {} by default.", - image_output_dir.display(), - image_output_dir.join(".png").display(), - ) -} - fn test_tool_runtime(session: Arc, turn_context: Arc) -> ToolCallRuntime { let router = Arc::new(ToolRouter::from_config( &turn_context.tools_config, @@ -3203,13 +3195,12 @@ async fn build_initial_context_omits_default_image_save_location_without_image_h } #[tokio::test] -async fn handle_output_item_done_records_image_save_message_after_successful_save() { +async fn handle_output_item_done_records_image_save_history_message() { let (session, turn_context) = make_session_and_context().await; let session = Arc::new(session); let turn_context = Arc::new(turn_context); let call_id = "ig_history_records_message"; - let expected_saved_path = crate::stream_events_utils::default_image_generation_output_dir() - .join(format!("{call_id}.png")); + let expected_saved_path = std::env::temp_dir().join(format!("{call_id}.png")); let _ = std::fs::remove_file(&expected_saved_path); let item = ResponseItem::ImageGenerationCall { id: call_id.to_string(), @@ -3229,9 +3220,13 @@ async fn handle_output_item_done_records_image_save_message_after_successful_sav .expect("image generation item should succeed"); let history = session.clone_history().await; - let expected_message: ResponseItem = - DeveloperInstructions::new(default_image_save_developer_message_text()).into(); - assert_eq!(history.raw_items(), &[expected_message, item]); + let save_message: ResponseItem = DeveloperInstructions::new(format!( + "Generated images are saved to {} as {} by default.", + std::env::temp_dir().display(), + std::env::temp_dir().join(".png").display(), + )) + .into(); + assert_eq!(history.raw_items(), &[save_message, item]); assert_eq!( std::fs::read(&expected_saved_path).expect("saved file"), b"foo" @@ -3245,8 +3240,7 @@ async fn handle_output_item_done_skips_image_save_message_when_save_fails() { let session = Arc::new(session); let turn_context = Arc::new(turn_context); let call_id = "ig_history_no_message"; - let expected_saved_path = crate::stream_events_utils::default_image_generation_output_dir() - .join(format!("{call_id}.png")); + let expected_saved_path = std::env::temp_dir().join(format!("{call_id}.png")); let _ = std::fs::remove_file(&expected_saved_path); let item = ResponseItem::ImageGenerationCall { id: call_id.to_string(), diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 05ac09ba1..19dbcf729 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -344,9 +344,6 @@ impl ContextManager { // all outputs must have a corresponding function/tool call normalize::remove_orphan_outputs(&mut self.items); - //rewrite image_gen_calls to messages to support stateless input - normalize::rewrite_image_generation_calls_for_stateless_input(&mut self.items); - // strip images when model does not support them normalize::strip_images_when_unsupported(input_modalities, &mut self.items); } diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index 29400ba31..066272748 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -398,7 +398,7 @@ fn for_prompt_strips_images_when_model_does_not_support_images() { } #[test] -fn for_prompt_rewrites_image_generation_calls_when_images_are_supported() { +fn for_prompt_preserves_image_generation_calls_when_images_are_supported() { let history = create_history_with_items(vec![ ResponseItem::ImageGenerationCall { id: "ig_123".to_string(), @@ -420,25 +420,11 @@ fn for_prompt_rewrites_image_generation_calls_when_images_are_supported() { assert_eq!( history.for_prompt(&default_input_modalities()), vec![ - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ - ContentItem::InputText { - text: "Image Generation Call".to_string(), - }, - ContentItem::InputText { - text: "Image ID: ig_123".to_string(), - }, - ContentItem::InputText { - text: "Prompt: lobster".to_string(), - }, - ContentItem::InputImage { - image_url: "data:image/png;base64,Zm9v".to_string(), - }, - ], - end_turn: None, - phase: None, + ResponseItem::ImageGenerationCall { + id: "ig_123".to_string(), + status: "generating".to_string(), + revised_prompt: Some("lobster".to_string()), + result: "Zm9v".to_string(), }, ResponseItem::Message { id: None, @@ -454,7 +440,7 @@ fn for_prompt_rewrites_image_generation_calls_when_images_are_supported() { } #[test] -fn for_prompt_rewrites_image_generation_calls_when_images_are_unsupported() { +fn for_prompt_clears_image_generation_result_when_images_are_unsupported() { let history = create_history_with_items(vec![ ResponseItem::Message { id: None, @@ -485,26 +471,11 @@ fn for_prompt_rewrites_image_generation_calls_when_images_are_unsupported() { end_turn: None, phase: None, }, - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ - ContentItem::InputText { - text: "Image Generation Call".to_string(), - }, - ContentItem::InputText { - text: "Image ID: ig_123".to_string(), - }, - ContentItem::InputText { - text: "Prompt: lobster".to_string(), - }, - ContentItem::InputText { - text: "image content omitted because you do not support image input" - .to_string(), - }, - ], - end_turn: None, - phase: None, + ResponseItem::ImageGenerationCall { + id: "ig_123".to_string(), + status: "completed".to_string(), + revised_prompt: Some("lobster".to_string()), + result: String::new(), }, ] ); diff --git a/codex-rs/core/src/context_manager/normalize.rs b/codex-rs/core/src/context_manager/normalize.rs index 3f13c5d9e..c217e0939 100644 --- a/codex-rs/core/src/context_manager/normalize.rs +++ b/codex-rs/core/src/context_manager/normalize.rs @@ -289,48 +289,6 @@ where } } -pub(crate) fn rewrite_image_generation_calls_for_stateless_input(items: &mut Vec) { - let original_items = std::mem::take(items); - *items = original_items - .into_iter() - .map(|item| match item { - ResponseItem::ImageGenerationCall { - id, - revised_prompt, - result, - .. - } => { - let image_url = if result.starts_with("data:") { - result - } else { - format!("data:image/png;base64,{result}") - }; - let revised_prompt = revised_prompt.unwrap_or_default(); - - ResponseItem::Message { - id: None, - role: "user".to_string(), - content: vec![ - ContentItem::InputText { - text: "Image Generation Call".to_string(), - }, - ContentItem::InputText { - text: format!("Image ID: {id}"), - }, - ContentItem::InputText { - text: format!("Prompt: {revised_prompt}"), - }, - ContentItem::InputImage { image_url }, - ], - end_turn: None, - phase: None, - } - } - _ => item, - }) - .collect(); -} - /// Strip image content from messages and tool outputs when the model does not support images. /// When `input_modalities` contains `InputModality::Image`, no stripping is performed. pub(crate) fn strip_images_when_unsupported( @@ -377,6 +335,9 @@ pub(crate) fn strip_images_when_unsupported( *content_items = normalized_content_items; } } + ResponseItem::ImageGenerationCall { result, .. } => { + result.clear(); + } _ => {} } } diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index e5231314e..a44bc01f5 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -73,15 +73,11 @@ async fn save_image_generation_result(call_id: &str, result: &str) -> Result PathBuf { - std::env::temp_dir() -} - /// Persist a completed model response item and record any cited memory usage. pub(crate) async fn record_completed_response_item( sess: &Session, @@ -214,7 +210,6 @@ pub(crate) async fn handle_output_item_done( .emit_turn_item_completed(&ctx.turn_context, turn_item) .await; } - record_completed_response_item(ctx.sess.as_ref(), ctx.turn_context.as_ref(), &item) .await; let last_agent_message = last_assistant_message_from_item(&item, plan_mode); @@ -310,7 +305,7 @@ pub(crate) async fn handle_non_tool_response_item( match save_image_generation_result(&image_item.id, &image_item.result).await { Ok(path) => { image_item.saved_path = Some(path.to_string_lossy().into_owned()); - let image_output_dir = default_image_generation_output_dir(); + let image_output_dir = std::env::temp_dir(); let message: ResponseItem = DeveloperInstructions::new(format!( "Generated images are saved to {} as {} by default.", image_output_dir.display(), @@ -324,7 +319,7 @@ pub(crate) async fn handle_non_tool_response_item( .await; } Err(err) => { - let output_dir = default_image_generation_output_dir(); + let output_dir = std::env::temp_dir(); tracing::warn!( call_id = %image_item.id, output_dir = %output_dir.display(), diff --git a/codex-rs/core/src/stream_events_utils_tests.rs b/codex-rs/core/src/stream_events_utils_tests.rs index b7f81be73..bfebb8902 100644 --- a/codex-rs/core/src/stream_events_utils_tests.rs +++ b/codex-rs/core/src/stream_events_utils_tests.rs @@ -1,4 +1,3 @@ -use super::default_image_generation_output_dir; use super::handle_non_tool_response_item; use super::last_assistant_message_from_item; use super::save_image_generation_result; @@ -71,7 +70,7 @@ fn last_assistant_message_from_item_returns_none_for_plan_only_hidden_message() #[tokio::test] async fn save_image_generation_result_saves_base64_to_png_in_temp_dir() { - let expected_path = default_image_generation_output_dir().join("ig_save_base64.png"); + let expected_path = std::env::temp_dir().join("ig_save_base64.png"); let _ = std::fs::remove_file(&expected_path); let saved_path = save_image_generation_result("ig_save_base64", "Zm9v") @@ -95,7 +94,7 @@ async fn save_image_generation_result_rejects_data_url_payload() { #[tokio::test] async fn save_image_generation_result_overwrites_existing_file() { - let existing_path = default_image_generation_output_dir().join("ig_overwrite.png"); + let existing_path = std::env::temp_dir().join("ig_overwrite.png"); std::fs::write(&existing_path, b"existing").expect("seed existing image"); let saved_path = save_image_generation_result("ig_overwrite", "Zm9v") @@ -109,7 +108,7 @@ async fn save_image_generation_result_overwrites_existing_file() { #[tokio::test] async fn save_image_generation_result_sanitizes_call_id_for_temp_dir_output_path() { - let expected_path = default_image_generation_output_dir().join("___ig___.png"); + let expected_path = std::env::temp_dir().join("___ig___.png"); let _ = std::fs::remove_file(&expected_path); let saved_path = save_image_generation_result("../ig/..", "Zm9v") diff --git a/codex-rs/core/tests/suite/model_switching.rs b/codex-rs/core/tests/suite/model_switching.rs index e4c9935d8..73e40f864 100644 --- a/codex-rs/core/tests/suite/model_switching.rs +++ b/codex-rs/core/tests/suite/model_switching.rs @@ -443,6 +443,9 @@ async fn model_change_from_image_to_text_strips_prior_image_content() -> Result< async fn generated_image_is_replayed_for_image_capable_models() -> Result<()> { skip_if_no_network!(Ok(())); + let saved_path = std::env::temp_dir().join("ig_123.png"); + let _ = std::fs::remove_file(&saved_path); + let server = MockServer::start().await; let image_model_slug = "test-image-model"; let image_model = test_model_info( @@ -527,19 +530,42 @@ async fn generated_image_is_replayed_for_image_capable_models() -> Result<()> { assert_eq!(requests.len(), 2, "expected two model requests"); let second_request = requests.last().expect("expected second request"); + let image_generation_calls = second_request.inputs_of_type("image_generation_call"); assert_eq!( - second_request.message_input_image_urls("user"), - vec!["data:image/png;base64,Zm9v".to_string()] + image_generation_calls.len(), + 1, + "expected generated image history to be replayed as an image_generation_call" ); + assert_eq!( + image_generation_calls[0]["id"].as_str(), + Some("ig_123"), + "expected the original image generation call id to be preserved" + ); + assert_eq!( + image_generation_calls[0]["result"].as_str(), + Some("Zm9v"), + "expected the original generated image payload to be preserved" + ); + assert!( + second_request + .message_input_texts("developer") + .iter() + .any(|text| text.contains("Generated images are saved to")), + "second request should include the saved-path note in model-visible history" + ); + let _ = std::fs::remove_file(&saved_path); Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn model_change_from_generated_image_to_text_strips_prior_generated_image_content() +async fn model_change_from_generated_image_to_text_preserves_prior_generated_image_call() -> Result<()> { skip_if_no_network!(Ok(())); + let saved_path = std::env::temp_dir().join("ig_123.png"); + let _ = std::fs::remove_file(&saved_path); + let server = MockServer::start().await; let image_model_slug = "test-image-model"; let text_model_slug = "test-text-only-model"; @@ -631,17 +657,164 @@ async fn model_change_from_generated_image_to_text_strips_prior_generated_image_ assert_eq!(requests.len(), 2, "expected two model requests"); let second_request = requests.last().expect("expected second request"); + let image_generation_calls = second_request.inputs_of_type("image_generation_call"); assert!( second_request.message_input_image_urls("user").is_empty(), - "second request should strip generated image content for text-only models" + "second request should not rewrite generated images into message input images" + ); + assert!( + image_generation_calls.len() == 1, + "second request should preserve the generated image call for text-only models" + ); + assert_eq!( + image_generation_calls[0]["id"].as_str(), + Some("ig_123"), + "second request should preserve the original generated image call id" + ); + assert_eq!( + image_generation_calls[0]["result"].as_str(), + Some(""), + "second request should strip generated image bytes for text-only models" ); assert!( second_request .message_input_texts("user") .iter() - .any(|text| text == "image content omitted because you do not support image input"), - "second request should include the image-omitted placeholder text" + .all(|text| text != "image content omitted because you do not support image input"), + "second request should not inject the image-omitted placeholder text" ); + assert!( + second_request + .message_input_texts("developer") + .iter() + .any(|text| text.contains("Generated images are saved to")), + "second request should include the saved-path note in model-visible history" + ); + let _ = std::fs::remove_file(&saved_path); + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn thread_rollback_after_generated_image_drops_entire_image_turn_history() -> Result<()> { + skip_if_no_network!(Ok(())); + + let saved_path = std::env::temp_dir().join("ig_rollback.png"); + let _ = std::fs::remove_file(&saved_path); + + let server = MockServer::start().await; + let image_model_slug = "test-image-model"; + let image_model = test_model_info( + image_model_slug, + "Test Image Model", + "supports image input", + default_input_modalities(), + ); + mount_models_once( + &server, + ModelsResponse { + models: vec![image_model], + }, + ) + .await; + + let responses = mount_sse_sequence( + &server, + vec![ + sse(vec![ + ev_response_created("resp-1"), + ev_image_generation_call("ig_rollback", "completed", "lobster", "Zm9v"), + ev_completed_with_tokens("resp-1", 10), + ]), + sse_completed("resp-2"), + ], + ) + .await; + + let mut builder = test_codex() + .with_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) + .with_config(move |config| { + config.model = Some(image_model_slug.to_string()); + }); + let test = builder.build(&server).await?; + let models_manager = test.thread_manager.get_models_manager(); + let _ = models_manager + .list_models(RefreshStrategy::OnlineIfUncached) + .await; + + test.codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "generate a lobster".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: test.cwd_path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: image_model_slug.to_string(), + effort: test.config.model_reasoning_effort, + service_tier: None, + summary: None, + collaboration_mode: None, + personality: None, + }) + .await?; + wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + test.codex + .submit(Op::ThreadRollback { num_turns: 1 }) + .await?; + wait_for_event(&test.codex, |ev| { + matches!(ev, EventMsg::ThreadRolledBack(_)) + }) + .await; + + test.codex + .submit(Op::UserTurn { + items: vec![UserInput::Text { + text: "after rollback".to_string(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + cwd: test.cwd_path().to_path_buf(), + approval_policy: AskForApproval::Never, + sandbox_policy: SandboxPolicy::new_read_only_policy(), + model: image_model_slug.to_string(), + effort: test.config.model_reasoning_effort, + service_tier: None, + summary: None, + collaboration_mode: None, + personality: None, + }) + .await?; + wait_for_event(&test.codex, |ev| matches!(ev, EventMsg::TurnComplete(_))).await; + + let requests = responses.requests(); + assert_eq!(requests.len(), 2, "expected two model requests"); + + let second_request = requests.last().expect("expected second request"); + assert!( + !second_request + .message_input_texts("user") + .iter() + .any(|text| text == "generate a lobster"), + "rollback should remove the rolled-back image-generation user turn" + ); + assert!( + !second_request + .message_input_texts("developer") + .iter() + .any(|text| text.contains("Generated images are saved to")), + "rollback should remove the generated-image save note with the rolled-back turn" + ); + assert!( + second_request + .inputs_of_type("image_generation_call") + .is_empty(), + "rollback should remove the generated image call with the rolled-back turn" + ); + let _ = std::fs::remove_file(&saved_path); Ok(()) }