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 {