diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index b1c6a2ac96..33387e620c 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3375,6 +3375,7 @@ dependencies = [ "codex-protocol", "codex-tools", "codex-utils-output-truncation", + "core_test_support", "http 1.4.0", "pretty_assertions", "serde_json", diff --git a/codex-rs/ext/history-notes/Cargo.toml b/codex-rs/ext/history-notes/Cargo.toml index 08fb92dd74..d043ee06e1 100644 --- a/codex-rs/ext/history-notes/Cargo.toml +++ b/codex-rs/ext/history-notes/Cargo.toml @@ -28,6 +28,7 @@ serde_json = { workspace = true } [dev-dependencies] codex-config = { workspace = true } codex-model-provider-info = { workspace = true } +core_test_support = { workspace = true } pretty_assertions = { workspace = true } tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt"] } diff --git a/codex-rs/ext/history-notes/src/tools.rs b/codex-rs/ext/history-notes/src/tools.rs index 6b390c69fa..ad7422c961 100644 --- a/codex-rs/ext/history-notes/src/tools.rs +++ b/codex-rs/ext/history-notes/src/tools.rs @@ -8,6 +8,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::models::FunctionCallOutputBody; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ResponseInputItem; @@ -280,7 +281,7 @@ impl HistoryNotesTool { .await .map_err(FunctionCallError::RespondToModel)?; - Ok(Box::new(HistoryNotesToolOutput { result })) + Ok(Box::new(HistoryNotesToolOutput::new(result)?)) } } @@ -324,6 +325,56 @@ impl<'call> ToolExecutor> for HistoryNotesTool { struct HistoryNotesToolOutput { result: Value, + output: FunctionCallOutputPayload, +} + +impl HistoryNotesToolOutput { + fn new(mut result: Value) -> Result { + // Separate attachments before serializing any text or retaining log output. + let images = result.as_object_mut().and_then(|map| map.remove("images")); + // The server applies the requested output budget before encryption. + let mut output = match result.get("encrypted_output").and_then(Value::as_str) { + Some(encrypted_content) => FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: encrypted_content.to_string(), + }, + ]), + None => FunctionCallOutputPayload::from_text(result.to_string()), + }; + if let Some(images) = images { + let invalid_image = || { + FunctionCallError::RespondToModel( + "History backend returned invalid image content.".to_string(), + ) + }; + let images = images.as_array().ok_or_else(invalid_image)?; + let mut content = match output.body { + FunctionCallOutputBody::Text(text) => { + vec![FunctionCallOutputContentItem::InputText { text }] + } + FunctionCallOutputBody::ContentItems(content) => content, + }; + for image in images { + let data = image + .get("data") + .and_then(Value::as_str) + .ok_or_else(invalid_image)?; + let mime_type = image + .get("mime_type") + .and_then(Value::as_str) + .ok_or_else(invalid_image)?; + let detail = + serde_json::from_value(image.get("detail").cloned().unwrap_or(Value::Null)) + .map_err(|_| invalid_image())?; + content.push(FunctionCallOutputContentItem::InputImage { + image_url: format!("data:{mime_type};base64,{data}"), + detail, + }); + } + output = FunctionCallOutputPayload::from_content_items(content); + } + Ok(Self { result, output }) + } } impl ToolOutput for HistoryNotesToolOutput { @@ -335,20 +386,15 @@ impl ToolOutput for HistoryNotesToolOutput { true } - fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { - // The server applies the requested output budget before encryption. - let output = match self.result.get("encrypted_output").and_then(Value::as_str) { - Some(encrypted_content) => FunctionCallOutputPayload::from_content_items(vec![ - FunctionCallOutputContentItem::EncryptedContent { - encrypted_content: encrypted_content.to_string(), - }, - ]), - None => FunctionCallOutputPayload::from_text(self.result.to_string()), - }; + fn post_tool_use_response(&self, _call_id: &str, _payload: &ToolPayload) -> Option { + // Hooks must not receive model-only image attachments. + Some(self.result.clone()) + } + fn to_response_item(&self, call_id: &str, _payload: &ToolPayload) -> ResponseInputItem { ResponseInputItem::FunctionCallOutput { call_id: call_id.to_string(), - output, + output: self.output.clone(), } } diff --git a/codex-rs/ext/history-notes/src/tools_tests.rs b/codex-rs/ext/history-notes/src/tools_tests.rs index 16cfb626a0..55b99cdc35 100644 --- a/codex-rs/ext/history-notes/src/tools_tests.rs +++ b/codex-rs/ext/history-notes/src/tools_tests.rs @@ -1,4 +1,6 @@ use codex_protocol::models::FunctionCallOutputContentItem; +use codex_protocol::models::FunctionCallOutputPayload; +use codex_protocol::models::ImageDetail; use codex_protocol::models::ResponseInputItem; use codex_tools::ToolOutput; use codex_tools::ToolPayload; @@ -9,15 +11,14 @@ use super::HistoryNotesToolOutput; #[test] fn preserves_encrypted_history_output() { - let result = HistoryNotesToolOutput { - result: json!({"encrypted_output": "enc_payload"}), - } - .to_response_item( - "call-1", - &ToolPayload::Function { - arguments: "{}".to_string(), - }, - ); + let result = HistoryNotesToolOutput::new(json!({"encrypted_output": "enc_payload"})) + .expect("valid output") + .to_response_item( + "call-1", + &ToolPayload::Function { + arguments: "{}".to_string(), + }, + ); let ResponseInputItem::FunctionCallOutput { output, .. } = result else { panic!("expected function-call output"); @@ -32,3 +33,114 @@ fn preserves_encrypted_history_output() { ) ); } + +#[test] +fn preserves_images_as_separate_output_items_without_logging_bytes() { + let output = HistoryNotesToolOutput::new(json!({ + "encrypted_output": "enc_payload", + "images": [ + {"data": "cG5n", "mime_type": "image/png", "detail": "original"}, + {"data": "anBlZw==", "mime_type": "image/jpeg", "detail": "low"}, + {"data": "Z2lm", "mime_type": "image/gif"}, + {"data": "d2VicA==", "mime_type": "image/webp", "detail": null} + ] + })) + .expect("valid image output"); + assert_eq!( + output.log_output(), + json!({"encrypted_output": "enc_payload"}).to_string() + ); + assert_eq!( + output.post_tool_use_response( + "call-1", + &ToolPayload::Function { + arguments: "{}".to_string() + } + ), + Some(json!({"encrypted_output": "enc_payload"})) + ); + assert_eq!( + output.to_response_item( + "call-1", + &ToolPayload::Function { + arguments: "{}".to_string() + } + ), + ResponseInputItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_payload".to_string() + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/png;base64,cG5n".to_string(), + detail: Some(ImageDetail::Original) + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/jpeg;base64,anBlZw==".to_string(), + detail: Some(ImageDetail::Low) + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/gif;base64,Z2lm".to_string(), + detail: None + }, + FunctionCallOutputContentItem::InputImage { + image_url: "data:image/webp;base64,d2VicA==".to_string(), + detail: None + }, + ]) + } + ); +} + +#[test] +fn accepts_empty_attachments_and_legacy_plaintext_results() { + for (result, expected) in [ + ( + json!({"encrypted_output": "enc_payload", "images": []}), + FunctionCallOutputPayload::from_content_items(vec![ + FunctionCallOutputContentItem::EncryptedContent { + encrypted_content: "enc_payload".to_string(), + }, + ]), + ), + ( + json!({"text": "legacy result"}), + FunctionCallOutputPayload::from_text(json!({"text": "legacy result"}).to_string()), + ), + ] { + let output = HistoryNotesToolOutput::new(result).expect("valid output"); + assert_eq!( + output.to_response_item( + "call-1", + &ToolPayload::Function { + arguments: "{}".to_string() + } + ), + ResponseInputItem::FunctionCallOutput { + call_id: "call-1".to_string(), + output: expected + } + ); + } +} + +#[test] +fn rejects_malformed_attachments_instead_of_silently_dropping_them() { + for images in [ + json!(null), + json!({}), + json!([null]), + json!([{"mime_type": "image/png"}]), + json!([{"data": "private-image-bytes"}]), + json!([{"data": "private-image-bytes", "mime_type": "image/png", "detail": "invalid"}]), + ] { + let result = HistoryNotesToolOutput::new( + json!({"encrypted_output": "enc_payload", "images": images}), + ); + let Err(codex_extension_api::FunctionCallError::RespondToModel(message)) = result else { + panic!("expected a model-facing image error"); + }; + assert_eq!(message, "History backend returned invalid image content."); + } +} diff --git a/codex-rs/ext/history-notes/tests/history_notes_extension.rs b/codex-rs/ext/history-notes/tests/history_notes_extension.rs index be5b31c3b5..00daaf7033 100644 --- a/codex-rs/ext/history-notes/tests/history_notes_extension.rs +++ b/codex-rs/ext/history-notes/tests/history_notes_extension.rs @@ -230,14 +230,27 @@ async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestRe ), ] { server.reset().await; + let mut response = json!({"encrypted_output": "enc_history"}); + let mut expected_output = + vec![json!({"type": "encrypted_content", "encrypted_content": "enc_history"})]; + if namespace == "history" && name == "read_item" { + // Images are not part of the text budget sent to the backend. + let data = "cG5n".repeat(1_000); + response["images"] = json!([ + {"data": data, "mime_type": "image/png", "detail": "original"}, + {"data": "anBlZw==", "mime_type": "image/jpeg", "detail": "high"} + ]); + expected_output.extend([ + json!({"type": "input_image", "image_url": format!("data:image/png;base64,{data}"), "detail": "original"}), + json!({"type": "input_image", "image_url": "data:image/jpeg;base64,anBlZw==", "detail": "high"}) + ]); + } Mock::given(method("POST")) .and(path(format!( "/backend-api/codex/alpha/{namespace}/v2/{name}" ))) .and(header("x-openai-actor-authorization", "actor-biscuit")) - .respond_with(ResponseTemplate::new(200).set_body_json(json!({ - "encrypted_output": "enc_history" - }))) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) .mount(&server) .await; let tool_name = ToolName::namespaced(namespace, name); @@ -245,12 +258,29 @@ async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestRe .iter() .find(|tool| tool.tool_name() == tool_name) .expect("exposed tool"); - tool.handle(tool_call(tool_name, arguments.clone())).await?; + let call = tool_call(tool_name, arguments.clone()); + let output = tool.handle(call.clone()).await?; + assert_eq!( + serde_json::to_value(output.to_response_item(&call.call_id, &call.payload))?, + json!({"type": "function_call_output", "call_id": call.call_id, "output": expected_output}) + ); arguments["context"] = json!({ "session_id": "session-123", "current_agent_name": "/root/worker", }); let requests = server.received_requests().await.expect("recorded requests"); + for request in &requests { + assert_eq!( + serde_json::from_str::( + request + .headers + .get("x-openai-tool-output-truncation-policy") + .expect("truncation policy header") + .to_str()? + )?, + json!({"mode": "bytes", "limit": 1024}) + ); + } assert_eq!( requests .iter() @@ -260,6 +290,39 @@ async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestRe ); } + // Plaintext backend responses must keep attachments out of text and logs too. + server.reset().await; + let text_result = json!({"content": "look: [image 1]", "n_chars": 15, "next_offset_chars": 15}); + let mut response = text_result.clone(); + response["images"] = json!([ + {"data": "cG5n", "mime_type": "image/png", "detail": "original"} + ]); + Mock::given(method("POST")) + .and(path("/backend-api/codex/alpha/history/v2/read_item")) + .and(header("x-openai-actor-authorization", "actor-biscuit")) + .respond_with(ResponseTemplate::new(200).set_body_json(response)) + .mount(&server) + .await; + let tool_name = ToolName::namespaced("history", "read_item"); + let read_item = tools + .iter() + .find(|tool| tool.tool_name() == tool_name) + .expect("history.read_item"); + let call = tool_call(tool_name, json!({"window_id": "window", "item_id": "item"})); + let output = read_item.handle(call.clone()).await?; + assert_eq!(output.log_output(), text_result.to_string()); + assert_eq!( + serde_json::to_value(output.to_response_item(&call.call_id, &call.payload))?, + json!({ + "type": "function_call_output", + "call_id": call.call_id, + "output": [ + {"type": "input_text", "text": text_result.to_string()}, + {"type": "input_image", "image_url": "data:image/png;base64,cG5n", "detail": "original"} + ] + }) + ); + for result in [ json!({"text": ""}), json!({"text": "x".repeat(4_001)}), diff --git a/codex-rs/ext/history-notes/tests/image_model_input.rs b/codex-rs/ext/history-notes/tests/image_model_input.rs new file mode 100644 index 0000000000..2a6a202546 --- /dev/null +++ b/codex-rs/ext/history-notes/tests/image_model_input.rs @@ -0,0 +1,91 @@ +//! Verify history attachments survive native preparation into the next model request. + +use std::sync::Arc; + +use codex_core::config::Config; +use codex_core::config::TokenBudgetConfig; +use codex_extension_api::ExtensionRegistryBuilder; +use codex_history_notes_extension::install; +use codex_login::AuthHeaders; +use codex_login::AuthManager; +use codex_login::CodexAuth; +use core_test_support::responses; +use core_test_support::test_codex::test_codex; +use http::HeaderMap; +use pretty_assertions::assert_eq; +use serde_json::json; +use wiremock::Mock; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; + +const PNG: &str = + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAAC0lEQVR4nGNgAAIAAAUAAXpeqz8AAAAASUVORK5CYII="; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn history_images_reach_the_next_model_request() -> Result<(), Box> { + let server = responses::start_mock_server().await; + Mock::given(method("POST")) + .and(path("/v1/alpha/history/v2/read_item")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "encrypted_output": "opaque-history-text", + "images": [ + {"data": PNG, "mime_type": "image/png", "detail": "original"}, + {"data": PNG, "mime_type": "image/png", "detail": "high"}, + {"data": PNG, "mime_type": "image/png", "detail": "auto"} + ] + }))) + .mount(&server) + .await; + let requests = responses::mount_sse_sequence( + &server, + vec![ + responses::sse(vec![ + responses::ev_function_call_with_namespace( + "read-history", + "history", + "read_item", + &json!({"window_id": "window", "item_id": "item"}).to_string(), + ), + responses::ev_completed("first-response"), + ]), + responses::sse(vec![ + responses::ev_assistant_message("assistant", "done"), + responses::ev_completed("second-response"), + ]), + ], + ) + .await; + let auth = CodexAuth::Headers(AuthHeaders::new(HeaderMap::new())); + let mut extensions = ExtensionRegistryBuilder::::new(); + install( + &mut extensions, + AuthManager::from_auth_for_testing(auth.clone()), + ); + let test = test_codex() + .with_auth(auth) + .with_extensions(Arc::new(extensions.build())) + .with_config(|config| { + config.model_provider.name = "OpenAI".to_string(); + config.token_budget = Some(TokenBudgetConfig { + use_history_notes_extension: true, + ..TokenBudgetConfig::default() + }); + }) + .build_with_auto_env(&server) + .await?; + test.submit_text_turn("Read the image from history.") + .await?; + let requests = requests.requests(); + assert_eq!(requests.len(), 2); + assert_eq!( + requests[1].function_call_output("read-history")["output"], + json!([ + {"type": "encrypted_content", "encrypted_content": "opaque-history-text"}, + {"type": "input_image", "image_url": format!("data:image/png;base64,{PNG}"), "detail": "original"}, + {"type": "input_image", "image_url": format!("data:image/png;base64,{PNG}"), "detail": "high"}, + {"type": "input_image", "image_url": format!("data:image/png;base64,{PNG}"), "detail": "auto"} + ]) + ); + Ok(()) +}