From 18d7ace2210a1cc65c0de85a2a6af6905e26d634 Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 14 Sep 2026 22:35:08 +0000 Subject: [PATCH] Move Guardian reviewer lifecycle into the extension (#45537) ## Why Guardian reviews must stop when their parent shuts down or their history is reset. Shutdown must also finish reviewer cleanup before closing the parent's persistent history, including when a review is waiting to retry after a rate limit. ## What changed - Let the Guardian extension own reviewer startup, prewarming, and shutdown through `ThreadManager`, tracking and joining outstanding work during teardown. - Cancel reviews on history reset or parent shutdown and reject decisions returned after cancellation, including cached extension decisions. - Keep background prewarm previews from overwriting the active turn's model metadata. ## Testing Extend regression coverage to verify reviewer cleanup after parent shutdown, prompt shutdown during a 60-second Guardian retry with the network request denied, and preservation of active model review requirements during prewarming. GitOrigin-RevId: 05efc369b9a1642a10365eac8a09da77f3f6e28c --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/codex_thread.rs | 2 +- codex-rs/core/src/exec.rs | 8 +- codex-rs/core/src/guardian/decision.rs | 39 ++++- codex-rs/core/src/guardian/mod.rs | 8 +- codex-rs/core/src/guardian/review.rs | 25 ++- codex-rs/core/src/guardian/review_request.rs | 31 ++-- codex-rs/core/src/guardian/review_session.rs | 28 +--- .../core/src/guardian/review_session_setup.rs | 123 +++++++-------- .../core/src/guardian/review_session_tests.rs | 39 +---- .../src/guardian/review_session_threads.rs | 96 ----------- codex-rs/core/src/guardian/runtime.rs | 2 + codex-rs/core/src/guardian/test_host.rs | 49 +++++- codex-rs/core/src/guardian/tests.rs | 18 ++- codex-rs/core/src/guardian_review.rs | 9 +- codex-rs/core/src/session/handlers.rs | 7 +- codex-rs/core/src/session/mod.rs | 21 ++- codex-rs/core/src/session/turn_context.rs | 8 +- codex-rs/core/src/session/turn_suspension.rs | 5 +- codex-rs/core/src/session_startup_prewarm.rs | 15 -- codex-rs/core/src/state/session.rs | 18 ++- codex-rs/core/tests/suite/guardian_review.rs | 13 ++ codex-rs/core/tests/suite/network_approval.rs | 68 +++++++- .../ext/extension-api/src/contributors.rs | 3 +- .../src/contributors/thread_lifecycle.rs | 2 +- codex-rs/ext/guardian-reviewer/Cargo.toml | 2 +- codex-rs/ext/guardian-reviewer/src/lib.rs | 1 + codex-rs/ext/guardian-reviewer/src/pool.rs | 8 + codex-rs/ext/guardian-v2/Cargo.toml | 3 +- .../src/async_scorer/extension_tests.rs | 2 + .../ext/guardian-v2/src/sync_reviewer/mod.rs | 149 +++++++++++++++++- 31 files changed, 495 insertions(+), 308 deletions(-) delete mode 100644 codex-rs/core/src/guardian/review_session_threads.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 65de5641b4..14d6dc3e5e 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3596,6 +3596,7 @@ dependencies = [ "codex-extension-api", "codex-features", "codex-guardian-context", + "codex-guardian-reviewer", "codex-history", "codex-http-client", "codex-login", diff --git a/codex-rs/core/src/codex_thread.rs b/codex-rs/core/src/codex_thread.rs index 50ad565c3c..76c5ad7df9 100644 --- a/codex-rs/core/src/codex_thread.rs +++ b/codex-rs/core/src/codex_thread.rs @@ -692,7 +692,7 @@ impl CodexThread { pub async fn guardian_trunk_rollout_path(&self) -> Option { self.session - .guardian_review_session() + .guardian_review_session()? .trunk() .await? .rollout_path() diff --git a/codex-rs/core/src/exec.rs b/codex-rs/core/src/exec.rs index 78c0076f0d..2b47786397 100644 --- a/codex-rs/core/src/exec.rs +++ b/codex-rs/core/src/exec.rs @@ -254,12 +254,16 @@ pub(crate) fn cancel_when_either( first: CancellationToken, second: CancellationToken, ) -> CancellationToken { - let combined = CancellationToken::new(); + let combined = first.child_token(); + if combined.is_cancelled() || second.is_cancelled() { + combined.cancel(); + return combined; + } let cancel = combined.clone(); tokio::spawn(async move { tokio::select! { - _ = first.cancelled() => {} _ = second.cancelled() => {} + _ = cancel.cancelled() => {} } cancel.cancel(); }); diff --git a/codex-rs/core/src/guardian/decision.rs b/codex-rs/core/src/guardian/decision.rs index a1d0026317..2779900d97 100644 --- a/codex-rs/core/src/guardian/decision.rs +++ b/codex-rs/core/src/guardian/decision.rs @@ -54,10 +54,35 @@ pub(crate) async fn decide_approval( review_id: String, request: impl Into, reasons: ApprovalRequestReasons, - options: GuardianReviewOptions, + mut options: GuardianReviewOptions, ) -> Option { + let runtime = session + .services + .thread_extension_data + .get::(); + let _task = runtime.as_ref().map(|runtime| runtime.tasks.token()); + if runtime + .as_ref() + .is_some_and(|runtime| runtime.cancellation.is_cancelled()) + { + return Some(ReviewDecision::Abort); + } let context = context.into(); let request = request.into(); + let (_, history_reset) = session.history_reset().await; + let cancellation = match options.external_cancel.take() { + Some(external) => crate::exec::cancel_when_either(external, history_reset.clone()), + None => history_reset.child_token(), + }; + let _cancel_on_drop = cancellation.clone().drop_guard(); + let review_cancellation = match runtime.as_ref() { + Some(runtime) => { + crate::exec::cancel_when_either(runtime.cancellation.clone(), cancellation.clone()) + } + None => cancellation.clone(), + }; + let _review_cancel_on_drop = review_cancellation.clone().drop_guard(); + options.external_cancel = Some(review_cancellation.clone()); let turn = context.turn(); let live_config = session.get_config().await; let requirements = live_config.config_layer_stack.requirements(); @@ -109,6 +134,7 @@ pub(crate) async fn decide_approval( }; let runtime = codex_guardian_reviewer::SynchronousReview::new(ReviewRuntime { session: Arc::clone(&session), + history_reset: history_reset.clone(), context: context.clone(), review_id: review_id.clone(), request: request.clone(), @@ -148,7 +174,7 @@ pub(crate) async fn decide_approval( )), synchronous_reviewer: &runtime, }; - match session.services.extensions.decide_approval(&input).await { + let decision = match session.services.extensions.decide_approval(&input).await { Some(ApprovalDecision::Reviewed(decision)) => Some(decision), Some(ApprovalDecision::Allow) if !require_fresh_review => { let request = match request.validate(&context) { @@ -193,5 +219,14 @@ pub(crate) async fn decide_approval( None | Some(ApprovalDecision::AskUser) => { runtime.review(GuardianReviewReason::Policy).await } + }; + // Enforce cancellation after extension callbacks too, including cached decisions. + if history_reset.is_cancelled() + || cancellation.is_cancelled() + || review_cancellation.is_cancelled() + { + Some(ReviewDecision::Abort) + } else { + decision } } diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 91a58efc57..a4e0e1429a 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -56,10 +56,12 @@ pub(crate) use review::new_guardian_review_id; #[cfg(test)] pub(crate) use review::record_guardian_denial_for_test; pub(crate) use review::routes_approval_policy_to_guardian; -pub(crate) use review::routes_approval_to_guardian; -pub use review_session::GuardianReviewSessionHost; +pub use review_session::GuardianReviewSession; pub(crate) use review_session::GuardianReviewSessionManager; -pub(crate) use review_session::prewarm_guardian_review_session; +pub use review_session::GuardianReviewState; +pub use review_session::PreparedGuardianContext; +pub use review_session::prepare_review_prewarm; + pub(crate) use review_session::prompt_cache_key_override_for_review_session; pub(crate) use runtime::ReviewAction; diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index ff2a94763a..fc7960d547 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -102,21 +102,6 @@ pub(crate) fn new_guardian_review_id() -> String { uuid::Uuid::new_v4().to_string() } -/// Whether this turn should route allowed approval prompts through the guardian -/// reviewer instead of surfacing them to the user. ARC may still block actions -/// earlier in the flow. -pub(crate) fn routes_approval_to_guardian(turn: &TurnContext) -> bool { - routes_approval_to_guardian_with_reviewer(turn, turn.config.approvals_reviewer) -} - -/// Whether an approval with its own reviewer selection should be routed through guardian. -pub(crate) fn routes_approval_to_guardian_with_reviewer( - turn: &TurnContext, - approvals_reviewer: ApprovalsReviewer, -) -> bool { - routes_approval_policy_to_guardian(turn.approval_policy(), approvals_reviewer) -} - /// Whether an exact approval policy and reviewer should route through Guardian. pub(crate) fn routes_approval_policy_to_guardian( approval_policy: AskForApproval, @@ -333,6 +318,14 @@ async fn run_guardian_review_session_before_deadline( external_cancel: Option, deadline: Instant, ) -> (GuardianReviewOutcome, GuardianReviewAnalyticsResult) { + let Some(pool) = session.guardian_review_session() else { + return ( + GuardianReviewOutcome::Error(GuardianReviewError::prompt_build(anyhow::anyhow!( + "Guardian extension is not installed for this thread" + ))), + GuardianReviewAnalyticsResult::without_session(), + ); + }; let session_config = match guardian_review_session_config(session.as_ref(), &context).await { Ok(session_config) => session_config, Err(err) => { @@ -344,7 +337,7 @@ async fn run_guardian_review_session_before_deadline( }; let (session_outcome, session_analytics_result) = Box::pin(super::review_session::run_guardian_review_session( - session.guardian_review_session(), + pool, GuardianReviewSessionParams { parent_session: Arc::clone(&session), parent_context: context.clone(), diff --git a/codex-rs/core/src/guardian/review_request.rs b/codex-rs/core/src/guardian/review_request.rs index 4d0426afd6..ce6c9c3f64 100644 --- a/codex-rs/core/src/guardian/review_request.rs +++ b/codex-rs/core/src/guardian/review_request.rs @@ -34,6 +34,7 @@ impl ReviewHost for super::super::runtime::ReviewRuntime { ) -> Result { let super::super::runtime::ReviewRuntime { session, + history_reset: _, context, review_id, request, @@ -204,20 +205,26 @@ impl ReviewHost for super::super::runtime::ReviewRuntime { review_evidence, } = prepared; let session = Arc::clone(&self.session); - if session.guardian_context_mode == GuardianContextMode::ThreadOwned - && matches!(&outcome, GuardianReviewOutcome::Completed(assessment) if assessment.outcome == GuardianAssessmentOutcome::Allow) - && (root_authorization_version - != session - .services - .agent_control - .root_user_authorization(session.thread_id) - .await - .map(|snapshot| snapshot.authorization_version) - || user_message_revision + if matches!(&outcome, GuardianReviewOutcome::Completed(assessment) if assessment.outcome == GuardianAssessmentOutcome::Allow) + && ((session.guardian_context_mode == GuardianContextMode::ThreadOwned + && (root_authorization_version != session - .conversation_history_snapshot() + .services + .agent_control + .root_user_authorization(session.thread_id) .await - .user_message_revision()) + .map(|snapshot| snapshot.authorization_version) + || user_message_revision + != session + .conversation_history_snapshot() + .await + .user_message_revision())) + || self.history_reset.is_cancelled() + || self + .options + .external_cancel + .as_ref() + .is_some_and(CancellationToken::is_cancelled)) { // A completed approval cannot outlive the owning-session or root evidence // it evaluated, including when either changed before prompt construction. diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 114889e669..23e756fcf4 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -4,12 +4,10 @@ #[path = "review_session_setup.rs"] mod setup; -pub(crate) use setup::prewarm_guardian_review_session; +pub use setup::PreparedGuardianContext; +pub use setup::prepare_review_prewarm; pub(crate) use setup::run_guardian_review_session; -#[path = "review_session_threads.rs"] -mod managed_threads; - #[path = "review_session_context.rs"] mod context_policy; use context_policy::ReviewContextPolicy; @@ -126,26 +124,11 @@ pub(crate) struct GuardianReviewSessionParams { pub(crate) deadline: tokio::time::Instant, } -/// Host capability used to spawn private reviewer runtimes for this parent. -/// The extension owns pooling; this adapter keeps the existing context and runtime paths. -#[derive(Default)] -pub struct GuardianReviewSessionHost { - managed_threads: Option, -} - -impl GuardianReviewSessionHost { - pub fn with_thread_manager(manager: std::sync::Weak) -> Self { - Self { - managed_threads: Some(managed_threads::ManagedReviewerThreads::new(manager)), - } - } -} - pub(crate) type GuardianReviewSessionManager = codex_guardian_reviewer::ReviewerPool; /// Opaque host session handle. Its state belongs to the existing context builder. -pub(crate) struct GuardianReviewSession { +pub struct GuardianReviewSession { session: Arc, io: SessionIo, cancel_token: CancellationToken, @@ -153,7 +136,8 @@ pub(crate) struct GuardianReviewSession { state: Mutex, } -struct GuardianReviewState { +/// Opaque conversation progress retained while ThreadManager starts a reviewer. +pub struct GuardianReviewState { conversation: ConversationState, last_admitted_node_repl_response_sequence: u64, pending_node_repl_evidence_admission: Option, @@ -882,7 +866,7 @@ impl codex_guardian_reviewer::ReviewerRuntime for GuardianReviewSession { mod tests; impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession { - type Setup = setup::PreparedGuardianContext; + type Setup = PreparedGuardianContext; type Context = GuardianReviewSessionReuseKey; type Snapshot = GuardianReviewForkSnapshot; diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index 394c184048..edf191bfe0 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -1,4 +1,4 @@ -//! Prepares opaque session inputs for the extension-owned reviewer pool. +//! Captures the context used by Guardian before and after it starts a reviewer agent. //! Context selection and assembly stay on the existing path pending its replacement. use super::*; @@ -6,14 +6,14 @@ use codex_guardian_reviewer::ReviewerPool; use codex_guardian_reviewer::ReviewerRequest; use codex_guardian_reviewer::SessionDisposition; -pub(crate) struct PreparedGuardianContext { +pub struct PreparedGuardianContext { parent: Arc, context: GuardianReviewContext, config: Config, context_policy: ReviewContextPolicy, key: GuardianReviewSessionReuseKey, parent_compaction: Option, - host: Arc, + pub history_reset: CancellationToken, } impl PreparedGuardianContext { @@ -25,6 +25,11 @@ impl PreparedGuardianContext { node_repl_policy: &GuardianNodeReplPolicy, compaction_model_hash: Option<&str>, ) -> anyhow::Result { + let (reset_version, history_reset) = parent.history_reset().await; + // Preparation may have raced with a history reset before capturing the reviewer context. + if reset_version != history.reset_version { + history_reset.cancel(); + } let context_policy = ReviewContextPolicy::for_context(parent.guardian_context_mode, &config.features); let root_authorization_version = context_policy.root_authorization_version(&parent).await; @@ -40,10 +45,6 @@ impl PreparedGuardianContext { .with_node_repl_policy(node_repl_policy); key.root_authorization_version = root_authorization_version; key.parent_reset_version = history.reset_version; - let host = parent - .services - .thread_extension_data - .get_or_init(GuardianReviewSessionHost::default); Ok(Self { parent, context, @@ -51,13 +52,16 @@ impl PreparedGuardianContext { context_policy, key, parent_compaction, - host, + history_reset, }) } } impl PreparedGuardianContext { - fn reuse_key(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey { + pub fn reuse_key( + &self, + previous: Option<&GuardianReviewSession>, + ) -> GuardianReviewSessionReuseKey { let mut key = self.key.clone(); if self.context_policy != ReviewContextPolicy::ThreadOwned && self.parent_compaction.is_none() @@ -70,12 +74,21 @@ impl PreparedGuardianContext { key } - /// Returns captured parent inputs; Guardian selects the agent's identity and lifecycle. - async fn thread_options( + /// Moves fork history into startup options and retains only its context bookkeeping. + /// Guardian selects the agent's identity and lifecycle. + pub async fn thread_options( &self, - initial_history: Option, - ) -> crate::StartThreadOptions { - let initial_history = initial_history.or_else(|| { + snapshot: Option, + ) -> (crate::StartThreadOptions, GuardianReviewState) { + let (conversation, history) = snapshot.map(ConversationState::fork).unzip(); + let state = GuardianReviewState { + conversation: conversation.unwrap_or_default(), + last_admitted_node_repl_response_sequence: history.as_ref().map_or(0, |history| { + history.last_admitted_node_repl_response_sequence + }), + pending_node_repl_evidence_admission: None, + }; + let initial_history = history.map(|history| history.initial_history).or_else(|| { self.parent_compaction .clone() .map(|item| InitialHistory::Forked(vec![RolloutItem::ResponseItem(item.into())])) @@ -86,7 +99,7 @@ impl PreparedGuardianContext { .services .model_client .responses_websocket_enabled(); - crate::StartThreadOptions { + let options = crate::StartThreadOptions { internal_parent: Some(crate::thread_manager::InternalSessionParent { thread_id: self.parent.thread_id(), auth_manager: Arc::clone(&self.parent.services.auth_manager), @@ -99,24 +112,40 @@ impl PreparedGuardianContext { inherited_environments: Some(self.context.environments().clone()), client_mcp_extensions: self.parent.services.client_mcp_extensions.clone(), ..crate::StartThreadOptions::new(config) - } + }; + (options, state) } /// Binds context bookkeeping to an agent that Guardian has already started. - async fn bind_session( + pub async fn bind_thread( &self, - session: Arc, - io: SessionIo, + thread: &crate::CodexThread, context: GuardianReviewSessionReuseKey, state: GuardianReviewState, cancellation: CancellationToken, ) -> GuardianReviewSession { + let session = Arc::clone(&thread.session); + let io = SessionIo { + tx_sub: thread.io.tx_sub.clone(), + rx_event: thread.io.rx_event.clone(), + agent_status: thread.io.agent_status.clone(), + session_loop_termination: thread.io.session_loop_termination.clone(), + }; let inherited = session.inherited_instructions().await; let context = GuardianReviewSessionReuseKey { user_instructions: inherited.user, thread_instructions: inherited.thread, ..context }; + crate::session::emit_subagent_session_started( + &self.parent.services.analytics_events_client, + self.parent.app_server_client_metadata().await, + session.session_id(), + session.thread_id(), + Some(self.parent.thread_id()), + session.thread_config_snapshot().await, + SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_owned()), + ); GuardianReviewSession { session, io, @@ -127,40 +156,6 @@ impl PreparedGuardianContext { } } -impl PreparedGuardianContext { - pub(crate) async fn spawn( - &self, - context: GuardianReviewSessionReuseKey, - kind: GuardianReviewSessionKind, - snapshot: Option, - cancellation: CancellationToken, - ) -> anyhow::Result { - let (conversation, history) = snapshot.map(ConversationState::fork).unzip(); - let state = GuardianReviewState { - conversation: conversation.unwrap_or_default(), - last_admitted_node_repl_response_sequence: history.as_ref().map_or(0, |history| { - history.last_admitted_node_repl_response_sequence - }), - pending_node_repl_evidence_admission: None, - }; - let mut options = self - .thread_options(history.map(|history| history.initial_history)) - .await; - if matches!(kind, GuardianReviewSessionKind::EphemeralForked) { - options.config.ephemeral = true; - } - let threads = self.host.managed_threads.as_ref().ok_or_else(|| { - anyhow::anyhow!("Guardian extension is not installed for this thread") - })?; - let (session, io) = threads - .spawn(&self.parent, options, cancellation.clone()) - .await?; - Ok(self - .bind_session(session, io, context, state, cancellation) - .await) - } -} - pub(super) struct PreparedReview { context: Arc, params: GuardianReviewSessionParams, @@ -239,18 +234,16 @@ pub(super) async fn prepare_review( }) } -pub(crate) fn prewarm_guardian_review_session( - parent: Arc, - turn: Arc, -) -> BoxFuture<'static, anyhow::Result<()>> { - Box::pin(async move { - let context = prepare_prewarm(Arc::clone(&parent), turn).await?; - let key = context.reuse_key(/*previous*/ None); - parent - .guardian_review_session() - .prewarm(Arc::new(context), key) - .await - }) +/// Captures the same startup context used by the existing prompt builder. +/// The caller owns scheduling, cancellation, and the reviewer pool. +pub async fn prepare_review_prewarm( + parent: &crate::CodexThread, +) -> anyhow::Result { + let turn = parent + .session + .new_startup_prewarm_turn_with_sub_id(crate::session::INITIAL_SUBMIT_ID.to_owned()) + .await; + prepare_prewarm(Arc::clone(&parent.session), turn).await } pub(super) fn prepare_prewarm( diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index e99129d364..6616a291b3 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -251,30 +251,6 @@ async fn test_review_params() -> GuardianReviewSessionParams { } } -#[tokio::test] -async fn spawned_guardian_session_preserves_windows_sandbox_proxy_settings() { - let params = test_review_params().await; - let manager = params.parent_session.guardian_review_session(); - prewarm_guardian_review_session( - params.parent_session, - Arc::clone(params.parent_context.turn()), - ) - .await - .expect("initialize Guardian session"); - let mode = manager - .trunk() - .await - .expect("Guardian session") - .session - .windows_sandbox_proxy_settings_mode; - - assert_eq!( - mode, - codex_sandboxing::WindowsSandboxProxySettingsMode::Preserve - ); - manager.shutdown().await; -} - #[tokio::test] async fn spawned_guardian_reuse_key_matches_inherited_instructions() { let mut params = test_review_params().await; @@ -310,17 +286,16 @@ async fn spawned_guardian_reuse_key_matches_inherited_instructions() { thread_instructions: latest.clone(), ..stale_key.clone() }; + let manager = params + .parent_session + .guardian_review_session() + .expect("Guardian pool installed"); let prepared = setup::prepare_review(params).await.expect("prepare review"); - let review = prepared - .setup() - .spawn( - stale_key, - GuardianReviewSessionKind::TrunkNew, - /*snapshot*/ None, - CancellationToken::new(), - ) + manager + .prewarm(prepared.setup(), stale_key) .await .expect("spawn reviewer after instruction update"); + let review = manager.trunk().await.expect("prewarmed reviewer"); assert_eq!(review.reuse_key, expected_key); assert_eq!(review.session.inherited_instructions().await.thread, latest); diff --git a/codex-rs/core/src/guardian/review_session_threads.rs b/codex-rs/core/src/guardian/review_session_threads.rs deleted file mode 100644 index 642912f219..0000000000 --- a/codex-rs/core/src/guardian/review_session_threads.rs +++ /dev/null @@ -1,96 +0,0 @@ -//! Adapts the existing reviewer lifecycle to the extension's ThreadManager. -//! Captured parent inputs support inline delegates; prompts and review outcomes stay shared. - -use std::sync::Arc; -use std::sync::Weak; - -use codex_async_utils::OrCancelExt; -use codex_history::InitialHistory; -use codex_protocol::protocol::AskForApproval; -use codex_protocol::protocol::InternalSessionSource; -use codex_protocol::protocol::SessionSource; -use codex_protocol::protocol::SubAgentSource; -use codex_protocol::protocol::ThreadSource; -use tokio_util::sync::CancellationToken; - -use super::GUARDIAN_REVIEWER_NAME; -use crate::StartThreadOptions; -use crate::ThreadManager; -use crate::codex_delegate::forward_session_io; -use crate::config::Constrained; -use crate::session::SessionIo; -use crate::session::emit_subagent_session_started; -use crate::session::session::Session; - -pub(super) struct ManagedReviewerThreads { - manager: Weak, -} - -impl ManagedReviewerThreads { - pub(super) fn new(manager: Weak) -> Self { - Self { manager } - } - - pub(super) async fn spawn( - &self, - parent: &Arc, - mut options: StartThreadOptions, - cancel: CancellationToken, - ) -> anyhow::Result<(Arc, SessionIo)> { - let manager = self - .manager - .upgrade() - .ok_or_else(|| anyhow::anyhow!("thread manager is no longer available"))?; - if options.config.permissions.approval_policy.value() != AskForApproval::Never { - anyhow::bail!("Codex delegates require approval policy `never`"); - } - options.config.permissions.approval_policy = Constrained::allow_only(AskForApproval::Never); - options.session_source = Some(SessionSource::Internal(InternalSessionSource::Guardian)); - options.thread_source = Some(ThreadSource::GuardianReview); - options - .thread_extension_init - .insert(codex_extension_api::SessionIsolation::Isolated); - options.initial_history = match options.initial_history { - InitialHistory::Forked(history) => InitialHistory::Forked(history), - InitialHistory::New | InitialHistory::Cleared => InitialHistory::New, - InitialHistory::Resumed(_) => { - anyhow::bail!("guardian review forks cannot resume an existing thread") - } - }; - let spawned = manager - .start_thread(options) - .or_cancel(&cancel) - .await - .map_err(codex_protocol::error::CodexErr::from)??; - let thread = spawned.thread; - let session = Arc::clone(&thread.session); - let io = forward_session_io( - Arc::new(SessionIo { - tx_sub: thread.io.tx_sub.clone(), - rx_event: thread.io.rx_event.clone(), - agent_status: thread.io.agent_status.clone(), - session_loop_termination: thread.io.session_loop_termination.clone(), - }), - cancel, - ); - let manager = Arc::downgrade(&manager); - drop(tokio::spawn(async move { - thread.wait_until_terminated().await; - if let Some(manager) = manager.upgrade() { - manager - .remove_thread_if_matches(&thread.session.thread_id(), &thread) - .await; - } - })); - emit_subagent_session_started( - &parent.services.analytics_events_client, - parent.app_server_client_metadata().await, - session.session_id(), - session.thread_id(), - Some(parent.thread_id()), - session.thread_config_snapshot().await, - SubAgentSource::Other(GUARDIAN_REVIEWER_NAME.to_owned()), - ); - Ok((session, io)) - } -} diff --git a/codex-rs/core/src/guardian/runtime.rs b/codex-rs/core/src/guardian/runtime.rs index 9abbb59d3d..75d20b96d3 100644 --- a/codex-rs/core/src/guardian/runtime.rs +++ b/codex-rs/core/src/guardian/runtime.rs @@ -2,6 +2,7 @@ use codex_protocol::protocol::ReviewDecision; use std::sync::Arc; +use tokio_util::sync::CancellationToken; use super::ApprovalRequestReasons; use super::GuardianApprovalRequest; @@ -74,6 +75,7 @@ impl ReviewAction { #[derive(Clone)] pub(super) struct ReviewRuntime { pub(super) session: Arc, + pub(super) history_reset: CancellationToken, pub(super) context: GuardianReviewContext, pub(super) review_id: String, pub(super) request: ReviewAction, diff --git a/codex-rs/core/src/guardian/test_host.rs b/codex-rs/core/src/guardian/test_host.rs index 87fcf0d983..dc0d547cf7 100644 --- a/codex-rs/core/src/guardian/test_host.rs +++ b/codex-rs/core/src/guardian/test_host.rs @@ -3,10 +3,16 @@ use std::sync::Arc; +use codex_extension_api::SessionIsolation; use codex_home::CodexHomeUserInstructionsProvider; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::InternalSessionSource; use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadSource; +use super::GuardianReviewSessionManager; use crate::config::Config; +use crate::config::Constrained; use crate::session::session::Session; pub(crate) fn install(session: &Session, config: &Config) { @@ -29,7 +35,44 @@ pub(crate) fn install(session: &Session, config: &Config) { /*attestation_provider*/ None, /*external_time_provider*/ None, )); - let host = super::GuardianReviewSessionHost::with_thread_manager(Arc::downgrade(&manager)); - session.services.thread_extension_data.insert(host); - session.services.thread_extension_data.insert(manager); + session + .services + .thread_extension_data + .insert(GuardianReviewSessionManager::new( + move |context, key, kind, snapshot, cancel| { + let manager = Arc::clone(&manager); + Box::pin(async move { + let history_reset = context.history_reset.clone(); + let (mut options, state) = context.thread_options(snapshot).await; + if matches!( + kind, + codex_analytics::GuardianReviewSessionKind::EphemeralForked + ) { + options.config.ephemeral = true; + } + options.config.permissions.approval_policy = + Constrained::allow_only(AskForApproval::Never); + options.session_source = + Some(SessionSource::Internal(InternalSessionSource::Guardian)); + options.thread_source = Some(ThreadSource::GuardianReview); + options + .thread_extension_init + .insert(SessionIsolation::Isolated); + let session_cancel = cancel.clone(); + let until = async move { + let _cancel_on_exit = cancel.clone().drop_guard(); + tokio::select! { + _ = cancel.cancelled() => {} + _ = history_reset.cancelled() => {} + } + }; + let spawned = manager + .start_thread_until(options, until, &tokio_util::task::TaskTracker::new()) + .await?; + Ok(context + .bind_thread(&spawned.thread, key, state, session_cancel) + .await) + }) + }, + )); } diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index c5f8a5bb53..777b77c54a 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -12,7 +12,6 @@ use crate::guardian::prompt::BUNDLED_GUARDIAN_POLICY; use crate::guardian::prompt::BUNDLED_GUARDIAN_POLICY_TEMPLATE; use crate::guardian::prompt::guardian_policy_prompt_with_config_and_template; use crate::guardian::review::guardian_review_session_config; -use crate::guardian::review::routes_approval_to_guardian_with_reviewer; use crate::session::session::Session; use crate::session::tests::update_turn_settings_for_test; use crate::session::turn_context::TurnContext; @@ -2592,6 +2591,7 @@ async fn guardian_reuses_prompt_cache_key_and_appends_prior_reviews() -> anyhow: .await; let committed_rollout_items = session .guardian_review_session() + .expect("Guardian pool installed") .trunk() .await .expect("reviewer") @@ -2927,6 +2927,7 @@ async fn guardian_reused_trunk_ignores_stale_prior_turn_completion() -> anyhow:: session .guardian_review_session() + .expect("Guardian pool installed") .trunk() .await .expect("reviewer") @@ -4255,3 +4256,18 @@ async fn review_approval_request( .await .expect("Guardian should handle the request") } + +/// Whether this turn should route allowed approval prompts through the guardian +/// reviewer instead of surfacing them to the user. ARC may still block actions +/// earlier in the flow. +fn routes_approval_to_guardian(turn: &crate::session::turn_context::TurnContext) -> bool { + routes_approval_to_guardian_with_reviewer(turn, turn.config.approvals_reviewer) +} + +/// Whether an approval with its own reviewer selection should be routed through guardian. +fn routes_approval_to_guardian_with_reviewer( + turn: &crate::session::turn_context::TurnContext, + approvals_reviewer: ApprovalsReviewer, +) -> bool { + routes_approval_policy_to_guardian(turn.approval_policy(), approvals_reviewer) +} diff --git a/codex-rs/core/src/guardian_review.rs b/codex-rs/core/src/guardian_review.rs index 764ab656f7..82edf2323d 100644 --- a/codex-rs/core/src/guardian_review.rs +++ b/codex-rs/core/src/guardian_review.rs @@ -1,4 +1,7 @@ -//! Opaque host capabilities for the synchronous Guardian extension. -//! Core supplies context and runtime operations; the extension owns review policy and pooling. +//! Context adapter types used by the host-installed synchronous Guardian extension. +//! The extension supplies agent startup and owns the reviewer pool and its lifecycle. -pub use crate::guardian::GuardianReviewSessionHost; +pub use crate::guardian::GuardianReviewSession; +pub use crate::guardian::GuardianReviewState; +pub use crate::guardian::PreparedGuardianContext; +pub use crate::guardian::prepare_review_prewarm; diff --git a/codex-rs/core/src/session/handlers.rs b/codex-rs/core/src/session/handlers.rs index 8c43908eee..dc739a59bf 100644 --- a/codex-rs/core/src/session/handlers.rs +++ b/codex-rs/core/src/session/handlers.rs @@ -308,12 +308,12 @@ pub(super) async fn shutdown_session_runtime(sess: &Arc) { sess.mcp_refresh.close(); sess.services.mcp_runtime.shutdown().await; } - sess.guardian_review_session().shutdown().await; crate::hook_runtime::run_session_end_hooks(sess).await; + emit_thread_stop_lifecycle(sess).await; } -pub(super) async fn emit_thread_stop_lifecycle(sess: &Session) { +async fn emit_thread_stop_lifecycle(sess: &Session) { for contributor in sess.services.extensions.thread_lifecycle_contributors() { contributor .on_thread_stop(codex_extension_api::ThreadStopInput { @@ -338,8 +338,6 @@ pub async fn shutdown(sess: &Arc, sub_id: String) -> bool { &[], ); - emit_thread_stop_lifecycle(sess.as_ref()).await; - // Gracefully flush and shutdown thread persistence on session end so tests // that inspect durable state do not race with the background writer. if let Some(live_thread) = sess.live_thread() @@ -607,7 +605,6 @@ pub(super) async fn submission_loop( // explicit shutdown op, still run session teardown. if !shutdown_received { shutdown_session_runtime(&sess).await; - emit_thread_stop_lifecycle(sess.as_ref()).await; if let Some(live_thread) = sess.live_thread() && let Err(err) = live_thread.shutdown().await { diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index f4343f54bc..df065f8feb 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2245,13 +2245,11 @@ impl Session { )); } - /// Uses the extension-owned reviewer, or the same manager for standalone hosts. - pub(crate) fn guardian_review_session(&self) -> Arc { - self.services.thread_extension_data.get_or_init(|| { - GuardianReviewSessionManager::new(|context, key, kind, snapshot, cancel| { - Box::pin(async move { context.spawn(key, kind, snapshot, cancel).await }) - }) - }) + /// Returns the reviewer pool installed by the Guardian extension. + pub(crate) fn guardian_review_session(&self) -> Option> { + self.services + .thread_extension_data + .get::() } pub(crate) async fn emit_turn_started(&self, turn_context: &TurnContext) { @@ -4393,6 +4391,15 @@ impl Session { state.clone_history() } + /// Captures the history lifetime without exposing mutable session state. + pub(crate) async fn history_reset(&self) -> (u64, CancellationToken) { + let state = self.state.lock().await; + ( + state.history.reset_version, + state.history_reset.child_token(), + ) + } + pub(crate) async fn conversation_history_snapshot( &self, ) -> Arc { diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index fdf3487abc..9f86627a45 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -1021,12 +1021,12 @@ impl Session { self.features.enabled(Feature::Personality), ) .await; - self.services - .thread_extension_data - .insert(model_info.clone()); - let multi_agent_version = match multi_agent_runtime { TurnMultiAgentRuntime::ResolveAndStore => { + // A background preview must not overwrite a newer turn's model metadata. + self.services + .thread_extension_data + .insert(model_info.clone()); self.resolve_multi_agent_version_for_model(&model_info, &per_turn_config) } TurnMultiAgentRuntime::Preview => per_turn_config.multi_agent_version_for_model( diff --git a/codex-rs/core/src/session/turn_suspension.rs b/codex-rs/core/src/session/turn_suspension.rs index 64a573b813..17f756c5fc 100644 --- a/codex-rs/core/src/session/turn_suspension.rs +++ b/codex-rs/core/src/session/turn_suspension.rs @@ -107,9 +107,8 @@ pub(super) async fn suspend_turn_and_shutdown( live_thread.shutdown().await.map_err(|error| { CodexErr::Fatal(format!("close suspended root turn writer failed: {error}")) })?; - // Announce thread shutdown only after its writer closes so a replacement worker - // cannot write the same thread concurrently. - handlers::emit_thread_stop_lifecycle(session.as_ref()).await; + // Announce completion only after extension cleanup and writer closure so a + // replacement worker cannot write the same thread concurrently. session .deliver_event_raw(Event { id: submission_id, diff --git a/codex-rs/core/src/session_startup_prewarm.rs b/codex-rs/core/src/session_startup_prewarm.rs index 9074162ea4..43b66cb97c 100644 --- a/codex-rs/core/src/session_startup_prewarm.rs +++ b/codex-rs/core/src/session_startup_prewarm.rs @@ -12,7 +12,6 @@ use tracing::trace_span; use tracing::warn; use crate::client::ModelClientSession; -use crate::guardian::routes_approval_to_guardian; use crate::responses_metadata::CodexResponsesRequestKind; use crate::session::INITIAL_SUBMIT_ID; use crate::session::RequestEffortUsage; @@ -273,20 +272,6 @@ async fn schedule_startup_prewarm_inner( prewarm_started_at.elapsed(), /*status*/ None, ); - if routes_approval_to_guardian(&startup_turn_context) { - let guardian_session = Arc::clone(&session); - let guardian_parent_turn = Arc::clone(&startup_turn_context); - drop(tokio::spawn(async move { - if let Err(err) = crate::guardian::prewarm_guardian_review_session( - guardian_session, - guardian_parent_turn, - ) - .await - { - warn!("failed to initialize guardian review session: {err:#}"); - } - })); - } let startup_cancellation_token = CancellationToken::new(); let built_tools_started_at = Instant::now(); // Startup prewarm runs before run_turn and needs its own tool-building snapshot. diff --git a/codex-rs/core/src/state/session.rs b/codex-rs/core/src/state/session.rs index b45bec136a..b3ea977da0 100644 --- a/codex-rs/core/src/state/session.rs +++ b/codex-rs/core/src/state/session.rs @@ -28,6 +28,7 @@ use codex_protocol::protocol::TokenUsageInfo; use codex_protocol::protocol::TokenUsageRecord; use codex_protocol::protocol::TurnContextItem; use codex_utils_output_truncation::TruncationPolicy; +use tokio_util::sync::CancellationToken; use tokio_util::task::AbortOnDropHandle; /// Runtime request effort, initially unset and established by prewarm or sampling. @@ -72,6 +73,8 @@ pub(crate) struct SessionState { /// Persisted origin of the session base instructions, when known. pub(crate) base_instructions_provenance: Option, pub(crate) history: ContextManager, + /// Cancels work bound to discarded history. Appends and compaction preserve it. + pub(crate) history_reset: CancellationToken, pub(crate) latest_rate_limits: Option, pub(crate) latest_token_usage_record: Option, pub(crate) server_reasoning_included: bool, @@ -117,6 +120,7 @@ impl SessionState { session_configuration, base_instructions_provenance: None, history, + history_reset: CancellationToken::new(), latest_rate_limits: None, latest_token_usage_record: None, server_reasoning_included: false, @@ -174,10 +178,11 @@ impl SessionState { items: Vec, reference_context_item: Option, ) { - self.history.replace(items); - self.history - .set_reference_context_item(reference_context_item); - self.auto_compact_window.clear_prefill(); + self.replace_annotated_history( + items.into_iter().map(ResponseItemEnvelope::new).collect(), + reference_context_item, + HistoryReplacement::Reset, + ); } pub(crate) fn replace_annotated_history( @@ -188,7 +193,10 @@ impl SessionState { ) { match replacement { HistoryReplacement::Compaction => self.history.replace_compacted(items), - HistoryReplacement::Reset => self.history.replace_annotated(items), + HistoryReplacement::Reset => { + self.history.replace_annotated(items); + std::mem::take(&mut self.history_reset).cancel(); + } } self.history .set_reference_context_item(reference_context_item); diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index 0040bfe9e5..166df45251 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -1052,6 +1052,19 @@ async fn guardian_session_prewarms_and_is_reused_for_first_review( .await .expect("guardian trunk rollout path"); test.codex.shutdown_and_wait().await?; + // Parent stop joins ThreadManager cleanup through the real extension registration. + assert!( + test.codex + .thread_extension_data() + .get::() + .expect("Guardian tasks") + .tasks + .is_empty() + ); + assert!(matches!( + test.thread_store.flush_thread(guardian_thread_id).await, + Err(codex_thread_store::ThreadStoreError::ThreadNotFound { .. }) + )); let guardian_rollout = fs::read_to_string(guardian_rollout_path)? .lines() .map(codex_rollout::parse_rollout_line) diff --git a/codex-rs/core/tests/suite/network_approval.rs b/codex-rs/core/tests/suite/network_approval.rs index 283a5724d2..dce430e5b7 100644 --- a/codex-rs/core/tests/suite/network_approval.rs +++ b/codex-rs/core/tests/suite/network_approval.rs @@ -15,6 +15,7 @@ use codex_exec_server::LOCAL_ENVIRONMENT_ID; use codex_exec_server::REMOTE_ENVIRONMENT_ID; use codex_exec_server::RemoveOptions; use codex_features::Feature; +use codex_history::RolloutItem; use codex_network_proxy::NetworkProxyConfig; use codex_protocol::approvals::NetworkApprovalContext; use codex_protocol::approvals::NetworkApprovalProtocol; @@ -45,6 +46,7 @@ use codex_protocol::request_permissions::PermissionGrantScope; use codex_protocol::request_permissions::RequestPermissionProfile; use codex_protocol::request_permissions::RequestPermissionsResponse; use codex_protocol::user_input::UserInput; +use codex_thread_store::LoadThreadHistoryParams; use codex_utils_path_uri::PathUri; use core_test_support::PathBufExt; use core_test_support::PathExt; @@ -1582,12 +1584,16 @@ async fn failed_network_policy_amendment_denies_request_and_does_not_approve_hos Ok(()) } +#[test_case(ApprovalsReviewer::User; "user_approval")] +#[test_case(ApprovalsReviewer::AutoReview; "shutdown_during_guardian_retry")] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[cfg_attr( not(target_os = "linux"), ignore = "requires the trusted Linux proxy bridge" )] -async fn unattributed_network_request_uses_active_turn_environment_fallback() -> Result<()> { +async fn unattributed_network_request_uses_active_turn_environment_fallback( + approvals_reviewer: ApprovalsReviewer, +) -> Result<()> { skip_if_target_windows!(Ok(()), "uses a raw TCP proxy fixture"); skip_if_host_windows!(Ok(())); skip_if_no_network!(Ok(())); @@ -1610,7 +1616,7 @@ async fn unattributed_network_request_uses_active_turn_environment_fallback() -> &test, "hold the active turn", vec![local(test.config.cwd.clone())], - ApprovalsReviewer::User, + approvals_reviewer, AskForApproval::OnRequest, ) .await?; @@ -1623,6 +1629,64 @@ async fn unattributed_network_request_uses_active_turn_environment_fallback() -> .context("expected managed network proxy")? .http_addr .clone(); + if approvals_reviewer == ApprovalsReviewer::AutoReview { + mount_sse_once_match( + &server, + is_guardian_request, + core_test_support::responses::sse_failed( + "guardian-first-failure", + "rate_limit_exceeded", + "Rate limit exceeded. Please try again in 0s.", + ), + ) + .await; + let retrying = mount_sse_once_match( + &server, + is_guardian_request, + core_test_support::responses::sse_failed( + "guardian-rate-limited", + "rate_limit_exceeded", + "Rate limit exceeded. Please try again in 60s.", + ), + ) + .await; + let proxy_request = tokio::spawn(raw_http_proxy_request(proxy_addr, NETWORK_TEST_HOST)); + wait_for_response_request(&retrying).await; + let request = retrying.single_request().body_json(); + let thread_id = request["client_metadata"]["thread_id"] + .as_str() + .context("Guardian thread id")?; + let thread_id = codex_protocol::ThreadId::from_string(thread_id)?; + // Internal reviewers are hidden from get_thread. Their persisted terminal event + // confirms the reviewer's own retries finished before testing Guardian's retry wait. + tokio::time::timeout(Duration::from_secs(5), async { + loop { + test.thread_store.flush_thread(thread_id).await?; + let history = test + .thread_store + .load_latest_model_context(LoadThreadHistoryParams { + thread_id, + include_archived: false, + }) + .await?; + if history.items.iter().any(|item| { + matches!(item, RolloutItem::EventMsg(EventMsg::TurnComplete(event)) if event.error.is_some()) + }) { + return Ok::<(), anyhow::Error>(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .context("Guardian reviewer did not finish its rate-limited turn")??; + tokio::time::timeout(Duration::from_secs(5), test.codex.shutdown_and_wait()) + .await + .context("parent shutdown waited for Guardian's 60-second retry")??; + let response = tokio::time::timeout(Duration::from_secs(5), proxy_request).await???; + assert!(response.starts_with("HTTP/1.1 403"), "{response}"); + assert_eq!(retrying.requests().len(), 1); + return Ok(()); + } let proxy_request = tokio::spawn(raw_http_proxy_request(proxy_addr, NETWORK_TEST_HOST)); let approval = expect_network_approval(&test, LOCAL_ENVIRONMENT_ID).await?; assert_eq!(approval.command, ["network-access", NETWORK_TEST_TARGET]); diff --git a/codex-rs/ext/extension-api/src/contributors.rs b/codex-rs/ext/extension-api/src/contributors.rs index 89480df2f1..61e7da448a 100644 --- a/codex-rs/ext/extension-api/src/contributors.rs +++ b/codex-rs/ext/extension-api/src/contributors.rs @@ -170,7 +170,8 @@ pub trait ThreadLifecycleContributor: Send + Sync { }) } - /// Called before the host drops the thread runtime and thread-scoped store. + /// Called during runtime teardown, before the host closes persistent history. + /// Contributors must cancel and join their background work before returning. fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> { Box::pin(async move { let _self = self; diff --git a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs index 38b303a6c9..b614d536f6 100644 --- a/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs +++ b/codex-rs/ext/extension-api/src/contributors/thread_lifecycle.rs @@ -75,7 +75,7 @@ pub struct ThreadIdleInput<'a> { pub thread_store: &'a ExtensionData, } -/// Input supplied when the host stops a thread runtime. +/// Input supplied during runtime teardown, before persistent history closes. pub struct ThreadStopInput<'a> { /// Store scoped to the host session runtime. pub session_store: &'a ExtensionData, diff --git a/codex-rs/ext/guardian-reviewer/Cargo.toml b/codex-rs/ext/guardian-reviewer/Cargo.toml index 5723c09d2e..982e35ceed 100644 --- a/codex-rs/ext/guardian-reviewer/Cargo.toml +++ b/codex-rs/ext/guardian-reviewer/Cargo.toml @@ -26,7 +26,7 @@ rand = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } tokio = { workspace = true } -tokio-util = { workspace = true } +tokio-util = { workspace = true, features = ["rt"] } tracing = { workspace = true } [dev-dependencies] diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 68d24d0896..2ac07168c8 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -44,6 +44,7 @@ pub(crate) use deadline::run_before_review_deadline_with_cancel; pub use pool::ReviewerPool; pub use pool::ReviewerRequest; pub use pool::ReviewerSession; +pub use pool::ReviewerTasks; pub use pool::SessionDisposition; pub use review::ReviewHost; diff --git a/codex-rs/ext/guardian-reviewer/src/pool.rs b/codex-rs/ext/guardian-reviewer/src/pool.rs index 49b973f2a0..8447bb97e3 100644 --- a/codex-rs/ext/guardian-reviewer/src/pool.rs +++ b/codex-rs/ext/guardian-reviewer/src/pool.rs @@ -18,6 +18,14 @@ use crate::GuardianReviewSessionOutcome; use crate::run_before_review_deadline; use crate::run_before_review_deadline_with_cancel; +/// Background work owned by Guardian for one parent runtime. +/// Stop cancels the work and joins this tracker before the parent closes its history. +#[derive(Default)] +pub struct ReviewerTasks { + pub tasks: tokio_util::task::TaskTracker, + pub cancellation: CancellationToken, +} + /// A host-owned reviewer session. Context and snapshots remain opaque to the pool. /// Shutdown must cancel the runtime and await its termination. pub trait ReviewerSession: Send + Sync + 'static { diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index b2e57173b9..1cdb168392 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -13,12 +13,14 @@ doctest = false workspace = true [dependencies] +anyhow = { workspace = true } codex-analytics = { workspace = true } codex-api = { workspace = true } codex-core = { workspace = true } codex-extension-api = { workspace = true } codex-features = { workspace = true } codex-guardian-context = { workspace = true } +codex-guardian-reviewer = { workspace = true } codex-history = { workspace = true } codex-http-client = { workspace = true } codex-login = { workspace = true } @@ -33,7 +35,6 @@ tracing = { workspace = true } uuid = { workspace = true } [dev-dependencies] -anyhow = { workspace = true } codex-model-provider-info = { workspace = true } core_test_support = { workspace = true } pretty_assertions = { workspace = true } 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 50540c617a..31aac3bed6 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 @@ -2203,6 +2203,8 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> { .await; model_info.slug = "protected-model".to_owned(); thread_store.insert(model_info); + // A late prewarm preview must leave the active model's review requirements intact. + let _ = codex_core::guardian_review::prepare_review_prewarm(&test.codex).await?; let authorization = ScoreAuthorization::current(&test.codex).await; let progress = thread_store .get::() diff --git a/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs b/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs index bf858e562a..84cec63902 100644 --- a/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs @@ -1,18 +1,31 @@ -//! Installs the synchronous reviewer independently of async scorer startup. -//! Reviewer policy and pool state live in the extension; core supplies context and runtime handles. +//! Runs Guardian agents through ThreadManager and owns their background tasks. +//! Context stays in the temporary core adapter. Parent stop joins all reviewer +//! work before the parent closes its history, including partial startup. use std::sync::Arc; use std::sync::Weak; use codex_core::ThreadManager; use codex_core::config::Config; -use codex_core::guardian_review::GuardianReviewSessionHost; +use codex_core::config::Constrained; +use codex_core::guardian_review::GuardianReviewSession; use codex_extension_api::ExtensionFuture; use codex_extension_api::ExtensionRegistryBuilder; +use codex_extension_api::SessionIsolation; use codex_extension_api::ThreadLifecycleContributor; +use codex_extension_api::ThreadReadyInput; use codex_extension_api::ThreadStartInput; +use codex_extension_api::ThreadStopInput; +use codex_guardian_reviewer::ReviewerPool; +use codex_guardian_reviewer::ReviewerTasks; +use codex_protocol::ThreadId; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::InternalSessionSource; +use codex_protocol::protocol::SessionSource; +use codex_protocol::protocol::ThreadSource; -/// Owns reviewer state through the same thread manager as the parent conversation. +/// Owns reviewer agents through the same thread manager as the parent conversation. #[derive(Clone, Debug)] pub struct GuardianExtension { thread_manager: Weak, @@ -33,11 +46,137 @@ impl ThreadLifecycleContributor for GuardianExtension { if input.session_source.is_internal() { return; } + let manager = self.thread_manager.clone(); + let runtime = input.thread_store.get_or_init(ReviewerTasks::default); input.thread_store.get_or_init(|| { - GuardianReviewSessionHost::with_thread_manager(self.thread_manager.clone()) + ReviewerPool::::new( + move |context, key, kind, snapshot, cancel| { + let manager = manager.clone(); + let runtime = Arc::clone(&runtime); + Box::pin(async move { + let history_reset = context.history_reset.clone(); + // Register before checking cancellation so parent stop also joins + // a spawn racing with shutdown. + let _task = runtime.tasks.token(); + anyhow::ensure!( + !runtime.cancellation.is_cancelled() + && !cancel.is_cancelled() + && !history_reset.is_cancelled(), + "Guardian is stopping" + ); + let manager = manager.upgrade().ok_or_else(|| { + anyhow::anyhow!("thread manager is no longer available") + })?; + let (mut options, state) = context.thread_options(snapshot).await; + if matches!( + kind, + codex_analytics::GuardianReviewSessionKind::EphemeralForked + ) { + options.config.ephemeral = true; + } + options.config.permissions.approval_policy = + Constrained::allow_only(AskForApproval::Never); + options.session_source = + Some(SessionSource::Internal(InternalSessionSource::Guardian)); + options.thread_source = Some(ThreadSource::GuardianReview); + options + .thread_extension_init + .insert(SessionIsolation::Isolated); + let parent_cancel = runtime.cancellation.clone(); + let session_cancel = cancel.clone(); + let until = async move { + let _cancel_on_exit = cancel.clone().drop_guard(); + tokio::select! { + _ = cancel.cancelled() => {} + _ = history_reset.cancelled() => {} + _ = parent_cancel.cancelled() => {} + } + }; + let spawned = manager + .start_thread_until(options, until, &runtime.tasks) + .await?; + Ok(context + .bind_thread(&spawned.thread, key, state, session_cancel) + .await) + }) + }, + ) }); }) } + + fn on_thread_ready<'a>( + &'a self, + input: ThreadReadyInput<'a, Config>, + ) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if input.session_source.is_internal() + || !matches!( + input.config.permissions.approval_policy.value(), + AskForApproval::OnRequest | AskForApproval::Granular(_) + ) + || input.config.approvals_reviewer != ApprovalsReviewer::AutoReview + { + return; + } + let Some(runtime) = input.thread_store.get::() else { + return; + }; + let task = runtime.tasks.token(); + if runtime.cancellation.is_cancelled() { + return; + } + let Some(manager) = self.thread_manager.upgrade() else { + return; + }; + let Ok(thread_id) = ThreadId::from_string(input.thread_store.level_id()) else { + return; + }; + let Ok(parent) = manager.get_thread(thread_id).await else { + return; + }; + let Some(pool) = input + .thread_store + .get::>() + else { + return; + }; + let cancel = runtime.cancellation.clone(); + tokio::spawn(async move { + let _task = task; + let prepare = async { + let context = + codex_core::guardian_review::prepare_review_prewarm(&parent).await?; + let key = context.reuse_key(/*previous*/ None); + pool.prewarm(Arc::new(context), key).await + }; + tokio::select! { + _ = cancel.cancelled() => {} + result = prepare => { + if let Err(error) = result { + tracing::warn!("failed to prewarm Guardian reviewer: {error:#}"); + } + } + } + }); + }) + } + + fn on_thread_stop<'a>(&'a self, input: ThreadStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + if let Some(runtime) = input.thread_store.get::() { + runtime.cancellation.cancel(); + if let Some(pool) = input + .thread_store + .get::>() + { + pool.shutdown().await; + } + runtime.tasks.close(); + runtime.tasks.wait().await; + } + }) + } } pub fn install(