diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 0a9abd4611..00fc096d22 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -757,7 +757,7 @@ impl ModelClient { } } - fn build_responses_request( + pub(crate) fn build_responses_request( &self, prompt: &Prompt, model_info: &ModelInfo, diff --git a/codex-rs/core/src/guardian/input_budget.rs b/codex-rs/core/src/guardian/input_budget.rs new file mode 100644 index 0000000000..e1f7b42445 --- /dev/null +++ b/codex-rs/core/src/guardian/input_budget.rs @@ -0,0 +1,188 @@ +//! Finalizes a pending reviewer input after tools and turn context are resolved. +//! The thread attachment is a single in-flight request, consumed before history +//! recording. Budget failures preserve it for a compaction retry; successful +//! finalization consumes it. It is not another retained history. + +use codex_features::Feature; +use codex_guardian_context::ComposedContext; +use codex_guardian_context::RequestBudget; +use codex_guardian_context::effective_input_token_limit; +use codex_protocol::config_types::ReasoningSummary; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::models::ResponseItem; + +use crate::context::ContextualUserFragment; +use crate::context::GuardianBudgetOmission; +use crate::context_manager::estimate_item_token_count; +use crate::responses_metadata::CodexResponsesRequestKind; +use crate::session::TurnInput; +use crate::session::session::Session; +use crate::session::step_context::StepContext; +use crate::session::turn::build_prompt; +use crate::session::turn_context::TurnContext; + +#[derive(Clone)] +pub(crate) struct PendingReviewContext(pub ComposedContext); + +/// Reject inputs that cannot fit even without history or tools before Core tries +/// pre-turn compaction. This is only a feasibility check; its selected copy is +/// discarded. Actual selection waits for the complete first-step overhead. +pub(crate) async fn check_pending(session: &Session, turn: &TurnContext) -> CodexResult<()> { + let Some(pending) = session + .services + .thread_extension_data + .get::() + else { + return Ok(()); + }; + let base = session.get_prompt_base_instructions().await; + let minimum_prefix = + codex_protocol::protocol::TruncationPolicy::Bytes(base.text.len()).token_budget(); + let maximum = effective_input_token_limit(turn.model_info(), turn.config.model_context_window) + .saturating_sub(super::request_budget::INPUT_TOKEN_MARGIN); + if pending.0.estimated_tokens().saturating_add(minimum_prefix) > maximum { + pending + .0 + .clone() + .enforce_budget( + RequestBudget { + max_input_tokens: maximum, + existing_context_tokens: minimum_prefix, + }, + GuardianBudgetOmission.render(), + ) + .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?; + } + Ok(()) +} + +pub(crate) async fn finalize( + session: &Session, + step: &StepContext, + input: &mut [TurnInput], +) -> CodexResult<()> { + let Some(pending) = session + .services + .thread_extension_data + .get::() + else { + return Ok(()); + }; + let [TurnInput::UserInput { content, .. }] = input else { + return Err(CodexErr::InvalidRequest( + "Guardian expects one review input".to_owned(), + )); + }; + let context = pending.0.clone(); + let model = &step.settings.model_info; + let history = session.clone_history().await; + let prompt = build_prompt( + history.for_prompt(&model.input_modalities), + step, + session.get_prompt_base_instructions().await, + ); + let request = session.services.model_client.build_responses_request( + &prompt, + model, + /*effort*/ None, + ReasoningSummary::None, + /*service_tier*/ None, + &session + .responses_metadata(&step.turn, CodexResponsesRequestKind::Turn) + .await, + )?; + let mut existing = super::request_budget::estimate_request_tokens(&request) + .max(usize::try_from(session.get_total_token_usage().await).unwrap_or(usize::MAX)); + // These reminders are appended after the accepted input. Reserve their full + // serialized cost before choosing optional evidence, even if the time + // reminder's interval ultimately suppresses it. + if step + .turn + .config + .features + .enabled(Feature::CurrentTimeReminder) + && step.turn.config.current_time_reminder.is_some() + { + let reminder = ContextualUserFragment::into(crate::context::CurrentTimeReminder::new( + chrono::Utc::now(), + )); + existing = existing.saturating_add( + usize::try_from(estimate_item_token_count(&reminder)).unwrap_or(usize::MAX), + ); + } + if let Some(reminder) = session + .services + .agent_control + .rollout_budget() + .pending_reminder(session.thread_id(), &session.current_window_id().await) + { + let reminder = ContextualUserFragment::into(crate::context::RolloutBudgetContext { + remaining_tokens: reminder.remaining_tokens, + }); + existing = existing.saturating_add( + usize::try_from(estimate_item_token_count(&reminder)).unwrap_or(usize::MAX), + ); + } + // Reserve the real input annotations as well as the shared content cost. + // Removing optional items can only reduce this metadata. The remaining + // margin covers IDs and timestamps assigned when recording the message. + let mut framing = session.response_item_from_user_input(content.clone()); + if let ResponseItem::Message { content, .. } = &mut framing { + content.clear(); + } + let budget = RequestBudget { + max_input_tokens: effective_input_token_limit(model, step.turn.config.model_context_window) + .saturating_sub(super::request_budget::INPUT_TOKEN_MARGIN), + existing_context_tokens: existing.saturating_add( + usize::try_from(estimate_item_token_count(&framing)).unwrap_or(usize::MAX), + ), + }; + let context = context + .enforce_budget(budget, GuardianBudgetOmission.render()) + .map_err(|error| { + session + .services + .thread_extension_data + .insert(super::request_budget::ExhaustedReviewBudget); + match error { + codex_guardian_context::SectionError::EvidenceLimitExceeded { .. } => { + CodexErr::ContextWindowExceeded + } + error => CodexErr::InvalidRequest(error.to_string()), + } + })?; + for (section, cost) in context.section_costs() { + for (measurement, value) in cost.measurements() { + session + .services + .session_telemetry + .histogram_with_boundaries( + codex_guardian_context::SECTION_COST_METRIC, + i64::try_from(value).unwrap_or(i64::MAX), + codex_guardian_context::SECTION_COST_BOUNDARIES, + &[ + ("target", "sync"), + ("section", section), + ("measurement", measurement), + ], + ); + } + } + *content = context + .into_user_inputs() + .map_err(|error| CodexErr::InvalidRequest(error.to_string()))?; + session + .services + .thread_extension_data + .remove::(); + session + .services + .thread_extension_data + .remove::(); + Ok(()) +} + +#[cfg(test)] +#[path = "input_budget_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/guardian/input_budget_tests.rs b/codex-rs/core/src/guardian/input_budget_tests.rs new file mode 100644 index 0000000000..6681d12315 --- /dev/null +++ b/codex-rs/core/src/guardian/input_budget_tests.rs @@ -0,0 +1,138 @@ +//! Exercises reviewer history safety at cancellation and budget boundaries. + +use super::*; +use codex_guardian_context::ContextPresentation; +use codex_guardian_context::ContextProfile; +use codex_guardian_context::PlannedAction; +use codex_guardian_context::PlannedActionKind; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::TurnAbortReason; +use std::sync::Arc; + +fn required_context(text: String) -> ComposedContext { + let action = PlannedAction { + json: text, + kind: PlannedActionKind::Command, + reason: None, + }; + let context = super::super::prompt::collect_guardian_context( + &Vec::::new(), + super::super::GUARDIAN_MAX_TOOL_ENTRY_TOKENS, + &[], + &[], + Some(&action), + /*permissions*/ None, + /*node_repl*/ None, + ) + .expect("collect required context"); + let transcript = ContextProfile::synchronous() + .render_transcript(context.transcript_entries(), /*entry_number_offset*/ 0); + context + .compose( + ContextPresentation::SyncFull { + session_id: "test-parent", + }, + transcript, + ) + .expect("compose required context") +} + +#[tokio::test] +async fn cancelled_startup_does_not_record_unselected_review_evidence() { + let (session, turn, events) = crate::session::tests::make_session_and_context_with_rx().await; + let context = required_context("unselected review evidence ".repeat(/*n*/ 10_000)); + let content = context.clone().into_user_inputs().unwrap(); + session + .services + .thread_extension_data + .insert(PendingReviewContext(context)); + session + .set_session_startup_prewarm( + crate::session_startup_prewarm::SessionStartupPrewarmHandle::new( + tokio::spawn(std::future::pending()), + std::time::Instant::now(), + crate::client::WEBSOCKET_CONNECT_TIMEOUT, + ), + ) + .await; + session + .spawn_task( + turn, + vec![TurnInput::UserInput { + acceptance_order: None, + content, + client_id: None, + }], + crate::tasks::RegularTask::new(), + ) + .await; + let started = tokio::time::timeout(std::time::Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert!(matches!(started.msg, EventMsg::TurnStarted(_))); + session.abort_all_tasks(TurnAbortReason::Interrupted).await; + + let history = session.clone_history().await; + let recorded = history.raw_items().collect::>(); + assert!( + !serde_json::to_string(&recorded) + .unwrap() + .contains("unselected review evidence") + ); + assert!( + session + .services + .thread_extension_data + .get::() + .is_some() + ); +} + +#[tokio::test] +async fn feasibility_rejects_required_evidence_inside_the_reserved_margin() { + let (session, mut turn) = crate::session::tests::make_session_and_context().await; + let context = required_context("required action".to_owned()); + let base = session.get_prompt_base_instructions().await; + let prefix = codex_protocol::protocol::TruncationPolicy::Bytes(base.text.len()).token_budget(); + let model = Arc::make_mut(&mut Arc::make_mut(&mut turn.initial_settings).model_info); + model.effective_context_window_percent = 100; + Arc::make_mut(&mut turn.config).model_context_window = + Some(i64::try_from(context.estimated_tokens() + prefix + 128).unwrap()); + session + .services + .thread_extension_data + .insert(PendingReviewContext(context)); + + assert!(check_pending(&session, &turn).await.is_err()); +} + +#[tokio::test] +async fn finalization_overflow_marks_the_reviewer_exhausted() { + let (session, mut turn) = crate::session::tests::make_session_and_context().await; + Arc::make_mut(&mut turn.config).model_context_window = Some(1); + let session = Arc::new(session); + let step = session + .capture_step_context(Arc::new(turn), &tokio_util::sync::CancellationToken::new()) + .await + .unwrap(); + let context = required_context("required action".to_owned()); + let mut input = vec![TurnInput::UserInput { + acceptance_order: None, + content: context.clone().into_user_inputs().unwrap(), + client_id: None, + }]; + session + .services + .thread_extension_data + .insert(PendingReviewContext(context)); + + assert!(finalize(&session, &step, &mut input).await.is_err()); + assert!( + session + .services + .thread_extension_data + .get::() + .is_some() + ); +} diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 32a3f2f5e3..c8796a9a87 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -6,9 +6,15 @@ mod approval_request; mod coverage; mod decision; mod feedback; +mod input_budget; mod metrics; mod prompt; +pub(crate) use input_budget::PendingReviewContext; +pub(crate) use input_budget::check_pending as check_pending_guardian_input; +pub(crate) use input_budget::finalize as finalize_guardian_input; mod request_budget; +pub(crate) use request_budget::ExhaustedReviewBudget; +pub(crate) use request_budget::check_prompt as check_guardian_prompt_budget; pub(crate) use request_budget::observe as observe_guardian_request; mod review; mod review_session; diff --git a/codex-rs/core/src/guardian/prompt.rs b/codex-rs/core/src/guardian/prompt.rs index ca89fb4b3a..559bca40b5 100644 --- a/codex-rs/core/src/guardian/prompt.rs +++ b/codex-rs/core/src/guardian/prompt.rs @@ -1,6 +1,7 @@ use codex_extension_api::ConversationHistorySnapshot; use codex_guardian_context::Budgeted; use codex_guardian_context::CollectedContext; +use codex_guardian_context::ComposedContext; use codex_guardian_context::ContextPresentation; use codex_guardian_context::ContextProfile; #[cfg(test)] @@ -14,7 +15,6 @@ use codex_guardian_context::SectionHistory; use codex_guardian_context::SectionInput; use codex_guardian_context::default_registry; use codex_protocol::models::ResponseItem; -use codex_protocol::user_input::UserInput; use crate::context::GuardianReviewEvidence; use crate::context::NodeReplReviewEvidence; @@ -37,7 +37,7 @@ const GUARDIAN_MAX_APPROVAL_REASON_TOKENS: usize = 512; pub(super) const GUARDIAN_TRANSCRIPT_START: &str = ">>> TRANSCRIPT START\n"; pub(crate) struct GuardianPromptItems { - pub(crate) items: Vec, + pub(crate) context: ComposedContext, pub(crate) transcript_cursor: GuardianTranscriptCursor, pub(crate) node_repl_evidence_sequence: u64, pub(crate) reviewed_action_truncated: bool, @@ -214,26 +214,8 @@ pub(crate) async fn build_guardian_prompt_items_with_parent_turn( .push(Budgeted::required(placeholder.to_owned())); } let context = sections.compose(presentation, transcript)?; - for (section, cost) in context.section_costs() { - for (measurement, value) in cost.measurements() { - session - .services - .session_telemetry - .histogram_with_boundaries( - codex_guardian_context::SECTION_COST_METRIC, - i64::try_from(value).unwrap_or(i64::MAX), - codex_guardian_context::SECTION_COST_BOUNDARIES, - &[ - ("target", "sync"), - ("section", section), - ("measurement", measurement), - ], - ); - } - } - let items = context.into_user_inputs()?; Ok(GuardianPromptItems { - items, + context, transcript_cursor, node_repl_evidence_sequence, reviewed_action_truncated: planned_action_json.truncated, diff --git a/codex-rs/core/src/guardian/request_budget.rs b/codex-rs/core/src/guardian/request_budget.rs index ad53549845..24cc7948d8 100644 --- a/codex-rs/core/src/guardian/request_budget.rs +++ b/codex-rs/core/src/guardian/request_budget.rs @@ -1,28 +1,28 @@ -//! Measures the complete synchronous request after image admission and wire -//! prefix assembly. Includes reused history, tool definitions and output format. +//! Measures and checks complete synchronous requests after wire prefix assembly. +//! Includes reused history, tool definitions, output format and continuations. +use codex_api::ResponsesApiRequest; use codex_guardian_context::REQUEST_TOKENS_BOUNDARIES; use codex_guardian_context::REQUEST_TOKENS_METRIC; +use codex_guardian_context::effective_input_token_limit; +use codex_otel::SessionTelemetry; +use codex_protocol::error::CodexErr; +use codex_protocol::error::Result as CodexResult; +use codex_protocol::openai_models::ModelInfo; use codex_protocol::protocol::TruncationPolicy; +use crate::config::Config; use crate::context_manager::estimate_item_token_count; -use codex_api::ResponsesApiRequest; -use codex_otel::SessionTelemetry; +use crate::session::session::Session; + +pub(super) const INPUT_TOKEN_MARGIN: usize = 256; + +// Keep a failed reviewer out of reuse unless compaction and a fresh budget +// check succeed. Also distinguishes local budget failures from backend errors. +pub(crate) struct ExhaustedReviewBudget; pub(crate) fn observe(telemetry: &SessionTelemetry, request: &ResponsesApiRequest) -> usize { - let input = request - .input - .iter() - .map(estimate_item_token_count) - .fold(0i64, i64::saturating_add); - let instructions = TruncationPolicy::Bytes(request.instructions.len()).token_budget(); - let metadata = serde_json::to_vec(&(&request.tools, &request.text)) - .map(|bytes| TruncationPolicy::Bytes(bytes.len()).token_budget()) - .unwrap_or(usize::MAX); - let total = usize::try_from(input) - .unwrap_or(usize::MAX) - .saturating_add(instructions) - .saturating_add(metadata); + let total = estimate_request_tokens(request); // The assembled input already includes inherited history and the current // review. Do not report a guessed old/new split after context injection. telemetry.histogram_with_boundaries( @@ -33,3 +33,53 @@ pub(crate) fn observe(telemetry: &SessionTelemetry, request: &ResponsesApiReques ); total } + +pub(super) fn estimate_request_tokens(request: &ResponsesApiRequest) -> usize { + let input = request + .input + .iter() + .map(estimate_item_token_count) + .fold(0i64, i64::saturating_add); + let instructions = TruncationPolicy::Bytes(request.instructions.len()).token_budget(); + let metadata = serde_json::to_vec(&(&request.tools, &request.text)) + .map(|bytes| TruncationPolicy::Bytes(bytes.len()).token_budget()) + .unwrap_or(usize::MAX); + usize::try_from(input) + .unwrap_or(usize::MAX) + .saturating_add(instructions) + .saturating_add(metadata) +} + +/// Checks the fully assembled prompt after turn context has been injected. This +/// final guard also covers retries and reviewer tool continuations. +pub(crate) fn check_prompt( + session: &Session, + prompt: &crate::client_common::Prompt, + config: &Config, + model: &ModelInfo, + metadata: &crate::responses_metadata::CodexResponsesMetadata, +) -> CodexResult<()> { + let request = session.services.model_client.build_responses_request( + prompt, + model, + /*effort*/ None, + codex_protocol::config_types::ReasoningSummary::None, + /*service_tier*/ None, + metadata, + )?; + if estimate_request_tokens(&request) + > effective_input_token_limit(model, config.model_context_window) + .saturating_sub(INPUT_TOKEN_MARGIN) + { + session + .services + .thread_extension_data + .insert(ExhaustedReviewBudget); + return Err(CodexErr::ContextWindowExceeded); + } + session + .services + .thread_extension_data + .remove::(); + Ok(()) +} diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index c753599571..f7a4bfb8d7 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -512,9 +512,9 @@ async fn run_review_on_session( .await?; if prompt_items - .items - .iter() - .any(|item| matches!(item, UserInput::Image { .. })) + .context + .section_costs() + .any(|(_, cost)| cost.image_count > 0) { let reviewer_history = review_session.session.clone_history().await; let reviewer_image_urls = reviewer_history @@ -548,10 +548,7 @@ async fn run_review_on_session( } else { ImagePreparationMode::DetailBased }; - prompt_items.items.retain_mut(|item| { - let UserInput::Image { detail, .. } = item else { - return true; - }; + prompt_items.context.retain_images(|image_url, detail| { *detail = match normalize_output_image_detail(&model_info, *detail) { _ if mode == ImagePreparationMode::UnifiedBudget => { Some(ImageDetail::Original) @@ -559,7 +556,13 @@ async fn run_review_on_session( Some(ImageDetail::Low) => Some(ImageDetail::High), detail => detail, }; - let mut prepared = vec![ResponseInputItem::from(vec![item.clone()]).into()]; + let mut prepared = vec![ + ResponseInputItem::from(vec![UserInput::Image { + image_url: image_url.to_owned(), + detail: *detail, + }]) + .into(), + ]; prepare_response_items( &mut prepared, mode, @@ -574,7 +577,8 @@ async fn run_review_on_session( }) }); let prompt: ResponseItem = - ResponseInputItem::from(prompt_items.items.clone()).into(); + ResponseInputItem::from(prompt_items.context.clone().into_user_inputs()?) + .into(); let prompt_tokens = crate::context_manager::estimate_item_token_count(&prompt); let base_instructions = review_session.session.get_base_instructions().await; let history_tokens = reviewer_history @@ -587,13 +591,12 @@ async fn run_review_on_session( false }; if !admit_images { - prompt_items - .items - .retain(|item| !matches!(item, UserInput::Image { .. })); + prompt_items.context.retain_images(|_, _| false); } } - Ok::<_, anyhow::Error>(prompt_items) + let items = prompt_items.context.clone().into_user_inputs()?; + Ok::<_, anyhow::Error>((prompt_items, items)) }), ) .await; @@ -601,7 +604,7 @@ async fn run_review_on_session( Ok(prompt_items) => prompt_items, Err(outcome) => return (outcome, false, analytics_result), }; - let prompt_items = match prompt_items { + let (prompt_items, items) = match prompt_items { Ok(prompt_items) => prompt_items, Err(err) => { return ( @@ -652,8 +655,15 @@ async fn run_review_on_session( .unwrap_or_else(|| params.parent_context.turn().config.cwd.clone()); let parent_turn = params.parent_context.turn(); + review_session + .session + .services + .thread_extension_data + .insert(super::input_budget::PendingReviewContext( + prompt_items.context, + )); let submission = review_session.io.submit_turn_input( - TurnInputRequest::user_input(prompt_items.items) + TurnInputRequest::user_input(items) .with_thread_settings(codex_protocol::protocol::ThreadSettingsOverrides { environments: Some(codex_protocol::protocol::TurnEnvironmentSelections::new( parent_turn_legacy_fallback_cwd, @@ -696,6 +706,13 @@ async fn run_review_on_session( Box::pin(submission), ) .await; + if !matches!(&submit_result, Ok(Ok(TurnInputSubmission::Started { .. }))) { + review_session + .session + .services + .thread_extension_data + .remove::(); + } let child_turn_id = match submit_result { Ok(Ok(TurnInputSubmission::Started { turn_id })) => turn_id, Ok(Ok(submission)) => { @@ -737,6 +754,11 @@ async fn run_review_on_session( &mut analytics_result, ) .await; + review_session + .session + .services + .thread_extension_data + .remove::(); if matches!(outcome.0, GuardianReviewSessionOutcome::Completed(_)) { if outcome.2 && let Some(total_token_usage) = review_session.session.total_token_usage().await @@ -750,7 +772,14 @@ async fn run_review_on_session( state.prior_review_count = state.prior_review_count.saturating_add(1); state.last_reviewed_transcript_cursor = Some(transcript_cursor); } - (outcome.0, outcome.1, analytics_result) + let keep_review_session = outcome.1 + && review_session + .session + .services + .thread_extension_data + .remove::() + .is_none(); + (outcome.0, keep_review_session, analytics_result) } async fn ensure_guardian_followup_reminder(review_session: &GuardianReviewSession) { diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 4489eecbae..1a88cc4659 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -423,7 +423,7 @@ async fn build_guardian_prompt_full_mode_preserves_initial_review_format() -> an ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("whose request action you are assessing")); assert!(text.contains(">>> TRANSCRIPT START\n")); assert!(text.contains(">>> TRANSCRIPT END\n")); @@ -464,7 +464,7 @@ async fn build_guardian_prompt_prefers_retry_reason_over_approval_reason() -> an ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("Retry reason:\nThe sandbox blocked the initial command.\n\n")); assert!(!text.contains("A policy rule requires approval.")); @@ -506,8 +506,8 @@ async fn build_guardian_prompt_truncates_oversized_approval_reason() -> anyhow:: ) .await?; - let reason_item = prompt - .items + let items = prompt.context.into_user_inputs()?; + let reason_item = items .iter() .find_map(|item| match item { codex_protocol::user_input::UserInput::Text { text, .. } @@ -600,7 +600,7 @@ async fn build_guardian_prompt_includes_parent_turn_denied_reads() -> anyhow::Re ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("PARENT TURN PERMISSION CONTEXT START")); assert!(text.contains("do not approve escalation whose purpose is to read them")); assert!(text.contains(denied_root.to_string_lossy().as_ref())); @@ -664,7 +664,7 @@ async fn build_guardian_prompt_delta_mode_preserves_original_numbering() -> anyh ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("added since your last approval assessment")); assert!(text.contains(">>> TRANSCRIPT DELTA START\n")); assert!(text.contains("[5] user: Please also push the second docs fix.")); @@ -705,7 +705,7 @@ async fn build_guardian_prompt_delta_mode_handles_empty_delta() -> anyhow::Resul ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains(">>> TRANSCRIPT DELTA START\n")); assert!(text.contains("")); assert!(text.contains(">>> TRANSCRIPT DELTA END\n")); @@ -743,7 +743,7 @@ async fn build_guardian_prompt_stale_delta_cursor_falls_back_to_full_prompt() -> ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("whose request action you are assessing")); assert!(text.contains(">>> TRANSCRIPT START\n")); assert!(!text.contains("TRANSCRIPT DELTA")); @@ -832,7 +832,7 @@ async fn build_guardian_prompt_stale_delta_version_falls_back_to_full_prompt() - ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("whose request action you are assessing")); assert!(text.contains(">>> TRANSCRIPT START\n")); assert!(!text.contains("TRANSCRIPT DELTA")); @@ -1239,7 +1239,7 @@ async fn build_guardian_prompt_items_keeps_required_node_repl_reviews_generic() ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("Assess the exact planned action below.")); assert!(text.contains("Retry reason:\nRetry the authorized browser inspection.")); assert!(text.contains("Planned action JSON:")); @@ -1275,7 +1275,7 @@ async fn build_guardian_prompt_items_keeps_other_requests_generic() -> anyhow::R ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("Assess the exact planned action below.")); assert!(text.contains("Planned action JSON:")); assert!(!text.contains("Node REPL action JSON:")); @@ -1360,7 +1360,7 @@ async fn build_guardian_prompt_items_explains_network_access_review_scope() -> a ) .await?; - let text = guardian_prompt_text(&prompt.items); + let text = guardian_prompt_text(&prompt.context.into_user_inputs()?); assert!(text.contains("Below is a proposed network access request under review.")); assert!(!text.contains("Network approval context:")); assert!( @@ -2330,7 +2330,8 @@ async fn build_guardian_prompt_items_includes_parent_session_id() -> anyhow::Res ) .await?; let prompt_text = prompt - .items + .context + .into_user_inputs()? .into_iter() .map(|item| match item { codex_protocol::user_input::UserInput::Text { text, .. } => text, diff --git a/codex-rs/core/src/session/turn.rs b/codex-rs/core/src/session/turn.rs index c68e4a72a4..3950d34d13 100644 --- a/codex-rs/core/src/session/turn.rs +++ b/codex-rs/core/src/session/turn.rs @@ -163,11 +163,14 @@ pub(crate) struct McpStartupRequirements { pub(crate) async fn run_turn( sess: Arc, turn_context: Arc, - input: Vec, + mut input: Vec, mcp_startup_requirements: &mut McpStartupRequirements, prewarmed_client_session: Option, cancellation_token: CancellationToken, ) -> CodexResult> { + if crate::guardian::is_basic_session_source(&turn_context.session_source) { + crate::guardian::check_pending_guardian_input(&sess, &turn_context).await?; + } // Record results from hooks that finished after the previous turn before this turn's user prompt. drain_async_hook_results(&sess, &turn_context, /*before_user_prompt*/ true).await; @@ -305,6 +308,34 @@ pub(crate) async fn run_turn( if run_pending_session_start_hooks(&sess, &turn_context).await { return Ok(None); } + if crate::guardian::is_basic_session_source(&turn_context.session_source) + && let Err(error) = + crate::guardian::finalize_guardian_input(&sess, &first_step_context, &mut input).await + { + // Token-budget compaction resets history, which can discard the evidence + // referenced by a pending delta review. Leave budget failures unreusable. + if !matches!(error.details(), CodexErrorDetails::ContextWindowExceeded) + || turn_context.config.features.enabled(Feature::TokenBudget) + { + return Err(error); + } + // Incoming evidence can overflow even below the normal history + // threshold. Keep it pending while compacting, then select once more. + run_auto_compact( + &sess, + Arc::clone(&first_step_context), + /*fallback_step_context*/ None, + &mut client_session, + InitialContextInjection::DoNotInject, + CompactionReason::ContextLimit, + CompactionPhase::PreTurn, + ) + .await?; + world_state = sess + .record_context_updates_and_set_reference_context_item(first_step_context.as_ref()) + .await?; + crate::guardian::finalize_guardian_input(&sess, &first_step_context, &mut input).await?; + } let mut can_drain_pending_input = input.is_empty(); if run_hooks_and_record_inputs( &sess, @@ -362,6 +393,7 @@ pub(crate) async fn run_turn( // 2. After auto-compact, when model/tool continuation needs to resume before any steer. let mut next_step_context = Some(first_step_context); + let mut guardian_budget_compacted = false; loop { // Note that pending_input would be something like a message the user // submitted through the UI while the model was running. Though the UI @@ -476,6 +508,7 @@ pub(crate) async fn run_turn( .await; match sampling_request_result { Ok((sampling_request_output, sampling_request_input)) => { + guardian_budget_compacted = false; let SamplingRequestResult { needs_follow_up: model_needs_follow_up, last_agent_message: sampling_request_last_agent_message, @@ -644,6 +677,37 @@ pub(crate) async fn run_turn( } continue; } + Err(err) + if matches!(err.details(), CodexErrorDetails::ContextWindowExceeded) + && !guardian_budget_compacted + && !turn_context.config.features.enabled(Feature::TokenBudget) + && sess + .services + .thread_extension_data + .get::() + .is_some() => + { + // Tool continuations can also cross the complete-request limit. + // Only summarizing compaction preserves the action and evidence; + // token-budget resets must fail closed and retire the reviewer. + // Retry once per model step, so ineffective compaction cannot loop. + guardian_budget_compacted = true; + run_auto_compact( + &sess, + Arc::clone(&step_context), + /*fallback_step_context*/ None, + &mut client_session, + InitialContextInjection::BeforeLastUserMessage { + world_state: Arc::clone(&world_state), + step_context: Arc::clone(&step_context), + }, + CompactionReason::ContextLimit, + CompactionPhase::MidTurn, + ) + .await?; + can_drain_pending_input = false; + continue; + } Err(err) if matches!(err.details(), CodexErrorDetails::TurnAborted) => { return Err(err); } @@ -720,6 +784,16 @@ pub(crate) async fn run_hooks_and_record_inputs( input: &[TurnInput], persist_context: PersistContext, ) -> bool { + // Cancellation can reach this path before Guardian's tools and context are + // resolved. Only finalized evidence may enter reusable reviewer history. + if sess + .services + .thread_extension_data + .get::() + .is_some() + { + return false; + } let mut blocked_input = false; let mut accepted_user_input = false; for input_item in input { @@ -1480,6 +1554,15 @@ async fn run_sampling_request( step_context.as_ref(), base_instructions.clone(), ); + if crate::guardian::is_basic_session_source(&turn_context.session_source) { + crate::guardian::check_guardian_prompt_budget( + &sess, + &prompt, + &turn_context.config, + &step_context.settings.model_info, + responses_metadata, + )?; + } let err = match try_run_sampling_request( tool_runtime.clone(), Arc::clone(&sess), diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index e1b0bc558f..f14c48e315 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -4795,8 +4795,10 @@ async fn code_mode_node_repl_text_evidence_is_visible_only_to_guardian( reviewer.input_modalities = vec![codex_protocol::openai_models::InputModality::Text]; } else { + // Fit the required text and request-only prefix, while leaving + // insufficient room for the additional image reservation. reviewer.context_window = - (reviewer_constraint == Some("small")).then_some(10_000); + (reviewer_constraint == Some("small")).then_some(20_000); reviewer.max_context_window = reviewer.context_window; } } diff --git a/codex-rs/core/tests/suite/guardian_context_budget.rs b/codex-rs/core/tests/suite/guardian_context_budget.rs new file mode 100644 index 0000000000..721096f1b3 --- /dev/null +++ b/codex-rs/core/tests/suite/guardian_context_budget.rs @@ -0,0 +1,327 @@ +//! Review requests preserve required evidence while fitting the aggregate budget. +//! Parent token-budget mode must not replace Guardian's summary compaction. + +use anyhow::Result; +use codex_core::config::Constrained; +use codex_core::config::CurrentTimeReminderConfig; +use codex_core::config::RolloutBudgetConfig; +use codex_features::Feature; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::openai_models::AutoReviewMessages; +use codex_protocol::protocol::AskForApproval; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_custom_tool_call; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::skip_if_no_network; +use core_test_support::skip_if_wine_exec; +use core_test_support::test_codex::test_codex; +use pretty_assertions::assert_eq; +use serde_json::json; +use test_case::test_case; + +#[derive(Clone, Copy)] +enum ReviewerResponse { + Decision, + ToolContinuation, + UncompactableContinuation, + NextReview, +} + +#[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::UncompactableContinuation; "ineffective_compaction_fails_closed")] +#[test_case(4_500, ReviewerResponse::NextReview; "incoming_review_compacts_existing_history")] +async fn review_respects_complete_context_budget( + window: i64, + reviewer_response: ReviewerResponse, +) -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_wine_exec!( + Ok(()), + "Guardian approval actions require host-native paths" + ); + let server = responses::start_mock_server().await; + let mut builder = test_codex() + .with_model_info_override("gpt-5.6-luna", move |model| { + model.context_window = Some(window); + model + .model_messages + .as_mut() + .expect("bundled reviewer model messages") + .auto_review = Some(AutoReviewMessages { + policy: Some("Review actions against user authorization.".to_owned()), + policy_template: Some("{{ tenant_policy_config }}".to_owned()), + node_repl_policy: None, + rejection_instructions: None, + timeout_instructions: None, + }); + }) + .with_model_info_override("gpt-5.5", |model| { + model.auto_review_model_override = Some("gpt-5.6-luna".to_owned()); + }) + .with_config(move |config| { + config + .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 + .features + .enable(Feature::CurrentTimeReminder) + .expect("test permits time reminders"); + config.current_time_reminder = Some(CurrentTimeReminderConfig::default()); + config.rollout_budget = Some(RolloutBudgetConfig { + limit_tokens: 1_000_000, + reminder_at_remaining_tokens: Vec::new(), + sampling_token_weight: 1.0, + prefill_token_weight: 1.0, + }); + }); + if matches!( + reviewer_response, + ReviewerResponse::ToolContinuation | ReviewerResponse::UncompactableContinuation + ) { + builder = builder + .with_code_mode_host_program(codex_utils_cargo_bin::cargo_bin("codex-code-mode-host")?); + } + 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 = + ev_assistant_message("commentary", &"optional old commentary ".repeat(/*n*/ 900)); + commentary["item"]["phase"] = json!("commentary"); + let mut events = vec![ + sse(vec![ + ev_response_created("parent-action"), + commentary, + ev_function_call("exec-over-budget", "exec_command", &command), + ev_completed("parent-action"), + ]), + sse(vec![ + ev_response_created("parent-done"), + ev_assistant_message("done", "done"), + ev_completed("parent-done"), + ]), + ]; + if window > 1 { + events.insert( + /*index*/ 1, + 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::ToolContinuation + | ReviewerResponse::UncompactableContinuation => ev_custom_tool_call( + "reviewer-inspect", + "exec", + "text('inspection-output'.repeat(600));", + ), + }, + ev_completed("review"), + ]), + ); + } + if matches!(reviewer_response, ReviewerResponse::ToolContinuation) { + events.insert( + /*index*/ 2, + sse(vec![ + ev_assistant_message( + "after-compaction", + r#"{"risk_level":"low","user_authorization":"high","outcome":"allow"}"#, + ), + ev_completed("after-compaction"), + ]), + ); + } + if matches!( + reviewer_response, + ReviewerResponse::UncompactableContinuation | ReviewerResponse::NextReview + ) { + events.extend([ + sse(vec![ + ev_function_call("retry-command", "exec_command", &command), + ev_completed("retry-action"), + ]), + sse(vec![ + ev_assistant_message( + "retry-decision", + r#"{"risk_level":"low","user_authorization":"high","outcome":"allow"}"#, + ), + ev_completed("retry-review"), + ]), + sse(vec![ + ev_assistant_message("retry-done", "done"), + ev_completed("retry-done"), + ]), + ]); + } + let compact = if matches!(reviewer_response, ReviewerResponse::Decision) { + None + } else { + let summary = if matches!( + reviewer_response, + ReviewerResponse::UncompactableContinuation + ) { + "still oversized ".repeat(/*n*/ 4_000) + } 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 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 guardian_requests = requests + .iter() + .filter(|request| { + request.body_json()["client_metadata"]["x-openai-subagent"].as_str() == Some("guardian") + }) + .collect::>(); + if window == 1 { + assert_eq!(requests.len(), 2); + assert!(guardian_requests.is_empty()); + let output = requests[1].function_call_output("exec-over-budget"); + assert!( + output + .to_string() + .contains("evidence exceeds limits for section request_budget"), + "expected required-evidence budget rejection: {output}" + ); + } else { + let recovered = matches!(reviewer_response, ReviewerResponse::ToolContinuation); + 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!( + compact + .body_json() + .to_string() + .contains("inspection-output") + ); + let recovered = guardian_requests[1].body_json(); + assert!( + recovered["input"] + .to_string() + .contains("Previous review evidence") + ); + assert_eq!( + guardian_requests[0].body_json()["client_metadata"]["thread_id"], + recovered["client_metadata"]["thread_id"] + ); + assert!( + requests + .last() + .expect("parent resumes after review") + .function_call_output("exec-over-budget") + .to_string() + .contains("guardian-budget-test") + ); + } + if matches!( + reviewer_response, + ReviewerResponse::UncompactableContinuation + ) { + let output = requests + .last() + .expect("parent continues after the rejected review") + .function_call_output("exec-over-budget"); + assert!( + output + .to_string() + .contains("Codex ran out of room in the model's context window"), + "expected context-budget rejection: {output}" + ); + } + let request = guardian_requests[0]; + let context = request.message_input_texts("user").join("\n"); + assert!(context.contains("")); + assert!(!context.contains("optional old commentary")); + assert!(context.contains("Run the command if the approval reviewer allows it.")); + assert!(context.contains("echo guardian-budget-test")); + let developer_context = request.message_input_texts("developer").join("\n"); + assert!(developer_context.contains("")); + assert!(developer_context.contains("")); + if matches!( + reviewer_response, + ReviewerResponse::UncompactableContinuation | ReviewerResponse::NextReview + ) { + let instruction = if matches!(reviewer_response, ReviewerResponse::NextReview) { + "New instructions: retry the requested command and keep all files private. " + .repeat(/*n*/ 12) + } else { + "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); + if matches!(reviewer_response, ReviewerResponse::NextReview) { + let compact = compact.expect("scenario includes compaction"); + assert!( + !compact + .body_json() + .to_string() + .contains("New instructions:"), + "incoming evidence must remain pending during compaction" + ); + assert!( + requests[4].body_json()["input"] + .to_string() + .contains("New instructions:") + ); + assert!( + requests[4].body_json()["input"] + .to_string() + .contains("Previous review evidence") + ); + assert_eq!( + request.body_json()["client_metadata"]["thread_id"], + requests[4].body_json()["client_metadata"]["thread_id"] + ); + } else { + assert_ne!( + request.body_json()["client_metadata"]["thread_id"], + requests[4].body_json()["client_metadata"]["thread_id"], + "failed compaction must retire the reviewer" + ); + } + assert!( + requests[5] + .function_call_output("retry-command") + .to_string() + .contains("guardian-budget-test") + ); + } + } + test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index 0791c5be2e..ddaabb2d8e 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -71,6 +71,9 @@ mod external_auth; mod fork_thread; mod git_enrichment; mod guardian_authorization; +// Uses the same command-approval harness as guardian_review below. +#[cfg(not(target_os = "windows"))] +mod guardian_context_budget; mod guardian_history; mod guardian_mcp_elicitation; mod guardian_retained_context;