diff --git a/codex-rs/core/src/guardian/input_budget_tests.rs b/codex-rs/core/src/guardian/input_budget_tests.rs index a6f7ce55bd..2f4496bde5 100644 --- a/codex-rs/core/src/guardian/input_budget_tests.rs +++ b/codex-rs/core/src/guardian/input_budget_tests.rs @@ -363,28 +363,8 @@ async fn compacted_review_restores_originals_once_and_persists_the_request_prefi &[], ) .unwrap(); - let elastic = (budgeting == ContextBudgeting::Elastic).then(|| { - let section = codex_guardian_context::ElasticSection::TrustedSkills; - let mut policy = codex_guardian_context::ElasticBudgetConfig::default(); - policy.sections.insert( - section, - codex_guardian_context::ElasticSectionLimits::new( - /*min_total_section_tokens*/ 1, /*max_section_entry_tokens*/ 1, - ) - .unwrap(), - ); - let minimum = codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS - - usize::try_from(policy.total_minimum_tokens()).unwrap() - + 1; - policy.sections.insert( - section, - codex_guardian_context::ElasticSectionLimits::new( - minimum, /*max_section_entry_tokens*/ 1, - ) - .unwrap(), - ); - policy - }); + let elastic = (budgeting == ContextBudgeting::Elastic) + .then(core_test_support::guardian::elastic_budget_config); step.turn.extension_data.insert(RetainedReviewContext { context: context.retained_instructions(), elastic, diff --git a/codex-rs/core/tests/suite/scenarios_guardian_orchestrator_connector.rs b/codex-rs/core/tests/suite/scenarios_guardian_orchestrator_connector.rs index 545d7b2f45..9f8096b48e 100644 --- a/codex-rs/core/tests/suite/scenarios_guardian_orchestrator_connector.rs +++ b/codex-rs/core/tests/suite/scenarios_guardian_orchestrator_connector.rs @@ -3,10 +3,15 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::time::Duration; use anyhow::Result; use codex_core::config::Constrained; use codex_features::Feature; +use codex_guardian_context::ContextBudgeting; +use codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS; +use codex_guardian_context::ElasticSection; +use codex_guardian_context::ElasticSectionLimits; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::protocol::AskForApproval; use core_test_support::apps_test_server::AppsTestServer; @@ -27,8 +32,12 @@ use wiremock::matchers::body_partial_json; use wiremock::matchers::method; use wiremock::matchers::path; +#[test_case::test_case(ContextBudgeting::Fixed; "fixed")] +#[test_case::test_case(ContextBudgeting::Elastic; "elastic_entry_cap")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn host_owned_orchestrator_connector_reaches_guardian_context() -> Result<()> { +async fn host_owned_orchestrator_connector_reaches_guardian_context( + budgeting: ContextBudgeting, +) -> Result<()> { skip_if_no_network!(Ok(())); let server = responses::start_mock_server().await; @@ -48,14 +57,15 @@ async fn host_owned_orchestrator_connector_reaches_guardian_context() -> Result< responses::ev_completed("done"), ]), ]; - Mock::given(method("POST")) + let classifier_mock = Mock::given(method("POST")) .and(path("/v1/responses")) .and(body_partial_json(json!({ "model": "gpt-5.6-luna" }))) .respond_with(responses::sse_response(responses::sse(vec![ responses::ev_assistant_message("score", "high"), responses::ev_completed("score"), ]))) - .mount(&server) + .expect(1..) + .mount_as_scoped(&server) .await; Mock::given(method("POST")) .and(path("/v1/responses")) @@ -98,7 +108,7 @@ async fn host_owned_orchestrator_connector_reaches_guardian_context() -> Result< ) .expect("write Guardian configuration"); }) - .with_config(|config| { + .with_config(move |config| { config.model_provider.supports_websockets = false; config .features @@ -114,12 +124,34 @@ async fn host_owned_orchestrator_connector_reaches_guardian_context() -> Result< .expect("enable Guardian approval"); config.approvals_reviewer = ApprovalsReviewer::AutoReview; config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + if budgeting == ContextBudgeting::Elastic { + const TOOL_ENTRY_CAP: usize = 1; + let mut policy = core_test_support::guardian::elastic_budget_config(); + let limits = |minimum| { + ElasticSectionLimits::new(minimum, TOOL_ENTRY_CAP).expect("valid tool limits") + }; + policy.sections.insert(ElasticSection::TrustedTools, limits(TOOL_ENTRY_CAP)); + let other_minimums = usize::try_from(policy.total_minimum_tokens()) + .expect("token total fits usize") + - TOOL_ENTRY_CAP; + policy.sections.insert( + ElasticSection::TrustedTools, + limits(DEFAULT_MAX_INPUT_TOKENS - other_minimums), + ); + config.guardian_async_elastic_budget = Some(policy); + } }) .build_with_auto_env(&server) .await?; wait_for_mcp_server(&test.codex, "codex_apps").await?; test.submit_text_turn("List today's calendar events.") .await?; + const CLASSIFIER_REQUEST_TIMEOUT: Duration = Duration::from_secs(30); + tokio::time::timeout( + CLASSIFIER_REQUEST_TIMEOUT, + classifier_mock.wait_until_satisfied(), + ) + .await?; assert_eq!(parent_index.load(Ordering::SeqCst), 2); let mut requests = server @@ -149,6 +181,19 @@ async fn host_owned_orchestrator_connector_reaches_guardian_context() -> Result< serde_json::from_str::(evidence) }) .collect::>>()?; + if budgeting == ContextBudgeting::Elastic { + assert!(trusted_messages.is_empty()); + let classifier = requests + .iter() + .find(|request| request["model"] == "gpt-5.6-luna") + .expect("classifier request"); + assert!( + classifier["input"] + .to_string() + .contains("") + ); + return Ok(()); + } assert_eq!( trusted_messages, vec![json!({ diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs index 0bb02fb85d..f86cd8122b 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/transcript.rs @@ -3,7 +3,6 @@ use codex_extension_api::ResponseItem; pub(crate) use codex_features::GuardianV2TranscriptSource as TranscriptSource; use codex_guardian_context::CollectedContext; use codex_guardian_context::ComposedContext; -use codex_guardian_context::ContextBudgeting; use codex_guardian_context::ContextPresentation; use codex_guardian_context::ContextProfile; use codex_guardian_context::ContextTarget; @@ -26,8 +25,7 @@ use codex_guardian_context::TranscriptRetentionConfig; use codex_guardian_context::TranscriptSelection; use codex_guardian_context::TrustedTool; use codex_guardian_context::default_registry; -use codex_guardian_context::prepare_context; -use std::borrow::Cow; +use codex_guardian_context::prepare_async_context; const DEFAULT_PROFILE: &ContextProfile = &ContextProfile::asynchronous(); pub(crate) const MAX_MESSAGE_ENTRY_TOKENS: usize = @@ -185,17 +183,13 @@ impl CollectedTranscript { TranscriptSelection::Full(_) => (0, ContextPresentation::Async), TranscriptSelection::Delta { offset, .. } => (offset, ContextPresentation::AsyncDelta), }; - let mut context = prepare_context( - Cow::Borrowed(&self.context), + let mut context = prepare_async_context( + &self.context, &self.profile, offset, presentation, - if self.elastic.is_some() { - ContextBudgeting::Elastic - } else { - ContextBudgeting::Fixed - }, reviewer_history, + self.elastic.as_ref(), )?; if self.elastic.is_none() { context.retain_new_native_messages(reviewer_history)?; diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index a3620c2560..977f6e93e1 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -108,6 +108,7 @@ pub use preparation::ContextBudgeting; pub use preparation::ContextProfile; pub use preparation::SyncImagePreparer; pub use preparation::SyncPreparation; +pub use preparation::prepare_async_context; pub use preparation::prepare_context; pub use preparation::prepare_sync_context_with_facts; pub use preparation::prepare_sync_delivery; diff --git a/codex-rs/guardian-context/src/preparation/async_entries.rs b/codex-rs/guardian-context/src/preparation/async_entries.rs new file mode 100644 index 0000000000..b5f3069414 --- /dev/null +++ b/codex-rs/guardian-context/src/preparation/async_entries.rs @@ -0,0 +1,76 @@ +//! Caps typed classifier evidence before it becomes trusted developer messages. +//! Delivery metadata must describe only the complete records kept for rendering. + +use crate::CollectedContext; +use crate::ContextSection; +use crate::ElasticBudgetConfig; +use crate::ElasticSection; +use crate::TruncationObservation; +use crate::budget::content_framing_tokens; +use crate::budget::content_tokens; +use crate::trusted_tool::MAX_TRUSTED_TOOL_CONTEXT_TOKENS; +use codex_protocol::models::ContentItem; +use codex_protocol::models::input_text; +use codex_protocol::protocol::TruncationPolicy; + +pub(super) fn apply( + context: &mut CollectedContext, + config: &ElasticBudgetConfig, + truncations: &mut Vec, +) { + let cap = |section| config.section_limits(section).max_section_entry_tokens(); + let mut record_omission = |component, original_bytes| { + truncations.push(TruncationObservation { + source: None, + component, + original_bytes, + retained_bytes: 0, + }); + }; + context.sections.retain_mut(|section| match section { + ContextSection::TrustedTool(tool) => { + let attestation = tool.attestation(); + let bytes = attestation.len(); + let rendering_limit = + TruncationPolicy::Tokens(MAX_TRUSTED_TOOL_CONTEXT_TOKENS).byte_budget(); + let keep = bytes <= rendering_limit + && fits(&input_text(attestation), cap(ElasticSection::TrustedTools)); + if !keep { + record_omission("trusted_tool", bytes); + } + keep + } + ContextSection::TrustedSkills(skills) => { + let maximum = cap(ElasticSection::TrustedSkills); + for path in std::mem::take(&mut skills.paths) { + let path_bytes = path.len(); + skills.paths.push(path); + let (rendered, count) = skills.render_paths(); + if count != skills.paths.len() || !fits(&rendered, maximum) { + drop(skills.paths.pop()); + record_omission("trusted_skills", path_bytes); + } + } + !skills.paths.is_empty() + } + ContextSection::PreviousReviews(reviews) => { + let maximum = cap(ElasticSection::PreviousReviews); + reviews.retain(|review| { + let keep = fits(&input_text(review.fragment.clone()), maximum); + if !keep { + record_omission("previous_reviews", review.fragment.len()); + } + keep + }) + } + _ => true, + }); +} + +fn fits(content: &ContentItem, maximum: usize) -> bool { + content_tokens(content).saturating_add(content_framing_tokens(/*item_count*/ 1)) <= maximum +} + +#[cfg(test)] +#[path = "async_entries_tests.rs"] +mod tests; diff --git a/codex-rs/guardian-context/src/preparation/async_entries_tests.rs b/codex-rs/guardian-context/src/preparation/async_entries_tests.rs new file mode 100644 index 0000000000..98f2976822 --- /dev/null +++ b/codex-rs/guardian-context/src/preparation/async_entries_tests.rs @@ -0,0 +1,143 @@ +#![cfg(test)] + +use super::*; +use crate::ContextPresentation; +use crate::ContextProfile; +use crate::ElasticSectionLimits; +use crate::PreviousReview; +use crate::PreviousReviews; +use crate::RequestBudget; +use crate::TrustedSkills; +use crate::TrustedTool; +use crate::TrustedToolSource; +use codex_protocol::ResponseItemId; +use codex_protocol::models::ResponseItem; +use pretty_assertions::assert_eq; + +#[test] +fn classifier_entry_caps_preserve_whole_evidence_and_only_sent_review_ids() { + const ENTRY_CAP: usize = 100; + const SECTION_MINIMUM: usize = 4_000; + const OVERSIZED_BYTES: usize = 500; + const NOTICE: &str = "Evidence omitted for this request."; + let kept_review = PreviousReview { + id: ResponseItemId::from_server("review_kept".to_owned()), + fragment: "Approved reading README.md.".to_owned(), + }; + let omitted_review = PreviousReview { + id: ResponseItemId::from_server("review_omitted".to_owned()), + fragment: format!("Reviewed: {}", "x".repeat(OVERSIZED_BYTES)), + }; + let reviews = + PreviousReviews::try_from_fragments(vec![omitted_review, kept_review.clone()]).unwrap(); + let tool = TrustedTool { + server: "files".to_owned(), + connector_id: None, + source: TrustedToolSource::PluginServiceOrchestrator, + }; + let skill = "/skills/read-only/SKILL.md"; + let context = CollectedContext { + sections: vec![ + ContextSection::TrustedTool(tool), + ContextSection::TrustedSkills(TrustedSkills { + paths: vec![ + format!("/skills/{}", "x".repeat(OVERSIZED_BYTES)), + skill.to_owned(), + ], + }), + ContextSection::PreviousReviews(reviews), + ], + }; + let profile = ContextProfile::asynchronous(); + let fixed = crate::prepare_async_context( + &context, + &profile, + /*entry_number_offset*/ 0, + ContextPresentation::Async, + &[], + /*elastic*/ None, + ) + .unwrap(); + assert!(fixed.truncations.is_empty()); + let fixed = fixed.into_annotated_messages(); + assert_eq!( + fixed + .iter() + .filter(|item| item + .metadata + .as_ref() + .is_some_and(|metadata| metadata.guardian_review_ids.len() > 1,)) + .count(), + 1 + ); + + let limits = ElasticSectionLimits::new(SECTION_MINIMUM, ENTRY_CAP).unwrap(); + let config = crate::budget::tests::complete_minimum_budget(ElasticBudgetConfig { + sections: [ + ( + ElasticSection::TrustedTools, + ElasticSectionLimits::new(SECTION_MINIMUM, /*max_section_entry_tokens*/ 1).unwrap(), + ), + (ElasticSection::TrustedSkills, limits), + (ElasticSection::PreviousReviews, limits), + ] + .into(), + ..Default::default() + }); + let capped = crate::prepare_async_context( + &context, + &profile, + /*entry_number_offset*/ 0, + ContextPresentation::Async, + &[], + Some(&config), + ) + .unwrap(); + assert_eq!( + capped + .truncations + .iter() + .map(|loss| loss.component) + .collect::>(), + vec!["trusted_tool", "trusted_skills", "previous_reviews"], + ); + let delivered = capped + .prepare_elastic_entries(&config) + .unwrap() + .prepare_delivery(&[]) + .unwrap() + .enforce_budget( + RequestBudget { + max_input_tokens: crate::DEFAULT_MAX_INPUT_TOKENS, + existing_context_tokens: 0, + }, + NOTICE.to_owned(), + ) + .unwrap() + .into_annotated_messages(); + let expected_review = PreviousReviews::try_from_fragments(vec![kept_review]) + .unwrap() + .into_annotated_message(); + assert!(delivered.contains(&expected_review)); + let developers = delivered + .iter() + .filter_map(|envelope| match &envelope.item { + ResponseItem::Message { role, .. } if role == "developer" => Some(&envelope.item), + _ => None, + }) + .collect::>(); + let expected_skill = codex_context_fragments::ContextualUserFragment::into(TrustedSkills { + paths: vec![skill.to_owned()], + }); + assert_eq!(developers, vec![&expected_review.item, &expected_skill]); + assert!( + serde_json::to_string( + &delivered + .iter() + .map(|envelope| &envelope.item) + .collect::>() + ) + .unwrap() + .contains(NOTICE) + ); +} diff --git a/codex-rs/guardian-context/src/preparation/mod.rs b/codex-rs/guardian-context/src/preparation/mod.rs index bc2f356f3a..fe338a9d7f 100644 --- a/codex-rs/guardian-context/src/preparation/mod.rs +++ b/codex-rs/guardian-context/src/preparation/mod.rs @@ -2,6 +2,7 @@ //! Collection keeps captured sources intact. This owner applies text and image policy; //! delivery checks repeat against current reviewer history after compaction. +mod async_entries; mod elastic; pub(crate) mod repl; pub(crate) mod root; @@ -13,6 +14,7 @@ use crate::CollectedContext; use crate::ComposedContext; use crate::ContextPresentation; use crate::ContextSection; +use crate::ElasticBudgetConfig; use crate::GuardianRootMessage; use crate::RetainedInstructionItem; use crate::SectionError; @@ -58,6 +60,51 @@ pub fn prepare_context( presentation: ContextPresentation<'_>, budgeting: ContextBudgeting, reviewer_history: &[ResponseItemEnvelope], +) -> Result { + prepare_context_with_async_caps( + context, + profile, + entry_number_offset, + presentation, + budgeting, + reviewer_history, + /*async_elastic*/ None, + ) +} + +/// Prepares the classifier suffix, applying typed async entry caps before rendering. +pub fn prepare_async_context( + context: &CollectedContext, + profile: &ContextProfile, + entry_number_offset: usize, + presentation: ContextPresentation<'_>, + reviewer_history: &[ResponseItemEnvelope], + elastic: Option<&ElasticBudgetConfig>, +) -> Result { + let budgeting = if elastic.is_some() { + ContextBudgeting::Elastic + } else { + ContextBudgeting::Fixed + }; + prepare_context_with_async_caps( + Cow::Borrowed(context), + profile, + entry_number_offset, + presentation, + budgeting, + reviewer_history, + elastic, + ) +} + +fn prepare_context_with_async_caps( + context: Cow<'_, CollectedContext>, + profile: &ContextProfile, + entry_number_offset: usize, + presentation: ContextPresentation<'_>, + budgeting: ContextBudgeting, + reviewer_history: &[ResponseItemEnvelope], + async_elastic: Option<&ElasticBudgetConfig>, ) -> Result { let (target, placeholder) = match &presentation { ContextPresentation::SyncFull { .. } => ( @@ -113,6 +160,9 @@ pub fn prepare_context( }, }; context.retain_new_reviews(reviewer_history); + if let Some(config) = async_elastic { + async_entries::apply(&mut context, config, &mut transcript.truncations); + } prepare_sections(&mut context, budgeting)?; context.compose_prepared(presentation, transcript) } diff --git a/codex-rs/guardian-context/src/reviews.rs b/codex-rs/guardian-context/src/reviews.rs index 841b666e3d..0d9309d15c 100644 --- a/codex-rs/guardian-context/src/reviews.rs +++ b/codex-rs/guardian-context/src/reviews.rs @@ -57,6 +57,11 @@ impl std::fmt::Debug for PreviousReviews { } impl PreviousReviews { + pub(crate) fn retain(&mut self, predicate: impl FnMut(&PreviousReview) -> bool) -> bool { + self.fragments.retain(predicate); + !self.fragments.is_empty() + } + /// Validates the count and size of complete fragments, including their markers. /// The host remains responsible for provenance and authorization validity. /// Rejects oversized evidence without truncating or dropping any records. @@ -121,13 +126,12 @@ impl CollectedContext { let ContextSection::PreviousReviews(reviews) = section else { return true; }; - reviews.fragments.retain(|review| { + reviews.retain(|review| { !history .iter() .filter_map(|item| item.metadata.as_ref()) .any(|metadata| metadata.guardian_review_ids.contains(&review.id)) - }); - !reviews.fragments.is_empty() + }) }); } } diff --git a/codex-rs/guardian-context/src/trusted_skills.rs b/codex-rs/guardian-context/src/trusted_skills.rs index 13c8053010..9e31034a8a 100644 --- a/codex-rs/guardian-context/src/trusted_skills.rs +++ b/codex-rs/guardian-context/src/trusted_skills.rs @@ -25,6 +25,32 @@ pub struct TrustedSkills { impl TrustedSkills { /// Classification attached to this fragment's rendered text. pub const KIND: &str = "guardian.trusted_skills"; + + pub(crate) fn render_paths(&self) -> (ContentItem, usize) { + let mut paths = String::from("["); + let mut count = 0; + let max_path_bytes = TruncationPolicy::Tokens(MAX_TRUSTED_SKILL_TOKENS) + .byte_budget() + .saturating_sub(TRUSTED_SKILLS_PREFIX.len()); + for path in &self.paths { + let separator = if paths.ends_with('[') { "" } else { "," }; + let encoded_path = json!(path).to_string(); + if paths + .len() + .saturating_add(separator.len()) + .saturating_add(encoded_path.len()) + .saturating_add(1) + > max_path_bytes + { + continue; + } + paths.push_str(separator); + paths.push_str(&encoded_path); + count += 1; + } + paths.push(']'); + (input_text(format!("{TRUSTED_SKILLS_PREFIX}{paths}")), count) + } } impl ContextualUserFragment for TrustedSkills { @@ -45,27 +71,7 @@ impl ContextualUserFragment for TrustedSkills { } fn body(&self) -> ContentItem { - let mut paths = String::from("["); - let max_path_bytes = TruncationPolicy::Tokens(MAX_TRUSTED_SKILL_TOKENS) - .byte_budget() - .saturating_sub(TRUSTED_SKILLS_PREFIX.len()); - for path in &self.paths { - let separator = if paths.ends_with('[') { "" } else { "," }; - let encoded_path = json!(path).to_string(); - if paths - .len() - .saturating_add(separator.len()) - .saturating_add(encoded_path.len()) - .saturating_add(1) - > max_path_bytes - { - continue; - } - paths.push_str(separator); - paths.push_str(&encoded_path); - } - paths.push(']'); - input_text(format!("{TRUSTED_SKILLS_PREFIX}{paths}")) + self.render_paths().0 } } diff --git a/codex-rs/guardian-context/src/trusted_tool.rs b/codex-rs/guardian-context/src/trusted_tool.rs index 3d05e1cf8a..b41149b304 100644 --- a/codex-rs/guardian-context/src/trusted_tool.rs +++ b/codex-rs/guardian-context/src/trusted_tool.rs @@ -15,7 +15,7 @@ use crate::SectionInput; use crate::SectionScope; use crate::truncate_text as truncate_entry; -const MAX_TRUSTED_TOOL_CONTEXT_TOKENS: usize = 512; +pub(crate) const MAX_TRUSTED_TOOL_CONTEXT_TOKENS: usize = 512; const USER_CONFIGURATION_PREFIX: &str = "Codex verified that this exact MCP tool or connector was declared in \ trusted user configuration. "; const PLUGIN_SERVICE_ORCHESTRATOR_PREFIX: &str = "Codex verified that this exact connector was provided by the \ @@ -44,6 +44,26 @@ pub struct TrustedTool { impl TrustedTool { /// Classification attached to this fragment's rendered text. pub const KIND: &str = "guardian.trusted_tool"; + + pub(crate) fn attestation(&self) -> String { + let (prefix, source) = match &self.source { + TrustedToolSource::UserConfiguration(path) => { + (USER_CONFIGURATION_PREFIX, path.as_str()) + } + TrustedToolSource::PluginServiceOrchestrator => ( + PLUGIN_SERVICE_ORCHESTRATOR_PREFIX, + "plugin_service_orchestrator", + ), + }; + format!( + "{prefix}{TRUSTED_TOOL_SCOPE}\n{}", + json!({ + "server": self.server, + "connector_id": self.connector_id, + "source": source, + }) + ) + } } impl ContextualUserFragment for TrustedTool { @@ -64,24 +84,8 @@ impl ContextualUserFragment for TrustedTool { } fn body(&self) -> ContentItem { - let (prefix, source) = match &self.source { - TrustedToolSource::UserConfiguration(path) => { - (USER_CONFIGURATION_PREFIX, path.as_str()) - } - TrustedToolSource::PluginServiceOrchestrator => ( - PLUGIN_SERVICE_ORCHESTRATOR_PREFIX, - "plugin_service_orchestrator", - ), - }; input_text(truncate_entry( - &format!( - "{prefix}{TRUSTED_TOOL_SCOPE}\n{}", - json!({ - "server": self.server, - "connector_id": self.connector_id, - "source": source, - }) - ), + &self.attestation(), MAX_TRUSTED_TOOL_CONTEXT_TOKENS, )) }