diff --git a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs index 025c6ec18..5c18aa8a4 100644 --- a/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs +++ b/codex-rs/app-server/tests/suite/v2/imagegen_extension.rs @@ -147,6 +147,84 @@ async fn standalone_image_generation_returns_saved_path_hint_to_model() -> Resul Ok(()) } +#[tokio::test] +async fn standalone_image_generation_failure_emits_terminal_item() -> Result<()> { + let call_id = "image-run-failed"; + let server = responses::start_mock_server().await; + Mock::given(method("POST")) + .and(path("/api/codex/images/generations")) + .respond_with(ResponseTemplate::new(500).set_body_string("image backend failed")) + .expect(1) + .mount(&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!({"prompt": "paint a blue whale"}).to_string(), + ), + responses::ev_completed("resp-1"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("msg-1", "I could not generate the image."), + responses::ev_completed("resp-2"), + ]), + ], + ) + .await; + + let codex_home = TempDir::new()?; + create_config_toml(codex_home.path(), &server.uri(), ImagegenTestMode::Direct)?; + write_chatgpt_auth( + codex_home.path(), + ChatGptAuthFixture::new("access-chatgpt"), + AuthCredentialsStoreMode::File, + )?; + let mut mcp = + TestAppServer::new_with_env(codex_home.path(), &[("OPENAI_API_KEY", None)]).await?; + timeout(DEFAULT_READ_TIMEOUT, mcp.initialize()).await??; + start_image_generation_turn(&mut mcp).await?; + + let completed = timeout( + DEFAULT_READ_TIMEOUT, + wait_for_image_generation_completed(&mut mcp), + ) + .await??; + assert_eq!( + completed.item, + ThreadItem::ImageGeneration { + id: call_id.to_string(), + status: "failed".to_string(), + revised_prompt: Some("paint a blue whale".to_string()), + result: String::new(), + saved_path: None, + } + ); + + timeout( + DEFAULT_READ_TIMEOUT, + mcp.read_stream_until_notification_message("turn/completed"), + ) + .await??; + let requests = response_mock.requests(); + assert_eq!(requests.len(), 2); + let (output, _) = requests[1] + .function_call_output_content_and_success(call_id) + .context("image generation function output should be present")?; + assert!( + output + .as_deref() + .is_some_and(|text| text.contains("image generation failed")) + ); + + Ok(()) +} + #[tokio::test] async fn standalone_image_edit_uses_attached_model_visible_image() -> Result<()> { let edit_request = run_image_edit_test(|codex_home| { diff --git a/codex-rs/core/src/stream_events_utils.rs b/codex-rs/core/src/stream_events_utils.rs index 81beea5be..9ae20df74 100644 --- a/codex-rs/core/src/stream_events_utils.rs +++ b/codex-rs/core/src/stream_events_utils.rs @@ -578,7 +578,9 @@ pub(crate) async fn finalize_turn_item( agent_message.memory_citation = memory_citation; } } - if let TurnItem::ImageGeneration(image_item) = &mut *turn_item { + if let TurnItem::ImageGeneration(image_item) = &mut *turn_item + && image_item.status == "completed" + { persist_image_generation_item(sess, turn_context, image_item).await; } } diff --git a/codex-rs/ext/image-generation/src/tool.rs b/codex-rs/ext/image-generation/src/tool.rs index 2cbb0ab9f..3712ebbb0 100644 --- a/codex-rs/ext/image-generation/src/tool.rs +++ b/codex-rs/ext/image-generation/src/tool.rs @@ -117,17 +117,33 @@ impl ImageGenerationTool { saved_path: None, })) .await; - let response = match request { + let result = match request { ImageRequest::Generate(request) => self.backend.generate(request).await, ImageRequest::Edit(request) => self.backend.edit(request).await, } - .map_err(|err| { - FunctionCallError::RespondToModel(format!("image generation failed: {err}")) - })?; - let Some(result) = response.data.into_iter().next().map(|data| data.b64_json) else { - return Err(FunctionCallError::RespondToModel( - "image generation returned no image data".to_string(), - )); + .map_err(|err| format!("image generation failed: {err}")) + .and_then(|response| { + response + .data + .into_iter() + .next() + .map(|data| data.b64_json) + .ok_or_else(|| "image generation returned no image data".to_string()) + }); + let result = match result { + Ok(result) => result, + Err(message) => { + call.turn_item_emitter + .emit_completed(ExtensionTurnItem::ImageGeneration(ImageGenerationItem { + id: call.call_id.clone(), + status: "failed".to_string(), + revised_prompt: Some(args.prompt.clone()), + result: String::new(), + saved_path: None, + })) + .await; + return Err(FunctionCallError::RespondToModel(message)); + } }; call.turn_item_emitter .emit_completed(ExtensionTurnItem::ImageGeneration(ImageGenerationItem { diff --git a/codex-rs/tui/src/chatwidget/replay.rs b/codex-rs/tui/src/chatwidget/replay.rs index bdb8b8d3a..ae05a7198 100644 --- a/codex-rs/tui/src/chatwidget/replay.rs +++ b/codex-rs/tui/src/chatwidget/replay.rs @@ -151,11 +151,12 @@ impl ChatWidget { } ThreadItem::ImageGeneration { id, + status, revised_prompt, saved_path, .. } => { - self.on_image_generation_end(id, revised_prompt, saved_path); + self.on_image_generation_end(id, status, revised_prompt, saved_path); } ThreadItem::EnteredReviewMode { review, .. } => { if from_replay { diff --git a/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__failed_image_generation_call_history_snapshot.snap b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__failed_image_generation_call_history_snapshot.snap new file mode 100644 index 000000000..c69c51f27 --- /dev/null +++ b/codex-rs/tui/src/chatwidget/snapshots/codex_tui__chatwidget__tests__failed_image_generation_call_history_snapshot.snap @@ -0,0 +1,6 @@ +--- +source: tui/src/chatwidget/tests/exec_flow.rs +expression: "lines_to_single_string(&cells[0])" +--- +✗ Image generation failed + └ A tiny blue square diff --git a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs index 8dbace437..ca1c9bc68 100644 --- a/codex-rs/tui/src/chatwidget/tests/exec_flow.rs +++ b/codex-rs/tui/src/chatwidget/tests/exec_flow.rs @@ -864,6 +864,7 @@ async fn image_generation_call_adds_history_cell() { handle_image_generation_end( &mut chat, "call-image-generation", + "completed", Some("A tiny blue square".into()), Some(test_path_buf("/tmp/ig-1.png").abs()), ); @@ -876,6 +877,21 @@ async fn image_generation_call_adds_history_cell() { let combined = lines_to_single_string(&cells[0]).replace(&platform_file_url, "file:///tmp/ig-1.png"); assert_chatwidget_snapshot!("image_generation_call_history_snapshot", combined); + + handle_image_generation_end( + &mut chat, + "call-image-generation-failed", + "failed", + Some("A tiny blue square".into()), + /*saved_path*/ None, + ); + + let cells = drain_insert_history(&mut rx); + assert_eq!(cells.len(), 1, "expected a single failure history cell"); + assert_chatwidget_snapshot!( + "failed_image_generation_call_history_snapshot", + lines_to_single_string(&cells[0]) + ); } #[tokio::test] diff --git a/codex-rs/tui/src/chatwidget/tests/helpers.rs b/codex-rs/tui/src/chatwidget/tests/helpers.rs index 517f8e439..d15537647 100644 --- a/codex-rs/tui/src/chatwidget/tests/helpers.rs +++ b/codex-rs/tui/src/chatwidget/tests/helpers.rs @@ -694,6 +694,7 @@ pub(super) fn handle_view_image_tool_call( pub(super) fn handle_image_generation_end( chat: &mut ChatWidget, call_id: impl Into, + status: impl Into, revised_prompt: Option, saved_path: Option, ) { @@ -704,7 +705,7 @@ pub(super) fn handle_image_generation_end( completed_at_ms: 0, item: AppServerThreadItem::ImageGeneration { id: call_id.into(), - status: "completed".to_string(), + status: status.into(), revised_prompt, result: String::new(), saved_path, diff --git a/codex-rs/tui/src/chatwidget/tool_lifecycle.rs b/codex-rs/tui/src/chatwidget/tool_lifecycle.rs index 43d4c2e3c..e37c44a4f 100644 --- a/codex-rs/tui/src/chatwidget/tool_lifecycle.rs +++ b/codex-rs/tui/src/chatwidget/tool_lifecycle.rs @@ -29,12 +29,14 @@ impl ChatWidget { pub(super) fn on_image_generation_end( &mut self, call_id: String, + status: String, revised_prompt: Option, saved_path: Option, ) { self.flush_answer_stream_with_separator(); self.add_to_history(history_cell::new_image_generation_call( call_id, + &status, revised_prompt, saved_path, )); diff --git a/codex-rs/tui/src/history_cell/patches.rs b/codex-rs/tui/src/history_cell/patches.rs index 00d36fded..8a2dcb33f 100644 --- a/codex-rs/tui/src/history_cell/patches.rs +++ b/codex-rs/tui/src/history_cell/patches.rs @@ -73,15 +73,17 @@ pub(crate) fn new_view_image_tool_call(path: AbsolutePathBuf, cwd: &Path) -> Pla pub(crate) fn new_image_generation_call( call_id: String, + status: &str, revised_prompt: Option, saved_path: Option, ) -> PlainHistoryCell { - let detail = revised_prompt.unwrap_or_else(|| call_id.clone()); - - let mut lines: Vec> = vec![ - vec!["• ".dim(), "Generated Image:".bold()].into(), - vec![" └ ".dim(), detail.dim()].into(), - ]; + let detail = revised_prompt.unwrap_or(call_id); + let heading = if status == "failed" { + vec!["✗ ".red().bold(), "Image generation failed".bold()].into() + } else { + vec!["• ".dim(), "Generated Image:".bold()].into() + }; + let mut lines: Vec> = vec![heading, vec![" └ ".dim(), detail.dim()].into()]; if let Some(saved_path) = saved_path { let saved_path = Url::from_file_path(saved_path.as_path()) .map(|url| url.to_string()) diff --git a/codex-rs/tui/src/history_cell/tests.rs b/codex-rs/tui/src/history_cell/tests.rs index ec3ae49f2..4d2339212 100644 --- a/codex-rs/tui/src/history_cell/tests.rs +++ b/codex-rs/tui/src/history_cell/tests.rs @@ -474,6 +474,7 @@ fn image_generation_call_renders_saved_path() { ); let cell = new_image_generation_call( "call-image-generation".to_string(), + "completed", Some("A tiny blue square".to_string()), Some(saved_path), );