diff --git a/codex-rs/core/tests/suite/guardian_context_budget.rs b/codex-rs/core/tests/suite/guardian_context_budget.rs index 721096f1b3..05441e50d1 100644 --- a/codex-rs/core/tests/suite/guardian_context_budget.rs +++ b/codex-rs/core/tests/suite/guardian_context_budget.rs @@ -35,7 +35,7 @@ enum ReviewerResponse { #[test_case(1, ReviewerResponse::Decision; "required_context_fails_closed")] #[test_case(4_500, ReviewerResponse::ToolContinuation; "oversized_tool_continuation_compacts")] #[test_case(4_500, ReviewerResponse::UncompactableContinuation; "ineffective_compaction_fails_closed")] -#[test_case(4_500, ReviewerResponse::NextReview; "incoming_review_compacts_existing_history")] +#[test_case(6_000, ReviewerResponse::NextReview; "incoming_review_compacts_existing_history")] async fn review_respects_complete_context_budget( window: i64, reviewer_response: ReviewerResponse, @@ -69,10 +69,6 @@ async fn review_respects_complete_context_budget( .features .enable(Feature::TokenBudget) .expect("parent uses token-budget mode while Guardian uses summary compaction"); - config - .features - .disable(Feature::RemoteCompactionV2) - .expect("test uses the legacy compact endpoint"); config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); config.approvals_reviewer = ApprovalsReviewer::AutoReview; config @@ -118,12 +114,22 @@ async fn review_respects_complete_context_budget( sse(vec![ ev_response_created("review"), match reviewer_response { - ReviewerResponse::Decision | ReviewerResponse::NextReview => { - ev_assistant_message( - "decision", - r#"{"risk_level":"low","user_authorization":"high","outcome":"allow"}"#, - ) - } + ReviewerResponse::Decision => ev_assistant_message( + "decision", + r#"{"risk_level":"low","user_authorization":"high","outcome":"allow"}"#, + ), + // V2 retains user review inputs. Give it assistant history to + // discard so compaction makes room for the next review. + ReviewerResponse::NextReview => ev_assistant_message( + "decision", + &json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "Previous review reasoning. ".repeat(/*n*/ 256), + }) + .to_string(), + ), ReviewerResponse::ToolContinuation | ReviewerResponse::UncompactableContinuation => ev_custom_tool_call( "reviewer-inspect", @@ -169,9 +175,7 @@ async fn review_respects_complete_context_budget( ]), ]); } - let compact = if matches!(reviewer_response, ReviewerResponse::Decision) { - None - } else { + if !matches!(reviewer_response, ReviewerResponse::Decision) { let summary = if matches!( reviewer_response, ReviewerResponse::UncompactableContinuation @@ -180,20 +184,29 @@ async fn review_respects_complete_context_budget( } else { "Previous review evidence and inspection results.".to_owned() }; - Some( - responses::mount_compact_json_once( - &server, - json!({"output": [{ - "type": "compaction", "encrypted_content": summary - }]}), - ) - .await, - ) - }; + let index = if matches!(reviewer_response, ReviewerResponse::NextReview) { + 4 + } else { + 2 + }; + events.insert( + index, + sse(vec![ + json!({ + "type": "response.output_item.done", + "item": {"type": "compaction", "encrypted_content": summary}, + }), + ev_completed("review-compaction"), + ]), + ); + } let responses = responses::mount_sse_sequence(&server, events).await; test.submit_text_turn("Run the command if the approval reviewer allows it.") .await?; - let requests = responses.requests(); + let (compact_requests, requests): (Vec<_>, Vec<_>) = responses + .requests() + .into_iter() + .partition(|request| !request.inputs_of_type("compaction_trigger").is_empty()); let guardian_requests = requests .iter() .filter(|request| { @@ -215,10 +228,8 @@ async fn review_respects_complete_context_budget( assert_eq!(requests.len(), if recovered { 4 } else { 3 }); assert_eq!(guardian_requests.len(), if recovered { 2 } else { 1 }); if recovered { - let compact = compact - .as_ref() - .expect("scenario includes compaction") - .single_request(); + assert_eq!(compact_requests.len(), 1); + let compact = &compact_requests[0]; assert!( compact .body_json() @@ -279,13 +290,22 @@ async fn review_respects_complete_context_budget( "Retry the requested command.".to_owned() }; test.submit_text_turn(&instruction).await?; - let requests = responses.requests(); - assert_eq!(requests.len(), 6); - let compact = compact - .as_ref() - .map(responses::ResponseMock::single_request); + let (compact_requests, requests): (Vec<_>, Vec<_>) = responses + .requests() + .into_iter() + .partition(|request| !request.inputs_of_type("compaction_trigger").is_empty()); + assert_eq!( + requests.len(), + 6, + "parent resumed with: {}", + requests + .last() + .expect("parent resumes after the retry") + .function_call_output("retry-command") + ); + assert_eq!(compact_requests.len(), 1); if matches!(reviewer_response, ReviewerResponse::NextReview) { - let compact = compact.expect("scenario includes compaction"); + let compact = &compact_requests[0]; assert!( !compact .body_json() diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs new file mode 100644 index 0000000000..6442089263 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs @@ -0,0 +1,219 @@ +//! Contributor integration tests use the normal catalog-backed lifecycle setup. + +use super::*; +use pretty_assertions::assert_eq; + +#[derive(Clone, Copy)] +enum BudgetOutcome { + Fits, + RequiresSync, +} + +async fn catalog_budget_fixture(base_url: String, window: i64) -> Result { + let server = responses::start_mock_server().await; + let test = test_codex() + .with_model_info_override(MODEL, move |model| { + model.context_window = Some(window); + model.effective_context_window_percent = 50; + model.comp_hash = Some("budget-checkpoint".to_owned()); + }) + .with_model_info_override("gpt-5.5", |model| { + model.comp_hash = Some("budget-checkpoint".to_owned()); + }) + .with_pre_build_hook(|home| { + std::fs::write(home.join("config.toml"), + "[features.guardianv2]\nenabled = true\nmax_parent_compaction_tokens = 4096\n[features.guardianv2.review_scope]\ncomputer_use_only = false\n[features.guardianv2.transcript]\ninclude_images = true\n").unwrap(); + }) + .with_config(move |config| { + // A smaller parent window must not replace Luna's catalog allowance. + config.model_context_window = Some(1_000); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config.guardian_policy_config = Some(TEST_GUARDIAN_POLICY.to_owned()); + config.model_provider.base_url = Some(base_url); + }) + .build_with_auto_env(&server).await?; + let mut builder = ExtensionRegistryBuilder::new(); + super::super::install( + &mut builder, + test.thread_manager.auth_manager(), + Arc::downgrade(&test.thread_manager), + ); + let registry = builder.build(); + let session_store = ExtensionData::new("session-1"); + registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &test.config, + session_source: &SessionSource::Exec, + persistent_thread_state_available: false, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &session_store, + thread_store: test.codex.thread_extension_data(), + }) + .await; + Ok(GuardianFailureFixture { + test, + registry, + session_store, + }) +} + +#[derive(Clone, Copy)] +enum BudgetEvidence { + Checkpoint, + Image, +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_accounts_for_parent_checkpoint() -> Result<()> { + skip_if_no_network!(Ok(())); + assert_catalog_budget(BudgetEvidence::Checkpoint).await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_preserves_optional_evidence_or_defers_to_sync() -> Result<()> { + skip_if_no_network!(Ok(())); + assert_catalog_budget(BudgetEvidence::Image).await +} + +async fn assert_catalog_budget(evidence: BudgetEvidence) -> Result<()> { + let windows = match evidence { + BudgetEvidence::Checkpoint => [18_000, 32_000], + // In the smaller window the text fits, but the 10K image reservation cannot. + BudgetEvidence::Image => [14_000, 40_000], + }; + for (window, outcome) in windows + .into_iter() + .zip([BudgetOutcome::RequiresSync, BudgetOutcome::Fits]) + { + let server = responses::start_mock_server().await; + let fixture = catalog_budget_fixture(server.uri(), window).await?; + let response = responses::mount_sse_once( + &server, + responses::sse(vec![ + ev_assistant_message("budget-score", "low"), + ev_completed("budget-score"), + ]), + ) + .await; + let instruction = "Inspect README.md; do not change files."; + let mut user = user_instruction(instruction); + let (checkpoint, commentary) = match evidence { + BudgetEvidence::Checkpoint => ( + Some(ResponseItem::ContextCompaction { + id: Some(ResponseItemId::from_server("cmp_budget".to_owned())), + encrypted_content: Some("checkpoint".repeat(/*n*/ 1_200)), + internal_chat_message_metadata_passthrough: None, + }), + (0..4) + .map(|index| { + format!( + "commentary {index}: {}", + "optional checkpoint commentary ".repeat(/*n*/ 180) + ) + }) + .collect::>(), + ), + BudgetEvidence::Image => { + let ResponseItem::Message { content, .. } = &mut user else { + unreachable!("user instruction is a message"); + }; + content.push(ContentItem::InputImage { + image_url: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGPgEpEDAABoAD1UCKP3AAAAAElFTkSuQmCC".to_owned(), + detail: None, + }); + (None, vec!["optional old commentary ".repeat(/*n*/ 500)]) + } + }; + let mut history = checkpoint.iter().cloned().collect::>(); + history.push(user); + history.extend(commentary.into_iter().map(|text| ResponseItem::Message { + id: None, + role: "assistant".to_owned(), + content: vec![ContentItem::InputText { text }], + phase: Some(MessagePhase::Commentary), + internal_chat_message_metadata_passthrough: None, + })); + fixture + .test + .codex + .inject_response_items(history.clone()) + .await?; + let history: Arc = match evidence { + BudgetEvidence::Checkpoint => Arc::new(TestRetainedHistory { + retained: history.clone(), + current: TestConversationHistory(history), + compaction_model_hash: Some("budget-checkpoint".to_owned()), + }), + BudgetEvidence::Image => fixture.test.codex.conversation_history_snapshot().await, + }; + let thread_store = fixture.test.codex.thread_extension_data(); + let turn_store = ExtensionData::new("turn-1"); + let tool_name = ToolName::plain("read_file"); + let payload = ToolPayload::Function { + arguments: r#"{"path":"README.md"}"#.to_owned(), + }; + fixture.registry.tool_lifecycle_contributors()[0] + .on_tool_start(ToolStartInput { + session_store: &fixture.session_store, + thread_store, + turn_store: &turn_store, + turn_id: "turn-1", + root_turn_id: None, + call_id: "budget-call", + originating_item_id: None, + tool_name: &tool_name, + mcp_tool: None, + payload: &payload, + conversation_history: history, + source: ToolCallSource::Direct, + }) + .await; + if matches!(outcome, BudgetOutcome::RequiresSync) { + fixture.assert_fails_closed("elevated_risk").await?; + assert!( + server + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.method.as_str() != "POST") + ); + fixture.test.codex.shutdown_and_wait().await?; + continue; + } + let progress = thread_store.get::().unwrap(); + tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { + while progress.latest_scored_tool_call.load(Ordering::Acquire) + < progress.latest_tool_call.load(Ordering::Acquire) + { + tokio::task::yield_now().await; + } + }) + .await?; + let request = response.single_request().body_json(); + let input = request["input"].to_string(); + if let Some(checkpoint) = checkpoint { + assert_eq!(request["input"][2], serde_json::to_value(checkpoint)?); + assert!(input.contains("commentary 0:")); + } else { + assert!(input.contains("input_image")); + assert!(input.contains("optional old commentary")); + } + assert!(input.contains(instruction)); + assert!(input.contains("read_file")); + assert_eq!( + cached_approval( + &fixture.registry, + thread_store, + "review action", + /*metrics*/ None + ) + .await, + Some(ReviewDecision::Approved) + ); + fixture.test.codex.shutdown_and_wait().await?; + } + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index 6860cf52cb..2d117f661a 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -27,7 +27,6 @@ use codex_extension_api::GuardianV2Enabled; use codex_extension_api::SkillInvocationContributor; use codex_extension_api::SkillInvocationInput; use codex_extension_api::ThreadLifecycleContributor; -use codex_extension_api::ThreadOriginator; use codex_extension_api::ThreadStartInput; use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; @@ -37,7 +36,6 @@ use codex_guardian_context::ContextTarget; use codex_guardian_context::PlannedAction; use codex_guardian_context::PlannedActionKind; use codex_history::RolloutItem; -use codex_login::AgentIdentityAuthPolicy; use codex_login::AuthManager; use codex_model_provider::create_model_provider; use codex_protocol::ThreadId; @@ -59,10 +57,8 @@ use super::metrics::sampler_failure_reason; use super::parent_compaction::ParentCompactionError; use super::parent_compaction::select_parent_compaction; use super::sampler::LunaSampler; -use super::sampler::LunaSamplerConfig; use super::sampler::LunaSamplerError; use super::sampler::LunaSamplingRequest; -use super::sampler::MODEL; use super::truncation::ClassificationTruncations; use super::trusted_skills::TrustedSkillInvocations; use super::trusted_skills::TrustedSkillRoots; @@ -132,38 +128,12 @@ impl ThreadLifecycleContributor for GuardianV2Extension { policy.enforce_required_model(); } let scoring_enabled = policy.scoring_enabled(); - let luna_compaction_hash = if let Some(thread_manager) = self.thread_manager.upgrade() { - thread_manager - .get_models_manager() - .get_model_info(MODEL, &input.config.to_models_manager_config()) - .await - .comp_hash - } else { - None - }; - let sampler_config = LunaSamplerConfig { - provider: create_model_provider( - input.config.model_provider.clone(), - Some(Arc::clone(&self.auth_manager)), - ), - http_client_factory: input.config.http_client_factory(), - agent_identity_policy: if input.config.features.enabled(Feature::UseAgentIdentity) { - AgentIdentityAuthPolicy::ChatGptAuth - } else { - AgentIdentityAuthPolicy::JwtOnly - }, - session_source: input.session_source.clone(), - session_id: input.session_store.level_id().to_string(), - thread_id: thread_id.clone(), - originator: input - .thread_store - .get::() - .map(|originator| originator.0.clone()), - free_guardian: input.config.free_guardian_enabled(), - service_tier: input.config.service_tier.clone(), - luna_compaction_hash, - metrics: input.extension_metrics.clone(), - }; + let sampler_config = super::startup::sampler_config( + &input, + Arc::clone(&self.auth_manager), + self.thread_manager.upgrade(), + ) + .await; if scoring_enabled && guardian_config.transcript.include_images { input 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 85c130050c..e6fd490624 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 @@ -3448,3 +3448,6 @@ fn review_scope(action: &serde_json::Value) -> Option { tool => GuardianScope::for_tool(&ToolName::plain(tool)), } } + +#[path = "budget_tests.rs"] +mod budget; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs index 9b60a471eb..138dbb08ee 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs @@ -25,6 +25,7 @@ pub(super) fn sampler_failure_reason(error: &LunaSamplerError) -> &'static str { LunaSamplerError::OutputTooLarge => "output_too_large", LunaSamplerError::Superseded => "superseded", LunaSamplerError::IncompatibleCompaction => "incompatible_compaction", + LunaSamplerError::InputTooLarge => "input_too_large", LunaSamplerError::Api(error) => match error { ApiError::Transport(TransportError::Http { status, .. }) | ApiError::Api { status, .. } => match status.as_u16() { diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs index 55dc263604..821d49485d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/mod.rs @@ -7,6 +7,7 @@ mod extension; mod metrics; mod parent_compaction; mod sampler; +mod startup; mod transcript; mod truncation; mod trusted_skills; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index 2f68b94e49..15e8de428c 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -69,6 +69,8 @@ pub struct LunaSamplerConfig { pub service_tier: Option, /// Luna model's host-resolved encrypted-compaction compatibility hash. pub luna_compaction_hash: Option, + /// Complete input allowance resolved for the classifier model. + pub max_input_tokens: usize, /// Host-provided metrics capability with the owning session's attribution. pub metrics: Option>, } @@ -117,6 +119,9 @@ pub enum LunaSamplerError { /// The supplied parent checkpoint cannot be consumed by this Luna configuration. #[error("parent compaction is incompatible with Luna")] IncompatibleCompaction, + /// The complete classifier input exceeded the model allowance. + #[error("Luna input exceeds the complete request budget")] + InputTooLarge, } struct ActiveRequest { @@ -224,6 +229,7 @@ impl LunaSampler { | LunaSamplerError::OutputTooLarge | LunaSamplerError::Superseded | LunaSamplerError::IncompatibleCompaction + | LunaSamplerError::InputTooLarge | LunaSamplerError::Api( ApiError::Transport(TransportError::Build(_)) | ApiError::ContextWindowExceeded @@ -310,6 +316,10 @@ impl LunaSampler { ); } } + // Oversized classifications defer to sync with the existing failure score. + if total_tokens > self.config.max_input_tokens.saturating_sub(/*rhs*/ 256) { + return Err(LunaSamplerError::InputTooLarge); + } let mut request = ResponsesApiRequest { model: MODEL.to_owned(), instructions: String::new(), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index 6dc562c92c..23635baf6e 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs @@ -208,6 +208,7 @@ pub(super) fn sampler_config(base_url: String) -> LunaSamplerConfig { free_guardian: false, service_tier: None, luna_compaction_hash: None, + max_input_tokens: codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS, metrics: None, } } @@ -490,6 +491,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications free_guardian: false, service_tier: None, luna_compaction_hash: None, + max_input_tokens: codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS, metrics: None, }) .await?; @@ -733,6 +735,7 @@ async fn sampler_returns_classification_token_before_terminal_response_events() free_guardian: false, service_tier: None, luna_compaction_hash: None, + max_input_tokens: codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS, metrics: None, }) .await?; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs new file mode 100644 index 0000000000..34a87d643a --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/startup.rs @@ -0,0 +1,65 @@ +//! Resolves the classifier's own model budget and transport configuration. +//! Parent-model context overrides do not change the classifier allowance. + +use std::sync::Arc; + +use codex_core::ThreadManager; +use codex_core::config::Config; +use codex_extension_api::ThreadOriginator; +use codex_extension_api::ThreadStartInput; +use codex_features::Feature; +use codex_login::AgentIdentityAuthPolicy; +use codex_login::AuthManager; +use codex_model_provider::create_model_provider; + +use super::sampler::LunaSamplerConfig; +use super::sampler::MODEL; + +pub(super) async fn sampler_config( + input: &ThreadStartInput<'_, Config>, + auth_manager: Arc, + manager: Option>, +) -> LunaSamplerConfig { + let luna_model = if let Some(thread_manager) = manager { + let mut model_config = input.config.to_models_manager_config(); + model_config.model_context_window = None; + Some( + thread_manager + .get_models_manager() + .get_model_info(MODEL, &model_config) + .await, + ) + } else { + None + }; + let max_input_tokens = + luna_model + .as_ref() + .map_or(codex_guardian_context::DEFAULT_MAX_INPUT_TOKENS, |model| { + codex_guardian_context::effective_input_token_limit( + model, /*configured_window*/ None, + ) + }); + let luna_compaction_hash = luna_model.and_then(|model| model.comp_hash); + LunaSamplerConfig { + provider: create_model_provider(input.config.model_provider.clone(), Some(auth_manager)), + http_client_factory: input.config.http_client_factory(), + agent_identity_policy: if input.config.features.enabled(Feature::UseAgentIdentity) { + AgentIdentityAuthPolicy::ChatGptAuth + } else { + AgentIdentityAuthPolicy::JwtOnly + }, + session_source: input.session_source.clone(), + session_id: input.session_store.level_id().to_string(), + thread_id: input.thread_store.level_id().to_string(), + originator: input + .thread_store + .get::() + .map(|originator| originator.0.clone()), + free_guardian: input.config.free_guardian_enabled(), + service_tier: input.config.service_tier.clone(), + luna_compaction_hash, + max_input_tokens, + metrics: input.extension_metrics.clone(), + } +}