diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index b06a2daad3..02fd1ca8d7 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -33,6 +33,7 @@ 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::new_guardian_review_id; +pub(crate) use review::prewarm_guardian_review_session; #[cfg(test)] pub(crate) use review::record_guardian_denial_for_test; pub(crate) use review::review_approval_request; diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 99f9c5d999..bdc5b7c03a 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -23,7 +23,9 @@ use tokio::sync::oneshot; use tokio::time::Instant; use tokio::time::sleep_until; 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; @@ -46,6 +48,7 @@ use super::prompt::guardian_output_schema; use super::prompt::parse_guardian_assessment; use super::review_session::GuardianReviewSessionOutcome; use super::review_session::GuardianReviewSessionParams; +use super::review_session::GuardianReviewSessionPrewarmParams; use super::review_session::build_guardian_review_session_config; const GUARDIAN_REJECTION_INSTRUCTIONS: &str = concat!( @@ -665,40 +668,19 @@ pub(crate) fn spawn_approval_request_review( rx } -/// Runs the guardian in a locked-down reusable review session. -/// -/// The guardian itself should not mutate state or trigger further approvals, so -/// it is pinned to a read-only sandbox with `approval_policy = never` and -/// nonessential agent features disabled. When the cached trunk session is idle, -/// later approvals append onto that same guardian conversation to preserve a -/// stable prompt-cache key. If the trunk is already busy, the review runs in an -/// ephemeral fork from the last committed trunk rollout so parallel approvals -/// do not block each other or mutate the cached thread. The trunk is recreated -/// when the effective review-session config changes, and any future compaction -/// must continue to preserve the guardian policy as exact top-level developer -/// context. It may still reuse the parent's managed-network allowlist for -/// read-only checks, but it intentionally runs without inherited exec-policy -/// rules. -async fn run_guardian_review_session_before_deadline( - session: Arc, - turn: Arc, - request: GuardianApprovalRequest, - retry_reason: Option, - schema: serde_json::Value, - external_cancel: Option, - deadline: Instant, -) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { +struct GuardianReviewSessionSetup { + spawn_config: Config, + model: String, + reasoning_effort: Option, +} + +async fn guardian_review_session_setup( + session: &Arc, + turn: &Arc, +) -> anyhow::Result { 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(), - ); - } - }, + Some(network_proxy) => Some(network_proxy.proxy().current_cfg().await?), None => None, }; let available_models = session @@ -745,14 +727,67 @@ async fn run_guardian_review_session_before_deadline( reasoning_effort, ) }; - let guardian_config = build_guardian_review_session_config( + let spawn_config = build_guardian_review_session_config( turn.config.as_ref(), - live_network_config.clone(), + live_network_config, guardian_model.as_str(), guardian_reasoning_effort.clone(), - ); - let guardian_config = match guardian_config { - Ok(config) => config, + )?; + + Ok(GuardianReviewSessionSetup { + spawn_config, + model: guardian_model, + reasoning_effort: guardian_reasoning_effort, + }) +} + +pub(crate) async fn prewarm_guardian_review_session(session: Arc, turn: Arc) { + if !routes_approval_to_guardian(turn.as_ref()) { + return; + } + + let setup = match guardian_review_session_setup(&session, &turn).await { + Ok(setup) => setup, + Err(err) => { + warn!("guardian trunk prewarm setup failed: {err:#}"); + return; + } + }; + session + .guardian_review_session + .prewarm_trunk(GuardianReviewSessionPrewarmParams { + parent_session: Arc::clone(&session), + parent_turn: turn, + spawn_config: setup.spawn_config, + }) + .await; +} + +/// Runs the guardian in a locked-down reusable review session. +/// +/// The guardian itself should not mutate state or trigger further approvals, so +/// it is pinned to a read-only sandbox with `approval_policy = never` and +/// nonessential agent features disabled. When the cached trunk session is idle, +/// later approvals append onto that same guardian conversation to preserve a +/// stable prompt-cache key. If the trunk is already busy, the review runs in an +/// ephemeral fork from the last committed trunk rollout so parallel approvals +/// do not block each other or mutate the cached thread. The trunk is recreated +/// when the effective review-session config changes, and any future compaction +/// must continue to preserve the guardian policy as exact top-level developer +/// context. It may still reuse the parent's managed-network allowlist for +/// read-only checks, but it intentionally runs without inherited exec-policy +/// rules. +async fn run_guardian_review_session_before_deadline( + session: Arc, + turn: Arc, + request: GuardianApprovalRequest, + retry_reason: Option, + schema: serde_json::Value, + external_cancel: Option, + deadline: Instant, +) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { + let setup = match guardian_review_session_setup(&session, &turn).await { + Ok(setup) => setup, Err(err) => { return ( GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(err)), @@ -767,12 +802,12 @@ async fn run_guardian_review_session_before_deadline( .run_review(GuardianReviewSessionParams { parent_session: Arc::clone(&session), parent_turn: turn.clone(), - spawn_config: guardian_config, + spawn_config: setup.spawn_config, request, retry_reason, schema, - model: guardian_model, - reasoning_effort: guardian_reasoning_effort, + model: setup.model, + reasoning_effort: setup.reasoning_effort, reasoning_summary: turn.reasoning_summary, personality: turn.personality, external_cancel, diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 63b778838e..6f2d8ecf73 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -30,6 +30,7 @@ use serde_json::Value; use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; +use tracing::info; use tracing::warn; use crate::codex_delegate::run_codex_thread_interactive; @@ -84,6 +85,12 @@ pub(crate) struct GuardianReviewSessionParams { pub(crate) deadline: tokio::time::Instant, } +pub(crate) struct GuardianReviewSessionPrewarmParams { + pub(crate) parent_session: Arc, + pub(crate) parent_turn: Arc, + pub(crate) spawn_config: Config, +} + #[derive(Default)] pub(crate) struct GuardianReviewSessionManager { state: Arc>, @@ -92,9 +99,15 @@ pub(crate) struct GuardianReviewSessionManager { #[derive(Default)] struct GuardianReviewSessionState { trunk: Option>, + prewarm: Option, ephemeral_reviews: Vec>, } +struct GuardianReviewSessionPrewarm { + reuse_key: GuardianReviewSessionReuseKey, + cancel_token: CancellationToken, +} + struct GuardianReviewSession { codex: Codex, cancel_token: CancellationToken, @@ -297,13 +310,17 @@ impl GuardianReviewSessionManager { } pub(crate) async fn shutdown(&self) { - let (review_session, ephemeral_reviews) = { + let (review_session, prewarm, ephemeral_reviews) = { let mut state = self.state.lock().await; ( state.trunk.take(), + state.prewarm.take(), std::mem::take(&mut state.ephemeral_reviews), ) }; + if let Some(prewarm) = prewarm { + prewarm.cancel_token.cancel(); + } if let Some(review_session) = review_session { review_session.shutdown().await; } @@ -312,6 +329,124 @@ impl GuardianReviewSessionManager { } } + pub(super) async fn prewarm_trunk(&self, params: GuardianReviewSessionPrewarmParams) { + let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( + ¶ms.spawn_config, + params.parent_session.user_instructions().await, + ); + let cancel_token = CancellationToken::new(); + let mut stale_trunk_to_shutdown = None; + let mut stale_prewarm_to_cancel = None; + let mut should_schedule_prewarm = true; + + { + let mut state = self.state.lock().await; + if let Some(prewarm) = state.prewarm.as_ref() { + if prewarm.reuse_key == reuse_key { + return; + } + stale_prewarm_to_cancel = state.prewarm.take(); + } + + if let Some(trunk) = state.trunk.as_ref() { + if trunk.reuse_key == reuse_key { + should_schedule_prewarm = false; + } + if should_schedule_prewarm && trunk.review_lock.try_acquire().is_ok() { + stale_trunk_to_shutdown = state.trunk.take(); + } else if should_schedule_prewarm { + info!("guardian trunk prewarm skipped because a stale trunk is currently busy"); + should_schedule_prewarm = false; + } + } + + if should_schedule_prewarm { + state.prewarm = Some(GuardianReviewSessionPrewarm { + reuse_key: reuse_key.clone(), + cancel_token: cancel_token.clone(), + }); + } + } + + if let Some(prewarm) = stale_prewarm_to_cancel { + info!("cancelling stale guardian trunk prewarm before scheduling a replacement"); + prewarm.cancel_token.cancel(); + } + if let Some(review_session) = stale_trunk_to_shutdown { + review_session.shutdown_in_background(); + } + if !should_schedule_prewarm { + return; + } + + let state = Arc::clone(&self.state); + drop(tokio::spawn(async move { + let spawned_session = spawn_guardian_review_session( + Arc::clone(¶ms.parent_session), + Arc::clone(¶ms.parent_turn), + params.spawn_config, + reuse_key.clone(), + cancel_token.clone(), + /*fork_snapshot*/ None, + ) + .await; + let review_session = match spawned_session { + Ok(review_session) => Arc::new(review_session), + Err(err) => { + if !cancel_token.is_cancelled() { + warn!("guardian trunk prewarm failed: {err:#}"); + } + let mut state = state.lock().await; + if state + .prewarm + .as_ref() + .is_some_and(|prewarm| prewarm.reuse_key == reuse_key) + { + state.prewarm = None; + } + return; + } + }; + + let mut spawned_session_to_shutdown = Some(review_session); + let mut stale_trunk_to_shutdown = None; + let installed = { + let mut state = state.lock().await; + if !state + .prewarm + .as_ref() + .is_some_and(|prewarm| prewarm.reuse_key == reuse_key) + { + false + } else { + state.prewarm = None; + if let Some(trunk) = state.trunk.as_ref() + && trunk.reuse_key != reuse_key + && trunk.review_lock.try_acquire().is_ok() + { + stale_trunk_to_shutdown = state.trunk.take(); + } + if state.trunk.is_none() { + state.trunk = spawned_session_to_shutdown.take(); + true + } else { + false + } + } + }; + + if installed { + info!("guardian trunk prewarm installed"); + } + if let Some(review_session) = stale_trunk_to_shutdown { + review_session.shutdown_in_background(); + } + if let Some(review_session) = spawned_session_to_shutdown { + review_session.shutdown().await; + } + })); + } + #[expect( clippy::await_holding_invalid_type, reason = "review session selection and trunk spawning must stay serialized" @@ -335,6 +470,14 @@ impl GuardianReviewSessionManager { .await { Ok(mut state) => { + if let Some(prewarm) = state.prewarm.as_ref() + && prewarm.reuse_key != next_reuse_key + { + info!("cancelling stale guardian trunk prewarm before review"); + if let Some(prewarm) = state.prewarm.take() { + prewarm.cancel_token.cancel(); + } + } if let Some(trunk) = state.trunk.as_ref() && trunk.reuse_key != next_reuse_key && trunk.review_lock.try_acquire().is_ok() @@ -343,13 +486,20 @@ impl GuardianReviewSessionManager { } if state.trunk.is_none() { + if let Some(prewarm) = state.prewarm.take() { + info!( + "guardian trunk prewarm was not ready before review; falling back to lazy spawn" + ); + prewarm.cancel_token.cancel(); + } let spawn_cancel_token = CancellationToken::new(); let review_session = match run_before_review_deadline_with_cancel( deadline, params.external_cancel.as_ref(), &spawn_cancel_token, Box::pin(spawn_guardian_review_session( - ¶ms, + Arc::clone(¶ms.parent_session), + Arc::clone(¶ms.parent_turn), params.spawn_config.clone(), next_reuse_key.clone(), spawn_cancel_token.clone(), @@ -559,7 +709,8 @@ impl GuardianReviewSessionManager { params.external_cancel.as_ref(), &spawn_cancel_token, Box::pin(spawn_guardian_review_session( - ¶ms, + Arc::clone(¶ms.parent_session), + Arc::clone(¶ms.parent_turn), fork_config, reuse_key, spawn_cancel_token.clone(), @@ -600,7 +751,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, @@ -616,10 +768,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, @@ -1072,6 +1224,7 @@ mod tests { use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::TurnAbortedEvent; use codex_protocol::protocol::TurnCompleteEvent; + use pretty_assertions::assert_eq; async fn test_review_session() -> ( GuardianReviewSession, @@ -1259,6 +1412,75 @@ mod tests { ); } + #[tokio::test] + async fn shutdown_cancels_pending_trunk_prewarm() { + let parent_config = crate::config::test_config().await; + let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( + &parent_config, + /*user_instructions*/ None, + ); + let cancel_token = CancellationToken::new(); + let manager = GuardianReviewSessionManager { + state: Arc::new(Mutex::new(GuardianReviewSessionState { + trunk: None, + prewarm: Some(GuardianReviewSessionPrewarm { + reuse_key, + cancel_token: cancel_token.clone(), + }), + ephemeral_reviews: Vec::new(), + })), + }; + + manager.shutdown().await; + + assert!(cancel_token.is_cancelled()); + assert!(manager.state.lock().await.prewarm.is_none()); + } + + #[tokio::test] + async fn prewarm_trunk_skips_existing_matching_trunk_and_cancels_stale_prewarm() { + let params = test_review_params().await; + let reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( + ¶ms.spawn_config, + params.parent_session.user_instructions().await, + ); + let mut stale_spawn_config = params.spawn_config.clone(); + stale_spawn_config.model = Some("different-guardian-model".to_string()); + let stale_reuse_key = GuardianReviewSessionReuseKey::from_spawn_config( + &stale_spawn_config, + params.parent_session.user_instructions().await, + ); + let stale_cancel_token = CancellationToken::new(); + let (mut review_session, _tx_event, _rx_sub) = test_review_session().await; + review_session.reuse_key = reuse_key; + let manager = GuardianReviewSessionManager { + state: Arc::new(Mutex::new(GuardianReviewSessionState { + trunk: Some(Arc::new(review_session)), + prewarm: Some(GuardianReviewSessionPrewarm { + reuse_key: stale_reuse_key, + cancel_token: stale_cancel_token.clone(), + }), + ephemeral_reviews: Vec::new(), + })), + }; + + manager + .prewarm_trunk(GuardianReviewSessionPrewarmParams { + parent_session: Arc::clone(¶ms.parent_session), + parent_turn: Arc::clone(¶ms.parent_turn), + spawn_config: params.spawn_config, + }) + .await; + + assert!(stale_cancel_token.is_cancelled()); + { + let state = manager.state.lock().await; + assert!(state.prewarm.is_none()); + assert!(state.trunk.is_some()); + } + manager.shutdown().await; + } + #[tokio::test] async fn guardian_review_session_compact_scope_change_invalidates_cached_session() { let parent_config = crate::config::test_config().await; @@ -1525,6 +1747,7 @@ mod tests { let manager = GuardianReviewSessionManager { state: Arc::new(Mutex::new(GuardianReviewSessionState { trunk: Some(Arc::new(review_session)), + prewarm: None, ephemeral_reviews: Vec::new(), })), }; diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 0d3c9c033b..4f77edf62f 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -43,6 +43,7 @@ use core_test_support::responses::sse_response; use core_test_support::responses::start_mock_server; use pretty_assertions::assert_eq; use std::fs; +use std::path::PathBuf; use std::sync::Arc; use std::time::Duration; use tempfile::tempdir; @@ -68,6 +69,19 @@ where } } +async fn wait_for_guardian_trunk_rollout_path(session: &Session) -> PathBuf { + timeout(Duration::from_secs(10), async { + loop { + if let Some(path) = session.guardian_review_session.trunk_rollout_path().await { + return path; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } + }) + .await + .expect("guardian trunk prewarm should complete") +} + #[tokio::test] async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { let server = start_mock_server().await; @@ -167,6 +181,111 @@ async fn request_permissions_routes_to_guardian_when_reviewer_is_enabled() { assert!(guardian_request.body_contains_text("need network")); } +#[tokio::test] +async fn prewarmed_guardian_trunk_is_reused_by_first_review() { + let server = start_mock_server().await; + let guardian_request_log = mount_sse_once( + &server, + sse(vec![ + ev_response_created("resp-guardian"), + ev_assistant_message( + "msg-guardian", + &serde_json::json!({ + "risk_level": "low", + "user_authorization": "high", + "outcome": "allow", + "rationale": "The request grants narrowly scoped network access for this turn.", + }) + .to_string(), + ), + ev_completed("resp-guardian"), + ]), + ) + .await; + + let (mut session, mut turn_context_raw) = make_session_and_context().await; + *session.active_turn.lock().await = Some(ActiveTurn::default()); + turn_context_raw + .approval_policy + .set(AskForApproval::OnRequest) + .expect("test setup should allow updating approval policy"); + turn_context_raw + .features + .enable(Feature::GuardianApproval) + .expect("test setup should allow enabling guardian approvals"); + let mut config = (*turn_context_raw.config).clone(); + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config.model_provider.base_url = Some(format!("{}/v1", server.uri())); + let config = Arc::new(config); + let models_manager = models_manager_with_provider( + config.codex_home.to_path_buf(), + Arc::clone(&session.services.auth_manager), + config.model_provider.clone(), + ); + session.services.models_manager = models_manager; + turn_context_raw.config = Arc::clone(&config); + turn_context_raw.provider = create_model_provider( + config.model_provider.clone(), + turn_context_raw.auth_manager.clone(), + ); + let session = Arc::new(session); + let turn_context = Arc::new(turn_context_raw); + + crate::guardian::prewarm_guardian_review_session( + Arc::clone(&session), + Arc::clone(&turn_context), + ) + .await; + let prewarmed_trunk_rollout_path = wait_for_guardian_trunk_rollout_path(&session).await; + + let requested_permissions = RequestPermissionProfile { + network: Some(NetworkPermissions { + enabled: Some(true), + }), + ..RequestPermissionProfile::default() + }; + let environment = turn_context + .environments + .primary() + .expect("primary environment") + .selection(); + let response = timeout( + Duration::from_secs(45), + session.request_permissions_for_environment( + &turn_context, + "perm-call-prewarmed".to_string(), + RequestPermissionsArgs { + environment_id: None, + reason: Some("need network".to_string()), + permissions: requested_permissions.clone(), + }, + environment, + CancellationToken::new(), + ), + ) + .await + .expect("request_permissions should not wait for a client approval"); + + assert_eq!( + response, + Some(RequestPermissionsResponse { + permissions: requested_permissions.clone(), + scope: PermissionGrantScope::Turn, + strict_auto_review: false, + }) + ); + assert_eq!( + wait_for_guardian_trunk_rollout_path(&session).await, + prewarmed_trunk_rollout_path + ); + + let guardian_request = guardian_request_log.single_request(); + assert_eq!(guardian_request.path(), "/v1/responses"); + assert!(guardian_request.body_contains_text("request_permissions")); + assert!(guardian_request.body_contains_text("need network")); + session.guardian_review_session.shutdown().await; +} + #[tokio::test] async fn request_permissions_guardian_review_stops_when_cancelled() { let server = start_mock_server().await; diff --git a/codex-rs/core/src/tasks/regular.rs b/codex-rs/core/src/tasks/regular.rs index c01cae9145..035f685c52 100644 --- a/codex-rs/core/src/tasks/regular.rs +++ b/codex-rs/core/src/tasks/regular.rs @@ -54,6 +54,13 @@ impl SessionTask for RegularTask { }); sess.send_event(ctx.as_ref(), event).await; sess.set_server_reasoning_included(/*included*/ false).await; + if crate::guardian::routes_approval_to_guardian(ctx.as_ref()) { + let prewarm_sess = Arc::clone(&sess); + let prewarm_ctx = Arc::clone(&ctx); + drop(tokio::spawn(async move { + crate::guardian::prewarm_guardian_review_session(prewarm_sess, prewarm_ctx).await; + })); + } let prewarmed_client_session = match sess .consume_startup_prewarm_for_regular_turn(&cancellation_token) .await