From e84a594636b47ecb1d9fb8db781915381c4a0e4b Mon Sep 17 00:00:00 2001 From: jif Date: Mon, 14 Sep 2026 21:21:52 +0000 Subject: [PATCH] Move Guardian reviewer startup into the pool (#45521) ## What changed Replace `ReviewerSessionFactory` with a startup callback installed through `ReviewerPool::new`. Review requests supply shared setup data and reuse context, while the pool uses its callback to create both reusable and forked reviewers. Update Guardian session setup, prewarming, and existing test fixtures to use the new pool API. GitOrigin-RevId: fa171503afcee8bdbdf6822573bea8ded50ce117 --- codex-rs/core/src/guardian/review_session.rs | 1 + .../core/src/guardian/review_session_setup.rs | 51 +++++---- .../core/src/guardian/review_session_tests.rs | 55 ++++------ codex-rs/core/src/session/mod.rs | 8 +- codex-rs/ext/guardian-reviewer/src/lib.rs | 1 - codex-rs/ext/guardian-reviewer/src/pool.rs | 100 ++++++++++-------- 6 files changed, 112 insertions(+), 104 deletions(-) diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 781aeeaffb..114889e669 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -882,6 +882,7 @@ impl codex_guardian_reviewer::ReviewerRuntime for GuardianReviewSession { mod tests; impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession { + type Setup = 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 c0eb43d67a..394c184048 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -4,10 +4,9 @@ use super::*; use codex_guardian_reviewer::ReviewerPool; use codex_guardian_reviewer::ReviewerRequest; -use codex_guardian_reviewer::ReviewerSessionFactory; use codex_guardian_reviewer::SessionDisposition; -pub(super) struct PreparedGuardianContext { +pub(crate) struct PreparedGuardianContext { parent: Arc, context: GuardianReviewContext, config: Config, @@ -128,14 +127,8 @@ impl PreparedGuardianContext { } } -impl ReviewerSessionFactory for PreparedGuardianContext { - type Session = GuardianReviewSession; - - fn context(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey { - self.reuse_key(previous) - } - - async fn spawn( +impl PreparedGuardianContext { + pub(crate) async fn spawn( &self, context: GuardianReviewSessionReuseKey, kind: GuardianReviewSessionKind, @@ -169,15 +162,18 @@ impl ReviewerSessionFactory for PreparedGuardianContext { } pub(super) struct PreparedReview { - factory: PreparedGuardianContext, + context: Arc, params: GuardianReviewSessionParams, } impl ReviewerRequest for PreparedReview { - type Factory = PreparedGuardianContext; + type Session = GuardianReviewSession; - fn factory(&self) -> &PreparedGuardianContext { - &self.factory + fn setup(&self) -> Arc { + Arc::clone(&self.context) + } + fn context(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey { + self.context.reuse_key(previous) } fn deadline(&self) -> tokio::time::Instant { self.params.deadline @@ -228,7 +224,7 @@ pub(crate) async fn run_guardian_review_session( pub(super) async fn prepare_review( params: GuardianReviewSessionParams, ) -> anyhow::Result { - let factory = PreparedGuardianContext::prepare( + let context = PreparedGuardianContext::prepare( Arc::clone(¶ms.parent_session), params.parent_context.clone(), params.spawn_config.clone(), @@ -237,19 +233,35 @@ pub(super) async fn prepare_review( params.compaction_model_hash.as_deref(), ) .await?; - Ok(PreparedReview { factory, params }) + Ok(PreparedReview { + context: Arc::new(context), + params, + }) } pub(crate) fn prewarm_guardian_review_session( parent: Arc, turn: Arc, ) -> BoxFuture<'static, anyhow::Result<()>> { - // Keep the Session -> Guardian -> Session startup future on the heap. + 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 + }) +} + +pub(super) fn prepare_prewarm( + parent: Arc, + turn: Arc, +) -> BoxFuture<'static, anyhow::Result> { Box::pin(async move { let context = GuardianReviewContext::from(turn); let config = guardian_review_session_config(&parent, &context).await?; let history = parent.clone_history().await; - let factory = PreparedGuardianContext::prepare( + PreparedGuardianContext::prepare( Arc::clone(&parent), context, config.spawn_config, @@ -257,7 +269,6 @@ pub(crate) fn prewarm_guardian_review_session( &config.node_repl_policy, config.compaction_model_hash.as_deref(), ) - .await?; - parent.guardian_review_session().prewarm(&factory).await + .await }) } diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index cd3e9ab71f..e99129d364 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -5,7 +5,6 @@ use crate::agents_md_manager::AgentsMdManager; use crate::context_manager::ContextManager; use codex_guardian_reviewer::ReviewerRequest; use codex_guardian_reviewer::ReviewerSession; -use codex_guardian_reviewer::ReviewerSessionFactory; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; use codex_protocol::openai_models::AutoReviewMessages; @@ -73,8 +72,7 @@ async fn run_review_preserves_evidence_during_parent_compaction() { .node_repl_auto_review_required, ) .with_node_repl_policy(¶ms.node_repl_policy); - let manager = GuardianReviewSessionManager::default(); - prewarm_test_session(&manager, reviewer).await; + let manager = prewarm_test_session(¶ms, reviewer).await; // Capture the review context, then compact the parent before the reviewer builds its prompt. let prepared = setup::prepare_review(params).await.unwrap(); let checkpoint: ResponseItem = serde_json::from_value(serde_json::json!({ @@ -314,7 +312,7 @@ async fn spawned_guardian_reuse_key_matches_inherited_instructions() { }; let prepared = setup::prepare_review(params).await.expect("prepare review"); let review = prepared - .factory() + .setup() .spawn( stale_key, GuardianReviewSessionKind::TrunkNew, @@ -901,8 +899,7 @@ async fn run_review_removes_trunk_when_event_stream_is_broken() { ) .with_environments(params.parent_context.environments()) .with_node_repl_policy(¶ms.node_repl_policy); - let manager = Arc::new(GuardianReviewSessionManager::default()); - prewarm_test_session(&manager, review_session).await; + let manager = Arc::new(prewarm_test_session(¶ms, review_session).await); let manager_for_review = Arc::clone(&manager); let review = tokio::spawn(async move { run_guardian_review_session(manager_for_review, params).await }); @@ -1174,32 +1171,22 @@ async fn interrupt_and_drain_turn_ignores_prior_turn_completion() { } // Reuse the existing in-memory reviewer fixture through the production prewarm path. -async fn prewarm_test_session(pool: &GuardianReviewSessionManager, session: GuardianReviewSession) { - struct ReadySession { - context: GuardianReviewSessionReuseKey, - session: Mutex>, - } - impl codex_guardian_reviewer::ReviewerSessionFactory for ReadySession { - type Session = GuardianReviewSession; - fn context( - &self, - _previous: Option<&GuardianReviewSession>, - ) -> GuardianReviewSessionReuseKey { - self.context.clone() - } - async fn spawn( - &self, - _context: GuardianReviewSessionReuseKey, - _kind: GuardianReviewSessionKind, - _snapshot: Option, - _cancellation: CancellationToken, - ) -> anyhow::Result { - Ok(self.session.lock().await.take().expect("one fixture spawn")) - } - } - let factory = ReadySession { - context: session.reuse_key.clone(), - session: Mutex::new(Some(session)), - }; - pool.prewarm(&factory).await.unwrap(); +async fn prewarm_test_session( + params: &GuardianReviewSessionParams, + session: GuardianReviewSession, +) -> GuardianReviewSessionManager { + let key = session.reuse_key.clone(); + let session = Arc::new(Mutex::new(Some(session))); + let pool = GuardianReviewSessionManager::new(move |_, _, _, _, _| { + let session = Arc::clone(&session); + Box::pin(async move { Ok(session.lock().await.take().expect("one fixture spawn")) }) + }); + let context = setup::prepare_prewarm( + Arc::clone(¶ms.parent_session), + Arc::clone(params.parent_context.turn()), + ) + .await + .unwrap(); + pool.prewarm(Arc::new(context), key).await.unwrap(); + pool } diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 6cd5d41403..f4343f54bc 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -2247,9 +2247,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::default) + 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 }) + }) + }) } pub(crate) async fn emit_turn_started(&self, turn_context: &TurnContext) { diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 61748f92e8..68d24d0896 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -44,7 +44,6 @@ 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::ReviewerSessionFactory; 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 e3f1a1894d..49b973f2a0 100644 --- a/codex-rs/ext/guardian-reviewer/src/pool.rs +++ b/codex-rs/ext/guardian-reviewer/src/pool.rs @@ -1,5 +1,5 @@ //! Owns the reusable reviewer and temporary forks for one parent thread. -//! The host owns session execution and context construction. Selection stays serialized; +//! Guardian supplies agent startup; the host supplies captured context. Selection stays serialized; //! concurrent reviews fork committed context and shutdown joins every tracked session. //! Startup and fork futures stay boxed to bound the orchestration stack frames. @@ -8,6 +8,7 @@ use std::sync::Arc; use codex_analytics::GuardianReviewAnalyticsResult; use codex_analytics::GuardianReviewSessionKind; +use codex_extension_api::ExtensionFuture; use tokio::sync::Mutex; use tokio::sync::Semaphore; use tokio::time::Instant; @@ -20,6 +21,7 @@ use crate::run_before_review_deadline_with_cancel; /// 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 { + type Setup: Send + Sync + 'static; type Context: Clone + PartialEq + Send + Sync; type Snapshot: Send + Sync; @@ -30,36 +32,21 @@ pub trait ReviewerSession: Send + Sync + 'static { fn shutdown(&self) -> impl Future + Send; } -/// Builds a session from one captured parent context. The host must preserve that -/// context's authority and honor cancellation during spawning, including partial startup. -pub trait ReviewerSessionFactory: Send + Sync { +/// Executes one approval on a selected session. The host must drain the submitted +/// turn before returning Reusable, and must keep the issuing action and permissions bound. +pub trait ReviewerRequest: Send + Sync { type Session: ReviewerSession; + fn setup(&self) -> Arc<::Setup>; fn context( &self, previous: Option<&Self::Session>, ) -> ::Context; - - fn spawn( - &self, - context: ::Context, - kind: GuardianReviewSessionKind, - snapshot: Option<::Snapshot>, - cancellation: CancellationToken, - ) -> impl Future> + Send; -} - -/// Executes one approval on a selected session. The host must drain the submitted -/// turn before returning Reusable, and must keep the issuing action and permissions bound. -pub trait ReviewerRequest: Send + Sync { - type Factory: ReviewerSessionFactory; - - fn factory(&self) -> &Self::Factory; fn deadline(&self) -> Instant; fn cancellation(&self) -> Option<&CancellationToken>; fn run( &self, - session: &::Session, + session: &Self::Session, kind: GuardianReviewSessionKind, ) -> impl Future< Output = ( @@ -81,8 +68,19 @@ pub enum SessionDisposition { pub struct ReviewerPool { state: Arc>>, cancellation: CancellationToken, + spawn: Box>, } +type SpawnReviewer = dyn Fn( + Arc<::Setup>, + ::Context, + GuardianReviewSessionKind, + Option<::Snapshot>, + CancellationToken, + ) -> ExtensionFuture<'static, anyhow::Result> + + Send + + Sync; + struct PoolState { trunk: Option>>, ephemeral_reviews: Vec>, @@ -93,14 +91,28 @@ struct Trunk { review_lock: Semaphore, } -impl Default for ReviewerPool { - fn default() -> Self { +impl ReviewerPool { + /// Installs Guardian's startup function once. It must finish or clean up partial startup + /// even when the caller drops its future, and preserve the supplied cancellation token. + pub fn new( + spawn: impl Fn( + Arc, + S::Context, + GuardianReviewSessionKind, + Option, + CancellationToken, + ) -> ExtensionFuture<'static, anyhow::Result> + + Send + + Sync + + 'static, + ) -> Self { Self { state: Arc::new(Mutex::new(PoolState { trunk: None, ephemeral_reviews: Vec::new(), })), cancellation: CancellationToken::new(), + spawn: Box::new(spawn), } } } @@ -117,20 +129,17 @@ impl ReviewerPool { } /// Prepares the first reviewer without replacing a review that won the startup race. - pub async fn prewarm( - &self, - factory: &impl ReviewerSessionFactory, - ) -> anyhow::Result<()> { + pub async fn prewarm(&self, setup: Arc, context: S::Context) -> anyhow::Result<()> { let cancellation = self.cancellation.child_token(); let guard = cancellation.clone().drop_guard(); - let session = factory - .spawn( - factory.context(/*previous*/ None), - GuardianReviewSessionKind::TrunkNew, - /*snapshot*/ None, - cancellation.clone(), - ) - .await?; + let session = (self.spawn)( + setup, + context, + GuardianReviewSessionKind::TrunkNew, + /*snapshot*/ None, + cancellation.clone(), + ) + .await?; let mut state = self.state.lock().await; if !cancellation.is_cancelled() && state.trunk.is_none() { state.trunk = Some(Arc::new(Trunk { @@ -181,8 +190,7 @@ impl ReviewerPool { request: R, ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) where - R: ReviewerRequest, - R::Factory: ReviewerSessionFactory, + R: ReviewerRequest, { let mut spawned_trunk = false; let (trunk, context) = match run_before_review_deadline( @@ -193,9 +201,8 @@ impl ReviewerPool { .await { Ok(mut state) => { - let context = request - .factory() - .context(state.trunk.as_ref().map(|trunk| trunk.session.as_ref())); + let context = + request.context(state.trunk.as_ref().map(|trunk| trunk.session.as_ref())); if let Some(trunk) = state.trunk.as_ref() && trunk.session.context() != &context && trunk.review_lock.try_acquire().is_ok() @@ -209,12 +216,13 @@ impl ReviewerPool { request.deadline(), request.cancellation(), &cancellation, - Box::pin(request.factory().spawn( + (self.spawn)( + request.setup(), context.clone(), GuardianReviewSessionKind::TrunkNew, /*snapshot*/ None, cancellation.clone(), - )), + ), ) .await { @@ -294,20 +302,20 @@ impl ReviewerPool { snapshot: Option, ) -> (GuardianReviewSessionOutcome, GuardianReviewAnalyticsResult) where - R: ReviewerRequest, - R::Factory: ReviewerSessionFactory, + R: ReviewerRequest, { let cancellation = self.cancellation.child_token(); let session = match run_before_review_deadline_with_cancel( request.deadline(), request.cancellation(), &cancellation, - Box::pin(request.factory().spawn( + (self.spawn)( + request.setup(), context, GuardianReviewSessionKind::EphemeralForked, snapshot, cancellation.clone(), - )), + ), ) .await {