diff --git a/codex-rs/core/src/guardian/review_session_setup.rs b/codex-rs/core/src/guardian/review_session_setup.rs index 08e785ad45..d27a412304 100644 --- a/codex-rs/core/src/guardian/review_session_setup.rs +++ b/codex-rs/core/src/guardian/review_session_setup.rs @@ -14,7 +14,7 @@ pub struct PreparedGuardianContext { context_policy: ReviewContextPolicy, key: GuardianReviewSessionReuseKey, parent_compaction: Option, - prefer_parent_checkpoint: bool, + restart_from_parent_checkpoint: bool, pub history_reset: CancellationToken, } @@ -54,7 +54,7 @@ impl PreparedGuardianContext { context_policy, key, parent_compaction, - prefer_parent_checkpoint: false, + restart_from_parent_checkpoint: false, history_reset, }) } @@ -83,7 +83,7 @@ impl PreparedGuardianContext { &self, snapshot: Option, ) -> (crate::StartThreadOptions, GuardianReviewState) { - let snapshot = snapshot.filter(|_| !self.prefer_parent_checkpoint); + let snapshot = snapshot.filter(|_| !self.restart_from_parent_checkpoint); let (mut conversation, mut history) = snapshot.map(ConversationState::fork).unzip(); if self.parent_compaction.is_some() && conversation @@ -212,6 +212,9 @@ impl ReviewerRequest for PreparedReview { fn context(&self, previous: Option<&GuardianReviewSession>) -> GuardianReviewSessionReuseKey { self.context.reuse_key(previous) } + fn requires_fresh_session(&self) -> bool { + self.context.restart_from_parent_checkpoint + } fn deadline(&self) -> tokio::time::Instant { self.params.deadline } @@ -256,7 +259,7 @@ impl ReviewerRequest for PreparedReview { if recovery_requested { result.disposition = SessionDisposition::Discard; } - if self.context.prefer_parent_checkpoint || !recovery_requested { + if self.context.restart_from_parent_checkpoint || !recovery_requested { record_failed_review(&session.session, &self.params, &result.outcome).await; } result @@ -283,7 +286,7 @@ pub(crate) async fn run_guardian_review_session( { // One restart only, under the original deadline. An oversized parent // checkpoint must fail rather than repeatedly compact and recreate. - Arc::make_mut(&mut prepared.context).prefer_parent_checkpoint = true; + Arc::make_mut(&mut prepared.context).restart_from_parent_checkpoint = true; pool.review(prepared).await } else { result diff --git a/codex-rs/ext/guardian-reviewer/src/pool.rs b/codex-rs/ext/guardian-reviewer/src/pool.rs index 750514ba7e..77f4ad4f18 100644 --- a/codex-rs/ext/guardian-reviewer/src/pool.rs +++ b/codex-rs/ext/guardian-reviewer/src/pool.rs @@ -48,6 +48,10 @@ pub trait ReviewerRequest: Send + Sync { &self, previous: Option<&Self::Session>, ) -> ::Context; + /// Requires a newly spawned session without a cached snapshot for this attempt. + fn requires_fresh_session(&self) -> bool { + false + } fn deadline(&self) -> Instant; fn cancellation(&self) -> Option<&CancellationToken>; fn run( @@ -90,7 +94,7 @@ type SpawnReviewer = dyn Fn( struct Trunk { session: Arc, - review_lock: Semaphore, + review_lock: Arc, cancellation: CancellationToken, } @@ -150,7 +154,7 @@ impl ReviewerPool { if !cancellation.is_cancelled() && trunk.is_none() { *trunk = Some(Arc::new(Trunk { session: Arc::new(session), - review_lock: Semaphore::new(/*permits*/ 1), + review_lock: Arc::new(Semaphore::new(/*permits*/ 1)), cancellation: guard.disarm(), })); } @@ -177,8 +181,9 @@ impl ReviewerPool { where R: ReviewerRequest, { + let requires_fresh_session = request.requires_fresh_session(); let mut spawned_trunk = false; - let (trunk, context) = match run_before_review_deadline( + let (trunk, context, reserved_permit) = match run_before_review_deadline( request.deadline(), request.cancellation(), self.trunk.lock(), @@ -187,11 +192,19 @@ impl ReviewerPool { { Ok(mut state) => { let context = request.context(state.as_ref().map(|trunk| trunk.session.as_ref())); + // Claim cached reviewers before validating them: retirement cancels + // the session before releasing its review permit. + let mut reserved_permit = state + .as_ref() + .and_then(|trunk| Arc::clone(&trunk.review_lock).try_acquire_owned().ok()); if let Some(trunk) = state.as_ref() - && (trunk.cancellation.is_cancelled() || trunk.session.context() != &context) - && trunk.review_lock.try_acquire().is_ok() + && (requires_fresh_session + || trunk.cancellation.is_cancelled() + || trunk.session.context() != &context) + && reserved_permit.is_some() { state.take(); + drop(reserved_permit.take()); } if state.is_none() { let cancellation = self.runtime.cancellation.child_token(); @@ -220,14 +233,25 @@ impl ReviewerPool { return (outcome, GuardianReviewAnalyticsResult::without_session()); } }; + let review_lock = Arc::new(Semaphore::new(/*permits*/ 1)); + // Reserve the first review before other requests can see this session. + reserved_permit = match Arc::clone(&review_lock).try_acquire_owned() { + Ok(permit) => Some(permit), + Err(error) => { + return ( + GuardianReviewSessionOutcome::PromptBuildFailed(error.into()), + GuardianReviewAnalyticsResult::without_session(), + ); + } + }; *state = Some(Arc::new(Trunk { session, - review_lock: Semaphore::new(/*permits*/ 1), + review_lock, cancellation: lifetime.disarm(), })); spawned_trunk = true; } - (state.as_ref().cloned(), context) + (state.as_ref().cloned(), context, reserved_permit) } Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), }; @@ -239,12 +263,12 @@ impl ReviewerPool { GuardianReviewAnalyticsResult::without_session(), ); }; - if trunk.session.context() != &context { + if trunk.session.context() != &context || (requires_fresh_session && !spawned_trunk) { return Box::pin(self.review_ephemeral(&request, context, /*snapshot*/ None)).await; } - let guard = match trunk.review_lock.try_acquire() { - Ok(guard) => guard, - Err(_) => { + let guard = match reserved_permit { + Some(guard) => guard, + None => { return Box::pin(self.review_ephemeral( &request, context, @@ -329,3 +353,7 @@ impl ReviewerPool { (outcome, analytics) } } + +#[cfg(test)] +#[path = "pool_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian-reviewer/src/pool_tests.rs b/codex-rs/ext/guardian-reviewer/src/pool_tests.rs new file mode 100644 index 0000000000..6acc17978f --- /dev/null +++ b/codex-rs/ext/guardian-reviewer/src/pool_tests.rs @@ -0,0 +1,115 @@ +use super::*; +use pretty_assertions::assert_eq; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; + +struct Session(usize); + +impl ReviewerSession for Session { + type Setup = (); + type Context = (); + type Snapshot = usize; + + fn context(&self) -> &() { + &() + } + async fn snapshot(&self) -> Option { + Some(self.0) + } + async fn commit_snapshot(&self) {} +} + +struct Request { + fresh: bool, + selected: tokio::sync::mpsc::UnboundedSender, + release: Option>, +} + +impl ReviewerRequest for Request { + type Session = Session; + + fn setup(&self) -> Arc<()> { + Arc::new(()) + } + fn context(&self, _previous: Option<&Session>) {} + fn requires_fresh_session(&self) -> bool { + self.fresh + } + fn deadline(&self) -> Instant { + Instant::now() + std::time::Duration::from_secs(5) + } + fn cancellation(&self) -> Option<&CancellationToken> { + None + } + async fn run( + &self, + session: &Session, + _kind: GuardianReviewSessionKind, + ) -> ReviewSessionResult { + self.selected + .send(session.0) + .expect("observe selected session"); + if let Some(release) = &self.release { + let _permit = release.acquire().await.expect("release busy review"); + } + ReviewSessionResult { + outcome: GuardianReviewSessionOutcome::Completed(Err(anyhow::anyhow!("test verdict"))), + disposition: SessionDisposition::Reusable, + analytics: GuardianReviewAnalyticsResult::without_session(), + } + } +} + +#[test_case::test_case(false; "idle trunk is replaced")] +#[test_case::test_case(true; "busy trunk is preserved")] +#[tokio::test] +async fn fresh_attempt_bypasses_cached_session_and_snapshot(busy: bool) { + let spawns = Arc::new(AtomicUsize::default()); + let pool = Arc::new(ReviewerPool::new( + Arc::new(ReviewerTasks::default()), + move |_, _, _, snapshot, _| { + assert_eq!( + snapshot, None, + "fresh recovery must not inherit a cached snapshot" + ); + let id = spawns.fetch_add(/*val*/ 1, Ordering::SeqCst); + Box::pin(async move { Ok(Session(id)) }) + }, + )); + pool.prewarm(Arc::new(()), ()).await.expect("prewarm trunk"); + let (selected, mut selections) = tokio::sync::mpsc::unbounded_channel(); + let release = Arc::new(Semaphore::new(/*permits*/ 0)); + let active = if busy { + let pool = Arc::clone(&pool); + let request = Request { + fresh: false, + selected: selected.clone(), + release: Some(Arc::clone(&release)), + }; + let task = tokio::spawn(async move { pool.review(request).await }); + assert_eq!(selections.recv().await, Some(0)); + Some(task) + } else { + None + }; + + pool.review(Request { + fresh: true, + selected: selected.clone(), + release: None, + }) + .await; + assert_eq!(selections.recv().await, Some(1)); + if let Some(active) = active { + release.add_permits(/*n*/ 1); + active.await.expect("finish concurrent review"); + } + // Idle recovery becomes the reusable trunk; busy recovery leaves the active trunk intact. + pool.review(Request { + fresh: false, + selected, + release: None, + }) + .await; + assert_eq!(selections.recv().await, Some(if busy { 0 } else { 1 })); +}