From 2d9a36327ac75fa592cb0912bb082292bff23d9b Mon Sep 17 00:00:00 2001 From: Saud Date: Wed, 3 Jun 2026 14:48:10 -0400 Subject: [PATCH] Warm Guardian review sessions before first approval --- codex-rs/core/src/guardian/mod.rs | 3 + codex-rs/core/src/guardian/review.rs | 180 ++++++++++++------- codex-rs/core/src/guardian/review_session.rs | 99 +++++++++- codex-rs/core/src/guardian/tests.rs | 56 ++++++ codex-rs/core/src/tasks/regular.rs | 1 + codex-rs/features/src/lib.rs | 8 + codex-rs/features/src/tests.rs | 8 + 7 files changed, 286 insertions(+), 69 deletions(-) diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index b4920f1ff6..b324920576 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -32,6 +32,7 @@ pub(crate) use approval_request::guardian_approval_request_to_json; pub(crate) use review::guardian_rejection_message; pub(crate) use review::guardian_timeout_message; pub(crate) use review::is_guardian_reviewer_source; +pub(crate) use review::maybe_warm_guardian_review_session; pub(crate) use review::new_guardian_review_id; #[cfg(test)] pub(crate) use review::record_guardian_denial_for_test; @@ -170,6 +171,8 @@ use review::GuardianReviewOutcome; #[cfg(test)] use review::run_guardian_review_session as run_guardian_review_session_for_test; #[cfg(test)] +use review::warm_guardian_review_session as warm_guardian_review_session_for_test; +#[cfg(test)] use review_session::build_guardian_review_session_config as build_guardian_review_session_config_for_test; #[cfg(test)] diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 06d1a78931..9f2dfd619f 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -5,6 +5,7 @@ use codex_analytics::GuardianReviewFailureReason; use codex_analytics::GuardianReviewTerminalStatus; use codex_analytics::GuardianReviewTrackContext; use codex_analytics::GuardianReviewedAction; +use codex_features::Feature; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; @@ -18,9 +19,12 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::WarningEvent; use std::sync::Arc; +use std::time::Instant; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; +use tracing::warn; +use crate::config::Config; use crate::session::session::Session; use crate::session::turn_context::TurnContext; use crate::turn_timing::now_unix_timestamp_ms; @@ -159,6 +163,33 @@ pub(crate) fn routes_approval_to_guardian_with_reviewer( ) && approvals_reviewer == ApprovalsReviewer::AutoReview } +pub(crate) fn maybe_warm_guardian_review_session(session: Arc, turn: Arc) { + if !turn.features.enabled(Feature::GuardianSessionWarmup) + || turn.session_source.is_non_root_agent() + || !routes_approval_to_guardian(turn.as_ref()) + { + return; + } + + drop(tokio::spawn(async move { + let started_at = Instant::now(); + let result = warm_guardian_review_session(Arc::clone(&session), turn).await; + let status = match &result { + Ok(true) => "ready", + Ok(false) => "already_initialized", + Err(_) => "failed", + }; + session.services.session_telemetry.record_startup_phase( + "guardian_review_session_warmup", + started_at.elapsed(), + Some(status), + ); + if let Err(err) = result { + warn!("guardian review session warmup failed: {err:#}"); + } + })); +} + pub(crate) fn is_guardian_reviewer_source( session_source: &codex_protocol::protocol::SessionSource, ) -> bool { @@ -665,69 +696,11 @@ pub(super) async fn run_guardian_review_session( schema: serde_json::Value, external_cancel: Option, ) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { - let network_proxy = session.services.network_proxy.load_full(); - let live_network_config = match network_proxy.as_ref() { - Some(network_proxy) => match network_proxy.proxy().current_cfg().await { - Ok(config) => Some(config), - Err(err) => { - return ( - GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(err)), - GuardianReviewAnalyticsResult::without_session(), - ); - } - }, - None => None, - }; - let available_models = session - .services - .models_manager - .list_models(codex_models_manager::manager::RefreshStrategy::Offline) - .await; - let preferred_reasoning_effort = |supports_low: bool, fallback| { - if supports_low { - Some(codex_protocol::openai_models::ReasoningEffort::Low) - } else { - fallback - } - }; - let model_override = turn.model_info.auto_review_model_override.as_deref(); - let review_model_id = - model_override.unwrap_or_else(|| turn.provider.approval_review_preferred_model()); - let review_model = available_models - .iter() - .find(|preset| preset.model == review_model_id); - let (guardian_model, guardian_reasoning_effort) = if let Some(preset) = review_model { - let reasoning_effort = preferred_reasoning_effort( - preset - .supported_reasoning_efforts - .iter() - .any(|effort| effort.effort == codex_protocol::openai_models::ReasoningEffort::Low), - Some(preset.default_reasoning_effort), - ); - (review_model_id.to_string(), reasoning_effort) - } else { - let reasoning_effort = preferred_reasoning_effort( - turn.model_info - .supported_reasoning_levels - .iter() - .any(|preset| preset.effort == codex_protocol::openai_models::ReasoningEffort::Low), - turn.reasoning_effort - .or(turn.model_info.default_reasoning_level), - ); - ( - model_override - .unwrap_or(turn.model_info.slug.as_str()) - .to_string(), - reasoning_effort, - ) - }; - let guardian_config = build_guardian_review_session_config( - turn.config.as_ref(), - live_network_config.clone(), - guardian_model.as_str(), + let ResolvedGuardianReviewSessionConfig { + guardian_config, + guardian_model, guardian_reasoning_effort, - ); - let guardian_config = match guardian_config { + } = match resolve_guardian_review_session_config(session.as_ref(), turn.as_ref()).await { Ok(config) => config, Err(err) => { return ( @@ -801,6 +774,89 @@ pub(super) async fn run_guardian_review_session( } } +struct ResolvedGuardianReviewSessionConfig { + guardian_config: Config, + guardian_model: String, + guardian_reasoning_effort: Option, +} + +pub(super) async fn warm_guardian_review_session( + session: Arc, + turn: Arc, +) -> anyhow::Result { + let resolved_config = + resolve_guardian_review_session_config(session.as_ref(), turn.as_ref()).await?; + session + .guardian_review_session + .warm_trunk(Arc::clone(&session), turn, resolved_config.guardian_config) + .await +} + +async fn resolve_guardian_review_session_config( + session: &Session, + turn: &TurnContext, +) -> anyhow::Result { + let network_proxy = session.services.network_proxy.load_full(); + let live_network_config = match network_proxy.as_ref() { + Some(network_proxy) => Some(network_proxy.proxy().current_cfg().await?), + None => None, + }; + let available_models = session + .services + .models_manager + .list_models(codex_models_manager::manager::RefreshStrategy::Offline) + .await; + let preferred_reasoning_effort = |supports_low: bool, fallback| { + if supports_low { + Some(codex_protocol::openai_models::ReasoningEffort::Low) + } else { + fallback + } + }; + let model_override = turn.model_info.auto_review_model_override.as_deref(); + let review_model_id = + model_override.unwrap_or_else(|| turn.provider.approval_review_preferred_model()); + let review_model = available_models + .iter() + .find(|preset| preset.model == review_model_id); + let (guardian_model, guardian_reasoning_effort) = if let Some(preset) = review_model { + let reasoning_effort = preferred_reasoning_effort( + preset + .supported_reasoning_efforts + .iter() + .any(|effort| effort.effort == codex_protocol::openai_models::ReasoningEffort::Low), + Some(preset.default_reasoning_effort), + ); + (review_model_id.to_string(), reasoning_effort) + } else { + let reasoning_effort = preferred_reasoning_effort( + turn.model_info + .supported_reasoning_levels + .iter() + .any(|preset| preset.effort == codex_protocol::openai_models::ReasoningEffort::Low), + turn.reasoning_effort + .or(turn.model_info.default_reasoning_level), + ); + ( + model_override + .unwrap_or(turn.model_info.slug.as_str()) + .to_string(), + reasoning_effort, + ) + }; + let guardian_config = build_guardian_review_session_config( + turn.config.as_ref(), + live_network_config, + guardian_model.as_str(), + guardian_reasoning_effort, + )?; + Ok(ResolvedGuardianReviewSessionConfig { + guardian_config, + guardian_model, + guardian_reasoning_effort, + }) +} + #[cfg(test)] mod review_tests { use super::*; diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index c139a910a7..f61f5ae469 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -87,6 +87,7 @@ pub(crate) struct GuardianReviewSessionManager { struct GuardianReviewSessionState { trunk: Option>, ephemeral_reviews: Vec>, + shutdown: bool, } struct GuardianReviewSession { @@ -290,6 +291,7 @@ impl GuardianReviewSessionManager { pub(crate) async fn shutdown(&self) { let (review_session, ephemeral_reviews) = { let mut state = self.state.lock().await; + state.shutdown = true; ( state.trunk.take(), std::mem::take(&mut state.ephemeral_reviews), @@ -303,6 +305,44 @@ impl GuardianReviewSessionManager { } } + #[expect( + clippy::await_holding_invalid_type, + reason = "warmup and review session trunk spawning must stay serialized" + )] + pub(crate) async fn warm_trunk( + &self, + parent_session: Arc, + parent_turn: Arc, + spawn_config: Config, + ) -> anyhow::Result { + let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config(&spawn_config); + let mut state = self.state.lock().await; + if state.shutdown || state.trunk.is_some() { + return Ok(false); + } + + let spawn_cancel_token = CancellationToken::new(); + let review_session = run_before_review_deadline_with_cancel( + tokio::time::Instant::now() + GUARDIAN_REVIEW_TIMEOUT, + /*external_cancel*/ None, + &spawn_cancel_token, + Box::pin(spawn_guardian_review_session( + &parent_session, + &parent_turn, + spawn_config, + reuse_key, + spawn_cancel_token.clone(), + /*fork_snapshot*/ None, + )), + ) + .await + .map_err(|outcome| { + anyhow!("guardian review session warmup did not complete: {outcome:?}") + })??; + state.trunk = Some(Arc::new(review_session)); + Ok(true) + } + #[expect( clippy::await_holding_invalid_type, reason = "review session selection and trunk spawning must stay serialized" @@ -337,7 +377,8 @@ impl GuardianReviewSessionManager { params.external_cancel.as_ref(), &spawn_cancel_token, Box::pin(spawn_guardian_review_session( - ¶ms, + ¶ms.parent_session, + ¶ms.parent_turn, params.spawn_config.clone(), next_reuse_key.clone(), spawn_cancel_token.clone(), @@ -543,7 +584,8 @@ impl GuardianReviewSessionManager { params.external_cancel.as_ref(), &spawn_cancel_token, Box::pin(spawn_guardian_review_session( - ¶ms, + ¶ms.parent_session, + ¶ms.parent_turn, fork_config, reuse_key, spawn_cancel_token.clone(), @@ -582,7 +624,8 @@ impl GuardianReviewSessionManager { } async fn spawn_guardian_review_session( - params: &GuardianReviewSessionParams, + parent_session: &Arc, + parent_turn: &Arc, spawn_config: Config, reuse_key: GuardianReviewSessionReuseKey, cancel_token: CancellationToken, @@ -598,10 +641,10 @@ async fn spawn_guardian_review_session( }; let codex = Box::pin(run_codex_thread_interactive( spawn_config, - params.parent_session.services.auth_manager.clone(), - params.parent_session.services.models_manager.clone(), - Arc::clone(¶ms.parent_session), - Arc::clone(¶ms.parent_turn), + parent_session.services.auth_manager.clone(), + parent_session.services.models_manager.clone(), + Arc::clone(parent_session), + Arc::clone(parent_turn), cancel_token.clone(), SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_string()), initial_history, @@ -958,6 +1001,7 @@ pub(crate) fn build_guardian_review_session_config( Feature::Plugins, Feature::WebSearchRequest, Feature::WebSearchCached, + Feature::GuardianSessionWarmup, ] { guardian_config.features.disable(feature).map_err(|err| { anyhow::anyhow!( @@ -1141,6 +1185,24 @@ mod tests { } } + #[tokio::test] + async fn warm_trunk_does_not_spawn_after_shutdown() { + let manager = GuardianReviewSessionManager::default(); + manager.shutdown().await; + let params = test_review_params().await; + + assert!( + !manager + .warm_trunk( + params.parent_session, + params.parent_turn, + params.spawn_config, + ) + .await + .expect("warmup after shutdown should be a no-op") + ); + } + #[tokio::test] async fn guardian_review_session_config_change_invalidates_cached_session() { let parent_config = crate::config::test_config().await; @@ -1260,6 +1322,29 @@ mod tests { assert!(!guardian_config.features.enabled(Feature::CodexHooks)); } + #[tokio::test] + async fn guardian_review_session_config_disables_warmup() { + let mut parent_config = crate::config::test_config().await; + parent_config + .features + .enable(Feature::GuardianSessionWarmup) + .expect("enable guardian session warmup on parent config"); + + let guardian_config = build_guardian_review_session_config( + &parent_config, + /*live_network_config*/ None, + "active-model", + /*reasoning_effort*/ None, + ) + .expect("guardian config"); + + assert!( + !guardian_config + .features + .enabled(Feature::GuardianSessionWarmup) + ); + } + #[tokio::test] async fn guardian_review_session_config_disables_skill_instructions() { let mut parent_config = crate::config::test_config().await; diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index ffadd801ca..0c916024de 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -1575,6 +1575,62 @@ async fn build_guardian_prompt_items_includes_parent_session_id() -> anyhow::Res Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn warmed_guardian_trunk_is_reused_by_first_review() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = start_mock_server().await; + let request_log = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-guardian"), + ev_assistant_message( + "msg-guardian", + "{\"risk_level\":\"low\",\"user_authorization\":\"high\",\"outcome\":\"allow\",\"rationale\":\"warmed guardian rationale\"}", + ), + ev_completed("resp-guardian"), + ]), + ) + .await; + let (session, turn) = guardian_test_session_and_turn(&server).await; + + assert!(warm_guardian_review_session_for_test(Arc::clone(&session), Arc::clone(&turn)).await?); + seed_guardian_parent_history(&session, &turn).await; + let outcome = run_guardian_review_session_for_test( + Arc::clone(&session), + Arc::clone(&turn), + GuardianApprovalRequest::Shell { + id: "shell-1".to_string(), + command: vec!["git".to_string(), "push".to_string()], + cwd: test_path_buf("/repo/codex-rs/core").abs(), + sandbox_permissions: crate::sandboxing::SandboxPermissions::UseDefault, + additional_permissions: None, + justification: Some("Need to push the docs fix.".to_string()), + }, + /*retry_reason*/ None, + guardian_output_schema(), + /*external_cancel*/ None, + ) + .await; + let (GuardianReviewOutcome::Completed(assessment), metadata) = outcome else { + panic!("expected guardian assessment"); + }; + + assert_eq!(assessment.outcome, GuardianAssessmentOutcome::Allow); + assert!(matches!( + metadata.guardian_session_kind, + Some(codex_analytics::GuardianReviewSessionKind::TrunkReused) + )); + assert_eq!(metadata.had_prior_review_context, Some(false)); + assert_eq!( + request_log.requests().len(), + 1, + "warmup should initialize the trunk without performing a review" + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow::Result<()> { skip_if_no_network!(Ok(())); diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index c01cae9145..2fbbb52e4d 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -54,6 +54,7 @@ impl SessionTask for RegularTask { }); sess.send_event(ctx.as_ref(), event).await; sess.set_server_reasoning_included(/*included*/ false).await; + crate::guardian::maybe_warm_guardian_review_session(Arc::clone(&sess), Arc::clone(&ctx)); let prewarmed_client_session = match sess .consume_startup_prewarm_for_regular_turn(&cancellation_token) .await diff --git a/codex-rs/features/src/lib.rs b/codex-rs/features/src/lib.rs index ce1ee6940e..63a8fbe027 100644 --- a/codex-rs/features/src/lib.rs +++ b/codex-rs/features/src/lib.rs @@ -190,6 +190,8 @@ pub enum Feature { DefaultModeRequestUserInput, /// Enable automatic review for approval prompts. GuardianApproval, + /// Warm the reusable automatic-review session before its first approval prompt. + GuardianSessionWarmup, /// Enable persisted thread goals and automatic goal continuation. Goals, /// Route MCP tool approval prompts through the MCP elicitation request path. @@ -1117,6 +1119,12 @@ pub const FEATURES: &[FeatureSpec] = &[ stage: Stage::Stable, default_enabled: true, }, + FeatureSpec { + id: Feature::GuardianSessionWarmup, + key: "guardian_session_warmup", + stage: Stage::UnderDevelopment, + default_enabled: false, + }, FeatureSpec { id: Feature::Goals, key: "goals", diff --git a/codex-rs/features/src/tests.rs b/codex-rs/features/src/tests.rs index 5d7087e90b..1c9d28e2c4 100644 --- a/codex-rs/features/src/tests.rs +++ b/codex-rs/features/src/tests.rs @@ -101,6 +101,14 @@ fn guardian_approval_is_stable_and_enabled_by_default() { assert_eq!(Feature::GuardianApproval.default_enabled(), true); } +#[test] +fn guardian_session_warmup_is_under_development_and_disabled_by_default() { + let spec = Feature::GuardianSessionWarmup.info(); + + assert_eq!(spec.stage, Stage::UnderDevelopment); + assert_eq!(Feature::GuardianSessionWarmup.default_enabled(), false); +} + #[test] fn external_migration_is_experimental_and_disabled_by_default() { let spec = Feature::ExternalMigration.info();