diff --git a/codex-rs/code-mode/src/runtime/value.rs b/codex-rs/code-mode/src/runtime/value.rs index cba9b4efef0b..661179963197 100644 --- a/codex-rs/code-mode/src/runtime/value.rs +++ b/codex-rs/code-mode/src/runtime/value.rs @@ -6,6 +6,8 @@ use codex_code_mode_protocol::ImageDetail; const IMAGE_HELPER_EXPECTS_MESSAGE: &str = "image expects a non-empty image URL string, an object with image_url and optional detail, or a raw MCP image block"; const REMOTE_IMAGE_URL_ERROR: &str = "Tool call failed: remote image URLs are not supported in tool outputs. Pass a base64 data URI instead"; +const INVALID_IMAGE_URL_ERROR: &str = + "Tool call failed: invalid image output. Pass a base64 data URI instead"; const CODEX_IMAGE_DETAIL_META_KEY: &str = "codex/imageDetail"; pub(super) fn serialize_output_text( @@ -59,10 +61,15 @@ pub(super) fn normalize_output_image( if image_url.is_empty() { return Err(IMAGE_HELPER_EXPECTS_MESSAGE.to_string()); } - let lower = image_url.to_ascii_lowercase(); - if lower.starts_with("http://") || lower.starts_with("https://") { + let Some((scheme, _)) = image_url.split_once(':') else { + return Err(INVALID_IMAGE_URL_ERROR.to_string()); + }; + if scheme.eq_ignore_ascii_case("http") || scheme.eq_ignore_ascii_case("https") { return Err(REMOTE_IMAGE_URL_ERROR.to_string()); } + if !scheme.eq_ignore_ascii_case("data") { + return Err(INVALID_IMAGE_URL_ERROR.to_string()); + } let detail = detail_override.or(detail); let detail = match detail { diff --git a/codex-rs/code-mode/src/service_tests.rs b/codex-rs/code-mode/src/service_tests.rs index 49cfc53d9912..707540e3f07b 100644 --- a/codex-rs/code-mode/src/service_tests.rs +++ b/codex-rs/code-mode/src/service_tests.rs @@ -873,6 +873,40 @@ async fn image_helpers_reject_remote_urls() { } } +#[tokio::test] +async fn image_helpers_reject_invalid_image_outputs() { + let image_url = + "Error executing tool exec: Expected at least one message to convert to CallToolResult"; + for source in [ + format!("image({image_url:?}, \"original\");"), + format!("generatedImage({{ image_url: {image_url:?} }});"), + ] { + let service = InProcessCodeModeSession::new(); + + let response = execute( + &service, + ExecuteRequest { + source, + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: Vec::new(), + error_text: Some( + "Tool call failed: invalid image output. Pass a base64 data URI instead" + .to_string(), + ), + } + ); + } +} + #[tokio::test] async fn image_helper_rejects_unsupported_detail() { let service = InProcessCodeModeSession::new(); diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index a20eaa267a8d..7d2f1e3205e9 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -2983,6 +2983,49 @@ async fn code_mode_image_helper_rejects_remote_url() -> Result<()> { Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_image_helper_rejects_invalid_image_output() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn( + &server, + "use exec to return an image", + r#" +const s = "Error executing tool exec: Expected at least one message to convert to CallToolResult"; +image(s.trim(), "original"); +"#, + ) + .await?; + + let req = second_mock.single_request(); + let items = custom_tool_output_items(&req, "call-1"); + let (_, success) = custom_tool_output_body_and_success(&req, "call-1"); + assert_ne!( + success, + Some(true), + "code_mode invalid image output unexpectedly succeeded" + ); + assert_eq!(items.len(), 2); + assert_regex_match( + concat!( + r"(?s)\A", + r"Script failed\nWall time \d+\.\d seconds\nOutput:\n\z" + ), + text_item(&items, /*index*/ 0), + ); + assert_eq!( + text_item(&items, /*index*/ 1), + concat!( + "Script error:\n", + "Tool call failed: invalid image output. ", + "Pass a base64 data URI instead" + ) + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_can_use_view_image_result_with_image_helper() -> Result<()> { skip_if_no_network!(Ok(()));