diff --git a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs b/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs index 7c7d94ac8b..ea37694606 100644 --- a/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs +++ b/codex-rs/core/src/compact_remote_v2_image_budget_tests.rs @@ -25,6 +25,15 @@ fn image() -> ContentItem { } } +fn file_image() -> ContentItem { + ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_123".to_string(), + }, + detail: Some(codex_protocol::models::ImageDetail::Original), + } +} + fn text(value: &str) -> ContentItem { ContentItem::InputText { text: value.to_string(), @@ -45,21 +54,23 @@ fn trim(items: Vec, max_tokens: usize) -> Vec { #[test] fn image_only_boundary_is_atomic_and_does_not_backfill_older_messages() { let newest = message(vec![text("new")]); - let items = vec![ - message(vec![text("old")]), - message(vec![image()]), - newest.clone(), - ]; - let image_tokens = images::content_item_token_count(&image()); - for (max_tokens, expected) in [ - ( - image_tokens + 1, - vec![message(vec![image()]), newest.clone()], - ), - (image_tokens, vec![newest.clone()]), - (1, vec![newest]), - ] { - assert_eq!(trim(items.clone(), max_tokens), expected); + for image in [image(), file_image()] { + let items = vec![ + message(vec![text("old")]), + message(vec![image.clone()]), + newest.clone(), + ]; + let image_tokens = images::content_item_token_count(&image); + for (max_tokens, expected) in [ + ( + image_tokens + 1, + vec![message(vec![image.clone()]), newest.clone()], + ), + (image_tokens, vec![newest.clone()]), + (1, vec![newest.clone()]), + ] { + assert_eq!(trim(items.clone(), max_tokens), expected); + } } } diff --git a/codex-rs/core/src/compact_remote_v2_images.rs b/codex-rs/core/src/compact_remote_v2_images.rs index 94cbdc239e..6e4914e9a4 100644 --- a/codex-rs/core/src/compact_remote_v2_images.rs +++ b/codex-rs/core/src/compact_remote_v2_images.rs @@ -1,10 +1,9 @@ -use crate::context_manager::estimate_image_bytes; +use crate::context_manager::estimate_image_reference_bytes; use codex_context_fragments::AnnotatedContent; use codex_context_fragments::set_annotated_content; use codex_context_fragments::to_annotated_content; use codex_history::ResponseItemEnvelope; use codex_protocol::models::ContentItem; -use codex_protocol::models::ImageReference; use codex_protocol::models::is_image_close_tag_text; use codex_protocol::models::is_image_open_tag_text; use codex_protocol::models::is_local_image_open_tag_text; @@ -18,17 +17,10 @@ pub(super) fn content_item_token_count(item: &ContentItem) -> usize { ContentItem::InputText { text } | ContentItem::OutputText { text } => { approx_token_count(text) } - ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => usize::try_from(approx_tokens_from_byte_count_i64(estimate_image_bytes( - image_url, *detail, - ))) + ContentItem::InputImage { image, detail } => usize::try_from( + approx_tokens_from_byte_count_i64(estimate_image_reference_bytes(image, *detail)), + ) .unwrap_or(usize::MAX), - ContentItem::InputImage { - image: ImageReference::File { .. }, - .. - } => 0, ContentItem::InputAudio { .. } => 0, } } diff --git a/codex-rs/core/src/context_manager/history.rs b/codex-rs/core/src/context_manager/history.rs index 883731f67c..b7ec2e4802 100644 --- a/codex-rs/core/src/context_manager/history.rs +++ b/codex-rs/core/src/context_manager/history.rs @@ -890,6 +890,7 @@ fn estimate_encrypted_function_output_length(encoded_len: usize) -> usize { /// Returns the same coarse, model-visible token estimate used for full history estimates. /// /// Counts content directly, excluding transport IDs, metadata, and outer JSON escaping. +/// Original-detail file images use the maximum patch count. pub(crate) fn estimate_item_token_count(item: &ResponseItem) -> i64 { let model_visible_bytes = estimate_response_item_model_visible_bytes(item); approx_tokens_from_byte_count_i64(model_visible_bytes) @@ -918,7 +919,6 @@ static ORIGINAL_IMAGE_ESTIMATE_CACHE: LazyLock i64 { - // TODO(kc) Account for file-backed image size after its token-cost contract is defined. match item { ResponseItem::Message { content, .. } => content .iter() @@ -926,14 +926,9 @@ fn estimate_response_item_model_visible_bytes(item: &ResponseItem) -> i64 { ContentItem::InputText { text } | ContentItem::OutputText { text } => { text_bytes(text) } - ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => estimate_image_bytes(image_url, *detail), - ContentItem::InputImage { - image: ImageReference::File { .. }, - .. - } => 0, + ContentItem::InputImage { image, detail } => { + estimate_image_reference_bytes(image, *detail) + } ContentItem::InputAudio { audio_url } => estimate_audio_bytes(audio_url), }) .fold(0i64, i64::saturating_add), @@ -1112,8 +1107,8 @@ fn estimate_original_image_bytes(image_url: &str) -> Option { }) } -/// Shared image estimate, excluding message framing. -pub(crate) fn estimate_image_bytes(image_url: &str, detail: Option) -> i64 { +/// Inline image estimate, excluding the data URL prefix and message framing. +fn estimate_image_bytes(image_url: &str, detail: Option) -> i64 { match detail { Some(ImageDetail::Original) => { estimate_original_image_bytes(image_url).unwrap_or(RESIZED_IMAGE_BYTES_ESTIMATE) @@ -1122,6 +1117,21 @@ pub(crate) fn estimate_image_bytes(image_url: &str, detail: Option) } } +/// Image estimate for callers that only have the reference. Original-detail file images use the +/// maximum patch count because their dimensions are not available from the reference. +pub(crate) fn estimate_image_reference_bytes( + image: &ImageReference, + detail: Option, +) -> i64 { + match image { + ImageReference::Inline { image_url } => estimate_image_bytes(image_url, detail), + ImageReference::File { .. } if detail == Some(ImageDetail::Original) => { + i64::try_from(approx_bytes_for_tokens(ORIGINAL_IMAGE_MAX_PATCHES)).unwrap_or(i64::MAX) + } + ImageReference::File { .. } => RESIZED_IMAGE_BYTES_ESTIMATE, + } +} + fn estimate_audio_bytes(audio_url: &str) -> i64 { i64::try_from(approx_bytes_for_tokens(estimate_audio_token_count( audio_url, @@ -1136,14 +1146,9 @@ fn estimate_function_output_bytes(output: &FunctionCallOutputBody) -> i64 { .iter() .map(|part| match part { FunctionCallOutputContentItem::InputText { text } => text_bytes(text), - FunctionCallOutputContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => estimate_image_bytes(image_url, *detail), - FunctionCallOutputContentItem::InputImage { - image: ImageReference::File { .. }, - .. - } => 0, + FunctionCallOutputContentItem::InputImage { image, detail } => { + estimate_image_reference_bytes(image, *detail) + } FunctionCallOutputContentItem::InputAudio { audio_url } => { estimate_audio_bytes(audio_url) } diff --git a/codex-rs/core/src/context_manager/history_tests.rs b/codex-rs/core/src/context_manager/history_tests.rs index a145ee20fb..704b8a56cf 100644 --- a/codex-rs/core/src/context_manager/history_tests.rs +++ b/codex-rs/core/src/context_manager/history_tests.rs @@ -2577,8 +2577,39 @@ fn image_data_url_payload_does_not_dominate_message_estimate() { assert!(estimated > text_only_estimated); } +/// File images use the fixed estimate for normal detail and the maximum patch count for original. #[test] -fn image_data_url_payload_does_not_dominate_function_call_output_estimate() { +fn file_images_use_detail_appropriate_estimates() { + let item = ResponseItem::Message { + id: None, + role: "user".to_string(), + content: vec![ + ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_high".to_string(), + }, + detail: Some(ImageDetail::High), + }, + ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_original".to_string(), + }, + detail: Some(ImageDetail::Original), + }, + ], + phase: None, + internal_chat_message_metadata_passthrough: None, + }; + + let estimated = estimate_response_item_model_visible_bytes(&item); + let expected = RESIZED_IMAGE_BYTES_ESTIMATE + .saturating_add(approx_bytes_for_tokens(/*tokens*/ 10_000) as i64); + + assert_eq!(estimated, expected); +} + +#[test] +fn function_call_output_estimates_inline_and_file_images() { let payload = "B".repeat(50_000); let image_url = format!("data:image/png;base64,{payload}"); let item = ResponseItem::FunctionCallOutput { @@ -2594,14 +2625,22 @@ fn image_data_url_payload_does_not_dominate_function_call_output_estimate() { image: ImageReference::Inline { image_url }, detail: Some(DEFAULT_IMAGE_DETAIL), }, + FunctionCallOutputContentItem::InputImage { + image: ImageReference::File { + file_id: "file_original".to_string(), + }, + detail: Some(ImageDetail::Original), + }, ]), internal_chat_message_metadata_passthrough: None, }; let raw_len = serde_json::to_string(&item).unwrap().len() as i64; let estimated = estimate_response_item_model_visible_bytes(&item); - let expected = - "call-abc".len() as i64 + "Screenshot captured".len() as i64 + RESIZED_IMAGE_BYTES_ESTIMATE; + let expected = "call-abc".len() as i64 + + "Screenshot captured".len() as i64 + + RESIZED_IMAGE_BYTES_ESTIMATE + + approx_bytes_for_tokens(/*tokens*/ 10_000) as i64; assert_eq!(estimated, expected); assert!(estimated < raw_len); diff --git a/codex-rs/core/src/context_manager/mod.rs b/codex-rs/core/src/context_manager/mod.rs index af520df6cf..8f3a347d20 100644 --- a/codex-rs/core/src/context_manager/mod.rs +++ b/codex-rs/core/src/context_manager/mod.rs @@ -4,6 +4,6 @@ pub(crate) mod updates; pub(crate) use history::ContextManager; pub(crate) use history::HistoryReplacement; -pub(crate) use history::estimate_image_bytes; +pub(crate) use history::estimate_image_reference_bytes; pub(crate) use history::estimate_item_token_count; pub(crate) use history::is_user_turn_boundary; diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 7179aa3175..d176191972 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -508,20 +508,26 @@ async fn run_review_on_session( .any(|(_, cost)| cost.image_count > 0) { let reviewer_history = review_session.session.clone_history().await; - let reviewer_image_urls = reviewer_history - .raw_items() - .flat_map(|item| match item { - ResponseItem::Message { content, .. } => content.as_slice(), - _ => &[], - }) - .filter_map(|item| match item { - ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - .. - } => Some(image_url.as_str()), - _ => None, - }) - .collect::>(); + let mut reviewer_image_urls = HashSet::new(); + let mut reviewer_file_ids = HashSet::new(); + for item in reviewer_history.raw_items() { + let ResponseItem::Message { content, .. } = item else { + continue; + }; + for item in content { + let ContentItem::InputImage { image, .. } = item else { + continue; + }; + match image { + ImageReference::Inline { image_url } => { + reviewer_image_urls.insert(image_url.as_str()); + } + ImageReference::File { file_id } => { + reviewer_file_ids.insert(file_id.as_str()); + } + } + } + } let context_window = model_info.resolved_context_window().map(|supported| { params .spawn_config @@ -542,7 +548,7 @@ async fn run_review_on_session( } else { ImagePreparationMode::DetailBased }; - prompt_items.context.retain_images(|image_url, detail| { + prompt_items.context.retain_images(|image, detail| { *detail = match normalize_output_image_detail(&model_info, *detail) { _ if mode == ImagePreparationMode::UnifiedBudget => { Some(ImageDetail::Original) @@ -550,15 +556,23 @@ async fn run_review_on_session( Some(ImageDetail::Low) => Some(ImageDetail::High), detail => detail, }; - let prepared_image_url = match resize_image(image_url, detail, mode) { - Ok(Some(prepared)) => Cow::Owned(prepared.into_data_url()), - Ok(None) => Cow::Borrowed(image_url), - Err(error) => { - warn!(%error, "failed to prepare guardian review image"); - return false; + match image { + ImageReference::Inline { image_url } => { + let prepared_image_url = match resize_image(image_url, detail, mode) + { + Ok(Some(prepared)) => Cow::Owned(prepared.into_data_url()), + Ok(None) => Cow::Borrowed(image_url), + Err(error) => { + warn!(%error, "failed to prepare guardian review image"); + return false; + } + }; + !reviewer_image_urls.contains(prepared_image_url.as_str()) } - }; - !reviewer_image_urls.contains(prepared_image_url.as_ref()) + ImageReference::File { file_id } => { + !reviewer_file_ids.contains(file_id.as_str()) + } + } }); let prompt: ResponseItem = ResponseInputItem::from(prompt_items.context.clone().into_user_inputs()?) diff --git a/codex-rs/core/tests/suite/compact_remote.rs b/codex-rs/core/tests/suite/compact_remote.rs index 82ed5f5092..f7dc16e13d 100644 --- a/codex-rs/core/tests/suite/compact_remote.rs +++ b/codex-rs/core/tests/suite/compact_remote.rs @@ -605,11 +605,14 @@ async fn amazon_bedrock_automatic_compaction_uses_v2_responses_endpoint() -> Res Ok(()) } -#[test_case(None; "default_trims_images")] -#[test_case(Some(false); "disabled_preserves_images")] +#[test_case(None, "image_url"; "default_trims_images")] +#[test_case(Some(false), "image_url"; "disabled_preserves_images")] +#[test_case(None, "file_id"; "default_trims_file_images")] +#[test_case(Some(false), "file_id"; "disabled_preserves_file_images")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn remote_compact_v2_charges_retained_images_to_token_budget( image_budget_enabled: Option, + image_field: &str, ) -> Result<()> { skip_if_no_network!(Ok(())); @@ -630,6 +633,14 @@ async fn remote_compact_v2_charges_retained_images_to_token_budget( // Each original-detail image costs 10,000 estimated patch tokens. let image_inputs = (1..=8) .map(|number| { + if image_field == "file_id" { + return Ok(UserInput::Image { + image: ImageReference::File { + file_id: format!("file_{number}"), + }, + detail: Some(codex_protocol::models::ImageDetail::Original), + }); + } let image = image::ImageBuffer::from_pixel( /*width*/ 3200, /*height*/ 3200, @@ -666,7 +677,24 @@ async fn remote_compact_v2_charges_retained_images_to_token_budget( .await?; wait_for_turn_complete(codex).await; let initial_request = initial_mock.single_request(); - let prepared_images = initial_request.message_input_image_urls("user"); + let prepared_images = initial_request + .inputs_of_type("message") + .into_iter() + .filter(|item| item["role"] == "user") + .flat_map(|item| { + item["content"] + .as_array() + .expect("message content is an array") + .clone() + }) + .filter(|item| item["type"] == "input_image") + .map(|item| { + item[image_field] + .as_str() + .expect("image input has a string reference") + .to_owned() + }) + .collect::>(); assert_eq!(prepared_images.len(), 7); for cycle in 1..=2 { @@ -717,16 +745,33 @@ async fn remote_compact_v2_charges_retained_images_to_token_budget( }; let mut expected_images = prepared_images[dropped..].to_vec(); if cycle == 2 { - let UserInput::Image { - image: ImageReference::Inline { image_url }, - .. - } = &image_inputs[7] - else { + let UserInput::Image { image, .. } = &image_inputs[7] else { unreachable!() }; - expected_images.push(image_url.clone()); + expected_images.push(match image { + ImageReference::Inline { image_url } => image_url.clone(), + ImageReference::File { file_id } => file_id.clone(), + }); } - assert_eq!(follow_up.message_input_image_urls("user"), expected_images); + let retained_images = follow_up + .inputs_of_type("message") + .into_iter() + .filter(|item| item["role"] == "user") + .flat_map(|item| { + item["content"] + .as_array() + .expect("message content is an array") + .clone() + }) + .filter(|item| item["type"] == "input_image") + .map(|item| { + item[image_field] + .as_str() + .expect("image input has a string reference") + .to_owned() + }) + .collect::>(); + assert_eq!(retained_images, expected_images); assert!( follow_up .message_input_texts("user") diff --git a/codex-rs/core/tests/suite/guardian_context_budget.rs b/codex-rs/core/tests/suite/guardian_context_budget.rs index 29be5e7a59..b2c658d730 100644 --- a/codex-rs/core/tests/suite/guardian_context_budget.rs +++ b/codex-rs/core/tests/suite/guardian_context_budget.rs @@ -30,8 +30,11 @@ use core_test_support::test_codex::test_codex; use core_test_support::wait_for_event; use pretty_assertions::assert_eq; use serde_json::json; +use std::sync::Arc; use test_case::test_case; +use super::image_rollout::RecordingFileAttachmentStore; + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_case(160_000, &[0, 0]; "complete_instructions_fit")] #[test_case(28_000, &[0, 1]; "oversized_history_compacts_then_truncates")] @@ -198,6 +201,7 @@ async fn review_preserves_user_instructions_until_request_budgeting( enum ReviewerResponse { Decision, ToolContinuation, + FileImageContinuation, UncompactableContinuation, CompactionError, NextReview, @@ -206,6 +210,7 @@ enum ReviewerResponse { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[test_case(1, ReviewerResponse::Decision; "required_context_fails_closed")] #[test_case(4_500, ReviewerResponse::ToolContinuation; "oversized_tool_continuation_compacts")] +#[test_case(4_500, ReviewerResponse::FileImageContinuation; "uploaded_original_image_history_compacts")] #[test_case(4_500, ReviewerResponse::UncompactableContinuation; "ineffective_compaction_fails_closed")] #[test_case(4_500, ReviewerResponse::CompactionError; "compaction_service_error_does_not_request_user_approval")] #[test_case(6_000, ReviewerResponse::NextReview; "incoming_review_compacts_existing_history")] @@ -231,6 +236,10 @@ async fn review_respects_complete_context_budget( ) .with_model_info_override("gpt-5.6-luna", move |model| { model.context_window = Some(window); + if matches!(reviewer_response, ReviewerResponse::FileImageContinuation) { + model.supports_image_detail_original = true; + model.use_responses_lite = false; + } model .model_messages .as_mut() @@ -268,12 +277,17 @@ async fn review_respects_complete_context_budget( if matches!( reviewer_response, ReviewerResponse::ToolContinuation + | ReviewerResponse::FileImageContinuation | ReviewerResponse::UncompactableContinuation | ReviewerResponse::CompactionError ) { builder = builder .with_code_mode_host_program(codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?); } + let image_store = Arc::new(RecordingFileAttachmentStore::default()); + if matches!(reviewer_response, ReviewerResponse::FileImageContinuation) { + builder = builder.with_image_store(image_store.clone()); + } let test = builder.build_with_auto_env(&server).await?; let command = json!({"cmd": "echo guardian-budget-test", "sandbox_permissions": "require_escalated", "justification": "Run the requested command."}).to_string(); let mut commentary = @@ -321,12 +335,22 @@ async fn review_respects_complete_context_budget( "exec", "text('inspection-output'.repeat(600));", ), + // A tiny inline image fits before upload. Its opaque original-detail file + // reference must reserve 10k tokens in reviewer history and force compaction. + ReviewerResponse::FileImageContinuation => ev_custom_tool_call( + "reviewer-inspect", + "exec", + r#"image("data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==", "original");"#, + ), }, ev_completed("review"), ]), ); } - if matches!(reviewer_response, ReviewerResponse::ToolContinuation) { + if matches!( + reviewer_response, + ReviewerResponse::ToolContinuation | ReviewerResponse::FileImageContinuation + ) { events.insert( /*index*/ 2, sse(vec![ @@ -424,18 +448,38 @@ async fn review_respects_complete_context_budget( "expected required-evidence budget rejection: {output}" ); } else { - let recovered = matches!(reviewer_response, ReviewerResponse::ToolContinuation); + let recovered = matches!( + reviewer_response, + ReviewerResponse::ToolContinuation | ReviewerResponse::FileImageContinuation + ); assert_eq!(requests.len(), if recovered { 4 } else { 3 }); assert_eq!(guardian_requests.len(), if recovered { 2 } else { 1 }); if recovered { assert_eq!(compact_requests.len(), 1); let compact = &compact_requests[0]; - assert!( - compact - .body_json() - .to_string() - .contains("inspection-output") - ); + if matches!(reviewer_response, ReviewerResponse::FileImageContinuation) { + assert_eq!( + image_store + .uploads + .lock() + .expect("image upload tracker lock is not poisoned") + .len(), + 1 + ); + // The file reservation also protects the compaction request itself: an output + // larger than its window is replaced before the summary request is sent. + assert_eq!( + compact.custom_tool_call_output("reviewer-inspect")["output"], + "Output exceeded the available model context and was truncated" + ); + } else { + assert!( + compact + .body_json() + .to_string() + .contains("inspection-output") + ); + } let recovered = guardian_requests[1].body_json(); assert!( recovered["input"] diff --git a/codex-rs/core/tests/suite/image_rollout.rs b/codex-rs/core/tests/suite/image_rollout.rs index b30ad11aa9..fb749dc813 100644 --- a/codex-rs/core/tests/suite/image_rollout.rs +++ b/codex-rs/core/tests/suite/image_rollout.rs @@ -57,8 +57,8 @@ use std::time::Duration; const UPLOADED_FILE_ID: &str = "file_uploaded_image"; #[derive(Default)] -struct RecordingFileAttachmentStore { - uploads: Mutex>, +pub(super) struct RecordingFileAttachmentStore { + pub(super) uploads: Mutex>, } impl AttachmentStore for RecordingFileAttachmentStore { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 2f4c531b29..c2454b8008 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -1677,6 +1677,8 @@ async fn contributor_includes_transcript_images_by_default() -> Result<()> { let user_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGPgEpEDAABoAD1UCKP3AAAAAElFTkSuQmCC"; let tool_image = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGOQE+ECAACQAD304kFaAAAAAElFTkSuQmCC"; + let user_file_id = "file_user"; + let tool_file_id = "file_tool"; let history = vec![ ResponseItem::Message { id: None, @@ -1691,6 +1693,12 @@ async fn contributor_includes_transcript_images_by_default() -> Result<()> { }, detail: Some(ImageDetail::High), }, + ContentItem::InputImage { + image: ImageReference::File { + file_id: user_file_id.to_owned(), + }, + detail: Some(ImageDetail::High), + }, ], phase: None, internal_chat_message_metadata_passthrough: None, @@ -1719,6 +1727,12 @@ async fn contributor_includes_transcript_images_by_default() -> Result<()> { }, detail: Some(ImageDetail::High), }, + FunctionCallOutputContentItem::InputImage { + image: ImageReference::File { + file_id: tool_file_id.to_owned(), + }, + detail: Some(ImageDetail::High), + }, ]), internal_chat_message_metadata_passthrough: None, }, @@ -1740,16 +1754,24 @@ enabled = true .expect("Luna user content should be an array"); assert_eq!( - content[content.len() - 2..], + content[content.len() - 4..], [ json!({ "type": "input_image", "image_url": user_image, }), + json!({ + "type": "input_image", + "file_id": user_file_id, + }), json!({ "type": "input_image", "image_url": tool_image, }), + json!({ + "type": "input_image", + "file_id": tool_file_id, + }), ] ); diff --git a/codex-rs/guardian-context/src/budget.rs b/codex-rs/guardian-context/src/budget.rs index 1ebec846e4..06db2a906d 100644 --- a/codex-rs/guardian-context/src/budget.rs +++ b/codex-rs/guardian-context/src/budget.rs @@ -108,7 +108,9 @@ impl SectionCost { ContentItem::InputImage { image: ImageReference::File { .. }, .. - } => {} + } => { + self.image_count = self.image_count.saturating_add(1); + } ContentItem::InputAudio { audio_url } => { // Guardian currently has no audio contributor. Count a future opaque // payload conservatively until its consumer supplies modality costs. @@ -214,17 +216,16 @@ pub(super) fn content_framing_tokens(item_count: usize) -> usize { fn adjusted_tokens(mut bytes: usize, content: &[ContentItem]) -> usize { for item in content { - if let ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - .. - } = item - { - let payload = ByteCount::measure(|counter| serde_json::to_writer(counter, image_url)); - if payload == usize::MAX { - return usize::MAX; + if let ContentItem::InputImage { image, .. } = item { + if let ImageReference::Inline { image_url } = image { + let payload = + ByteCount::measure(|counter| serde_json::to_writer(counter, image_url)); + if payload == usize::MAX { + return usize::MAX; + } + bytes = bytes.saturating_sub(payload.saturating_sub(2)); } bytes = bytes - .saturating_sub(payload.saturating_sub(2)) .saturating_add(TruncationPolicy::Tokens(IMAGE_TOKEN_RESERVATION).byte_budget()); } } diff --git a/codex-rs/guardian-context/src/budget_tests.rs b/codex-rs/guardian-context/src/budget_tests.rs index 33fc174ad7..8ce4fbb4a6 100644 --- a/codex-rs/guardian-context/src/budget_tests.rs +++ b/codex-rs/guardian-context/src/budget_tests.rs @@ -23,6 +23,12 @@ fn section_costs_keep_multimodal_payloads_separate() { }, detail: None, }), + Budgeted::required(ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_123".to_owned(), + }, + detail: None, + }), ]), }, SectionOutput { @@ -44,7 +50,7 @@ fn section_costs_keep_multimodal_payloads_separate() { SectionCost { text_bytes: "évidence".len(), image_bytes: "data:image/png;base64,AAAA".len(), - image_count: 1 + image_count: 2 } ), ( @@ -71,6 +77,14 @@ fn request_estimate_reserves_images_independently_of_encoded_size() { let short = estimate_input_tokens(&image("AAAA")); assert_eq!(short, estimate_input_tokens(&image(&"A".repeat(200_000)))); assert!(short >= IMAGE_TOKEN_RESERVATION); + + let file = message(vec![ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_123".to_owned(), + }, + detail: Some(codex_protocol::models::ImageDetail::Original), + }]); + assert!(estimate_input_tokens(&file) >= IMAGE_TOKEN_RESERVATION); } #[test] diff --git a/codex-rs/guardian-context/src/composition.rs b/codex-rs/guardian-context/src/composition.rs index 66892875fd..25450d989d 100644 --- a/codex-rs/guardian-context/src/composition.rs +++ b/codex-rs/guardian-context/src/composition.rs @@ -292,17 +292,7 @@ impl ComposedContext { })); continue; } - ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => UserInput::Image { - image: ImageReference::Inline { image_url }, - detail, - }, - ContentItem::InputImage { - image: ImageReference::File { .. }, - .. - } => continue, + ContentItem::InputImage { image, detail } => UserInput::Image { image, detail }, ContentItem::InputAudio { .. } | ContentItem::OutputText { .. } => { return Err(SectionError::UnsupportedDelivery { section: section.id, @@ -330,10 +320,6 @@ impl ComposedContext { } })); } - ContentItem::InputImage { - image: ImageReference::File { .. }, - .. - } => {} content => user_content.push(content), } } diff --git a/codex-rs/guardian-context/src/composition_tests.rs b/codex-rs/guardian-context/src/composition_tests.rs index 4095f00748..ed0fdf41c1 100644 --- a/codex-rs/guardian-context/src/composition_tests.rs +++ b/codex-rs/guardian-context/src/composition_tests.rs @@ -1,4 +1,6 @@ use super::*; +use codex_protocol::models::ImageDetail; +use codex_protocol::models::ImageReference; use pretty_assertions::assert_eq; #[test] @@ -56,12 +58,21 @@ fn delivery_preserves_arbitrary_message_boundaries_and_rejects_them_for_sync() { } #[test] -fn long_text_delivery_is_lossless_bounded_and_fully_budgeted() { +fn long_text_and_file_image_delivery_is_lossless_bounded_and_fully_budgeted() { let text = "é🙂\"\n".repeat(/*n*/ 20_000); + let file_image = ImageReference::File { + file_id: "file_123".to_owned(), + }; let context = ComposedContext { sections: vec![SectionOutput { id: "planned_action", - delivery: text_content(vec![text.clone()]), + delivery: SectionDelivery::UserContent(vec![ + Budgeted::required(ContentItem::InputText { text: text.clone() }), + Budgeted::required(ContentItem::InputImage { + image: file_image.clone(), + detail: Some(ImageDetail::High), + }), + ]), }], truncations: Vec::new(), }; @@ -70,7 +81,15 @@ fn long_text_delivery_is_lossless_bounded_and_fully_budgeted() { let ResponseItem::Message { content, .. } = &messages[0] else { panic!("user message") }; - let parts = content + let (image, text_content) = content.split_last().expect("message content"); + assert_eq!( + image, + &ContentItem::InputImage { + image: file_image.clone(), + detail: Some(ImageDetail::High), + } + ); + let parts = text_content .iter() .map(|item| match item { ContentItem::InputText { text } => text.as_str(), @@ -84,14 +103,16 @@ fn long_text_delivery_is_lossless_bounded_and_fully_budgeted() { .all(|part| part.len() <= TruncationPolicy::Tokens(9_000).byte_budget()) ); assert!(estimated >= crate::estimate_input_tokens(&messages[0])); - assert_eq!( - context.into_user_inputs().unwrap(), - parts - .into_iter() - .map(|part| UserInput::Text { - text: part.to_owned(), - text_elements: Vec::new(), - }) - .collect::>() - ); + let mut expected_inputs = parts + .into_iter() + .map(|part| UserInput::Text { + text: part.to_owned(), + text_elements: Vec::new(), + }) + .collect::>(); + expected_inputs.push(UserInput::Image { + image: file_image, + detail: Some(ImageDetail::High), + }); + assert_eq!(context.into_user_inputs().unwrap(), expected_inputs); } diff --git a/codex-rs/guardian-context/src/enforcement.rs b/codex-rs/guardian-context/src/enforcement.rs index 4e59d73463..234f0d2fbd 100644 --- a/codex-rs/guardian-context/src/enforcement.rs +++ b/codex-rs/guardian-context/src/enforcement.rs @@ -88,17 +88,13 @@ pub enum HistoryTruncation { impl ComposedContext { /// Applies host image admission before aggregate selection, preserving section /// identity and each retained item's selection policy. - pub fn retain_images(&mut self, mut admit: impl FnMut(&str, &mut Option) -> bool) { + pub fn retain_images( + &mut self, + mut admit: impl FnMut(&ImageReference, &mut Option) -> bool, + ) { for section in &mut self.sections { retain_content(section, &mut self.truncations, |_, item| match item { - ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => admit(image_url, detail), - ContentItem::InputImage { - image: ImageReference::File { .. }, - .. - } => false, + ContentItem::InputImage { image, detail } => admit(image, detail), _ => true, }); } diff --git a/codex-rs/guardian-context/src/enforcement_tests.rs b/codex-rs/guardian-context/src/enforcement_tests.rs index 88599eb177..4c3045753b 100644 --- a/codex-rs/guardian-context/src/enforcement_tests.rs +++ b/codex-rs/guardian-context/src/enforcement_tests.rs @@ -236,8 +236,14 @@ fn budget_reserves_existing_context_and_preserves_required_messages() { } #[test] -fn image_omission_preserves_text_and_later_eviction_policy() { +fn image_accounting_preserves_later_eviction_policy() { let evidence = text(&"optional commentary ".repeat(/*n*/ 100)); + let file_image = ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_123".to_owned(), + }, + detail: None, + }; let mut context = ComposedContext { sections: vec![SectionOutput { id: "evidence", @@ -251,6 +257,7 @@ fn image_omission_preserves_text_and_later_eviction_policy() { }, BudgetPriority::Image, ), + Budgeted::optional(file_image.clone(), BudgetPriority::Image), Budgeted::optional(evidence.clone(), BudgetPriority::Commentary), Budgeted::required(text("user restriction")), ]), @@ -261,7 +268,7 @@ fn image_omission_preserves_text_and_later_eviction_policy() { .clone() .enforce_budget( RequestBudget { - max_input_tokens: 1_000, + max_input_tokens: content_tokens(&file_image).saturating_add(1_000), existing_context_tokens: 0, }, "evidence omitted".to_owned(), @@ -271,12 +278,12 @@ fn image_omission_preserves_text_and_later_eviction_policy() { assert_eq!( without_oversized_image.into_messages(), vec![user_message(vec![ - evidence.clone(), + file_image.clone(), text("user restriction"), text("evidence omitted") ])] ); - context.retain_images(|_, _| false); + context.retain_images(|image, _| matches!(image, ImageReference::File { .. })); let available = context.estimated_tokens(); let retained = context .clone() @@ -291,12 +298,16 @@ fn image_omission_preserves_text_and_later_eviction_policy() { .unwrap(); assert_eq!( retained.into_messages(), - vec![user_message(vec![evidence, text("user restriction")])] + vec![user_message(vec![ + file_image.clone(), + evidence, + text("user restriction") + ])] ); let smaller = context .enforce_budget( RequestBudget { - max_input_tokens: 100, + max_input_tokens: content_tokens(&file_image).saturating_add(100), existing_context_tokens: 0, }, "evidence omitted".to_owned(), @@ -306,6 +317,7 @@ fn image_omission_preserves_text_and_later_eviction_policy() { assert_eq!( smaller.into_messages(), vec![user_message(vec![ + file_image, text("user restriction"), text("evidence omitted") ])] diff --git a/codex-rs/guardian-context/src/images.rs b/codex-rs/guardian-context/src/images.rs index 27314df4db..507e688215 100644 --- a/codex-rs/guardian-context/src/images.rs +++ b/codex-rs/guardian-context/src/images.rs @@ -1,5 +1,6 @@ //! Bounded image selection shared by Guardian consumers. //! Keeps source order and evicts oldest images using the existing count/byte caps. +//! Byte caps cover inline URLs and file ID strings, not referenced file contents. //! Consumers still choose source visibility and model-specific image detail. use crate::ContextSection; @@ -29,6 +30,7 @@ pub struct TranscriptImageInput<'a> { #[derive(Clone, Default, PartialEq)] pub struct TranscriptImages { pub images: Vec, + /// Omitted inline payload or file ID bytes, not the size of referenced files. pub omitted_bytes: usize, } @@ -56,37 +58,37 @@ impl TranscriptImages { } let mut images = VecDeque::new(); - let mut image_bytes = 0usize; + let mut image_reference_bytes = 0usize; let mut omitted_bytes = 0usize; - let mut include_image = |image_url: &str, detail: Option| { - if image_url.len() > MAX_TRANSCRIPT_IMAGE_BYTES { - omitted_bytes = omitted_bytes.saturating_add(image_url.len()); + let mut include_image = |image: &ImageReference, detail: Option| { + let reference_bytes = match image { + ImageReference::Inline { image_url } => image_url.len(), + ImageReference::File { file_id } => file_id.len(), + }; + if reference_bytes > MAX_TRANSCRIPT_IMAGE_BYTES { + omitted_bytes = omitted_bytes.saturating_add(reference_bytes); return; } while images.len() >= MAX_TRANSCRIPT_IMAGES - || image_bytes + image_url.len() > MAX_TRANSCRIPT_IMAGE_BYTES + || image_reference_bytes + reference_bytes > MAX_TRANSCRIPT_IMAGE_BYTES { - let Some(ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - .. - }) = images.pop_front() - else { + let Some(ContentItem::InputImage { image, .. }) = images.pop_front() else { break; }; - image_bytes -= image_url.len(); - omitted_bytes = omitted_bytes.saturating_add(image_url.len()); + let evicted_bytes = match image { + ImageReference::Inline { image_url } => image_url.len(), + ImageReference::File { file_id } => file_id.len(), + }; + image_reference_bytes = image_reference_bytes.saturating_sub(evicted_bytes); + omitted_bytes = omitted_bytes.saturating_add(evicted_bytes); } - image_bytes += image_url.len(); + image_reference_bytes = image_reference_bytes.saturating_add(reference_bytes); images.push_back(ContentItem::InputImage { - image: ImageReference::Inline { - image_url: image_url.to_owned(), - }, + image: image.clone(), detail, }); }; - // TODO(kc) Include file-backed images once Guardian can admit them without resolving or - // applying byte-based checks in Core. for item in items { match item { ResponseItem::Message { role, content, .. } @@ -94,15 +96,10 @@ impl TranscriptImages { { for item in content { match item { - ContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => include_image(image_url, *detail), - ContentItem::InputImage { - image: ImageReference::File { .. }, - .. + ContentItem::InputImage { image, detail } => { + include_image(image, *detail) } - | ContentItem::InputText { .. } + ContentItem::InputText { .. } | ContentItem::InputAudio { .. } | ContentItem::OutputText { .. } => {} } @@ -115,15 +112,10 @@ impl TranscriptImages { if let Some(content) = output.content_items() { for item in content { match item { - FunctionCallOutputContentItem::InputImage { - image: ImageReference::Inline { image_url }, - detail, - } => include_image(image_url, *detail), - FunctionCallOutputContentItem::InputImage { - image: ImageReference::File { .. }, - .. + FunctionCallOutputContentItem::InputImage { image, detail } => { + include_image(image, *detail) } - | FunctionCallOutputContentItem::InputText { .. } + FunctionCallOutputContentItem::InputText { .. } | FunctionCallOutputContentItem::InputAudio { .. } | FunctionCallOutputContentItem::EncryptedContent { .. } => {} } @@ -137,9 +129,9 @@ impl TranscriptImages { for image in input.node_repl_images { match image { ContentItem::InputImage { - image: ImageReference::Inline { image_url }, + image: image @ ImageReference::Inline { .. }, detail, - } => include_image(image_url, *detail), + } => include_image(image, *detail), ContentItem::InputImage { image: ImageReference::File { .. }, .. diff --git a/codex-rs/guardian-context/src/images_tests.rs b/codex-rs/guardian-context/src/images_tests.rs index 5c3fa54d04..89361ffe83 100644 --- a/codex-rs/guardian-context/src/images_tests.rs +++ b/codex-rs/guardian-context/src/images_tests.rs @@ -1,6 +1,11 @@ use super::MAX_TRANSCRIPT_IMAGE_BYTES; use super::TranscriptImageInput; use super::TranscriptImages; +use crate::CollectedContext; +use crate::ContextPresentation; +use crate::ContextSection; +use crate::RenderedTranscript; +use crate::composition::user_message; use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputContentItem; use codex_protocol::models::FunctionCallOutputPayload; @@ -99,20 +104,26 @@ fn image_selection_preserves_source_policy_order_and_both_limits() { ); } -/// File-backed images stay out of Guardian without consuming its inline-image capacity. +/// File-backed history images share count and reference-byte limits with inline images. #[test] -fn file_images_are_omitted_without_displacing_inline_images() { +fn file_images_are_selected_from_history_sources() { let file_image = ContentItem::InputImage { image: ImageReference::File { file_id: "file_123".to_string(), }, detail: Some(ImageDetail::High), }; + let tool_file_image = ContentItem::InputImage { + image: ImageReference::File { + file_id: "file_tool".to_string(), + }, + detail: Some(ImageDetail::High), + }; let history = [ ResponseItem::Message { id: None, role: "user".into(), - content: vec![image("first"), file_image.clone(), image("second")], + content: vec![image("first"), file_image.clone()], phase: None, internal_chat_message_metadata_passthrough: None, }, @@ -138,16 +149,7 @@ fn file_images_are_omitted_without_displacing_inline_images() { internal_chat_message_metadata_passthrough: None, }, ]; - let repl = [ - file_image, - image("fourth"), - ContentItem::InputImage { - image: ImageReference::File { - file_id: "file_repl".to_string(), - }, - detail: Some(ImageDetail::High), - }, - ]; + let repl = [image("fourth")]; assert_eq!( TranscriptImages::collect( @@ -159,8 +161,94 @@ fn file_images_are_omitted_without_displacing_inline_images() { } ), TranscriptImages { - images: ["first", "second", "third", "fourth"].map(image).to_vec(), - omitted_bytes: 0, + images: vec![file_image, tool_file_image, image("third"), image("fourth")], + omitted_bytes: "first".len(), + } + ); + + // Count-only eviction must remain observable even when no inline bytes were omitted. + let file_images = + ["file_1", "file_2", "file_3", "file_4", "file_5"].map(|file_id| ContentItem::InputImage { + image: ImageReference::File { + file_id: file_id.to_owned(), + }, + detail: Some(ImageDetail::High), + }); + let history = [ResponseItem::Message { + id: None, + role: "user".into(), + content: file_images.to_vec(), + phase: None, + internal_chat_message_metadata_passthrough: None, + }]; + let selected = TranscriptImages::collect( + &history, + TranscriptImageInput { + enabled: true, + include_tool_outputs: true, + node_repl_images: &[], + }, + ); + assert_eq!( + selected, + TranscriptImages { + images: file_images[1..].to_vec(), + omitted_bytes: "file_1".len(), + } + ); + let context = CollectedContext { + sections: vec![ContextSection::TranscriptImages(selected)], + } + .compose( + ContextPresentation::Async, + RenderedTranscript { + items: Vec::new(), + omission_note: None, + truncations: Vec::new(), + }, + ) + .unwrap(); + assert_eq!( + context + .truncations + .iter() + .map(|observation| ( + observation.component, + observation.original_bytes, + observation.retained_bytes, + )) + .collect::>(), + vec![("transcript_image", "file_1".len(), 0)] + ); + + // Newer evidence evicts the full-budget ID; the oversized ID must not displace it. + let content = [ + ImageReference::File { + file_id: "x".repeat(MAX_TRANSCRIPT_IMAGE_BYTES), + }, + ImageReference::Inline { + image_url: "recent".into(), + }, + ImageReference::File { + file_id: "y".repeat(MAX_TRANSCRIPT_IMAGE_BYTES + 1), + }, + ] + .map(|image| ContentItem::InputImage { + image, + detail: Some(ImageDetail::High), + }); + assert_eq!( + TranscriptImages::collect( + &[user_message(content.to_vec())], + TranscriptImageInput { + enabled: true, + include_tool_outputs: true, + node_repl_images: &[], + }, + ), + TranscriptImages { + images: vec![image("recent")], + omitted_bytes: 2 * MAX_TRANSCRIPT_IMAGE_BYTES + 1, } ); }