From a113f3e063dcb4b0a7d8a4e63cfb2fea83996116 Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 15 Sep 2026 10:27:12 +0000 Subject: [PATCH] Consolidate Guardian reviewer lifecycle ownership (#45672) ## What changed - Use cancellation guards to tie reusable reviewers and temporary forks to their lifetimes, including when a review future is dropped. Replace cancelled reusable reviewers before reuse. - Share `ReviewerTasks` between the pool and `ThreadManager`, and wait for tracked cleanup during pool shutdown. Remove the separate session shutdown protocol. - Move denial cleanup from core task handling into Guardian's turn start, stop, and abort hooks. Expose `install_reviewer` to register both thread and turn lifecycle contributors together. ## Testing Add an integration test that exercises stale denial cleanup at turn start, completion, and interruption, including starting another turn after interruption. GitOrigin-RevId: 85badb235309c1e1576547d1746f7ee9fc66d545 --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/guardian/review_session.rs | 9 - .../core/src/guardian/review_session_tests.rs | 14 +- codex-rs/core/src/guardian/test_host.rs | 5 +- .../core/src/session/tests/guardian_tests.rs | 45 ----- codex-rs/core/src/tasks/mod.rs | 16 -- codex-rs/core/tests/common/test_codex.rs | 5 +- .../ext/guardian-reviewer/src/deadline.rs | 17 -- .../guardian-reviewer/src/deadline_tests.rs | 69 ------- codex-rs/ext/guardian-reviewer/src/lib.rs | 1 - codex-rs/ext/guardian-reviewer/src/pool.rs | 174 +++++------------- codex-rs/ext/guardian-v2/Cargo.toml | 1 + codex-rs/ext/guardian-v2/src/lib.rs | 4 +- .../src/sync_reviewer/lifecycle_tests.rs | 125 +++++++++++++ .../ext/guardian-v2/src/sync_reviewer/mod.rs | 60 +++--- 15 files changed, 231 insertions(+), 315 deletions(-) delete mode 100644 codex-rs/ext/guardian-reviewer/src/deadline_tests.rs create mode 100644 codex-rs/ext/guardian-v2/src/sync_reviewer/lifecycle_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 26a070df9a..8f8cb83eee 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3609,6 +3609,7 @@ dependencies = [ "http 1.4.0", "pretty_assertions", "serde_json", + "tempfile", "thiserror 2.0.18", "tokio", "tracing", diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index 397f1eb635..6c0ef314ff 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -874,15 +874,6 @@ impl codex_guardian_reviewer::ReviewerSession for GuardianReviewSession { fn context(&self) -> &Self::Context { &self.reuse_key } - fn cancel(&self) { - self.cancel_token.cancel(); - } - - async fn shutdown(&self) { - self.cancel_token.cancel(); - let _ = self.io.shutdown_and_wait().await; - } - async fn snapshot(&self) -> Option { self.state.lock().await.conversation.snapshot().cloned() } diff --git a/codex-rs/core/src/guardian/review_session_tests.rs b/codex-rs/core/src/guardian/review_session_tests.rs index 6616a291b3..34ff8fd968 100644 --- a/codex-rs/core/src/guardian/review_session_tests.rs +++ b/codex-rs/core/src/guardian/review_session_tests.rs @@ -4,7 +4,6 @@ use super::*; use crate::agents_md_manager::AgentsMdManager; use crate::context_manager::ContextManager; use codex_guardian_reviewer::ReviewerRequest; -use codex_guardian_reviewer::ReviewerSession; use codex_history::CodexHarnessMetadata; use codex_history::ResponseItemEnvelope; use codex_protocol::openai_models::AutoReviewMessages; @@ -299,7 +298,7 @@ async fn spawned_guardian_reuse_key_matches_inherited_instructions() { assert_eq!(review.reuse_key, expected_key); assert_eq!(review.session.inherited_instructions().await.thread, latest); - review.shutdown().await; + manager.shutdown().await; } #[tokio::test] @@ -1152,10 +1151,13 @@ async fn prewarm_test_session( ) -> 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 pool = GuardianReviewSessionManager::new( + Arc::new(codex_guardian_reviewer::ReviewerTasks::default()), + 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()), diff --git a/codex-rs/core/src/guardian/test_host.rs b/codex-rs/core/src/guardian/test_host.rs index dc0d547cf7..9529fb499f 100644 --- a/codex-rs/core/src/guardian/test_host.rs +++ b/codex-rs/core/src/guardian/test_host.rs @@ -35,12 +35,15 @@ pub(crate) fn install(session: &Session, config: &Config) { /*attestation_provider*/ None, /*external_time_provider*/ None, )); + let runtime = Arc::new(codex_guardian_reviewer::ReviewerTasks::default()); session .services .thread_extension_data .insert(GuardianReviewSessionManager::new( + Arc::clone(&runtime), move |context, key, kind, snapshot, cancel| { let manager = Arc::clone(&manager); + let runtime = Arc::clone(&runtime); Box::pin(async move { let history_reset = context.history_reset.clone(); let (mut options, state) = context.thread_options(snapshot).await; @@ -67,7 +70,7 @@ pub(crate) fn install(session: &Session, config: &Config) { } }; let spawned = manager - .start_thread_until(options, until, &tokio_util::task::TaskTracker::new()) + .start_thread_until(options, until, &runtime.tasks) .await?; Ok(context .bind_thread(&spawned.thread, key, state, session_cancel) diff --git a/codex-rs/core/src/session/tests/guardian_tests.rs b/codex-rs/core/src/session/tests/guardian_tests.rs index 72ae069562..b1baa6609a 100644 --- a/codex-rs/core/src/session/tests/guardian_tests.rs +++ b/codex-rs/core/src/session/tests/guardian_tests.rs @@ -1330,48 +1330,3 @@ async fn guardian_subagent_does_not_inherit_parent_exec_policy_rules() { ); drop(io); } - -#[test_case(TerminalEventKind::TurnComplete; "completion")] -#[test_case(TerminalEventKind::TurnAborted; "interruption")] -#[tokio::test] -async fn terminal_turn_clears_extension_owned_denials(terminal: TerminalEventKind) { - let (session, turn, events) = make_session_and_context_with_rx().await; - let finish = Arc::new(tokio::sync::Notify::new()); - session - .spawn_task( - Arc::clone(&turn), - Vec::new(), - HeldStepTask { - kind: TaskKind::Regular, - finish: Arc::clone(&finish), - }, - ) - .await; - let denials = - codex_guardian_reviewer::ReviewDenials::for_thread(&session.services.thread_extension_data); - for _ in 0..2 { - assert_eq!( - denials.record_denial(&turn.sub_id, turn.model_info()).await, - None - ); - } - match terminal { - TerminalEventKind::TurnComplete => finish.notify_one(), - TerminalEventKind::TurnAborted => { - session.abort_all_tasks(TurnAbortReason::Interrupted).await - } - } - recv_terminal_event(&events, terminal).await; - // Delivery precedes accounting cleanup. Wait for the runtime to finish the turn. - timeout(Duration::from_secs(/*secs*/ 5), async { - while session.active_turn.lock().await.is_some() { - tokio::task::yield_now().await; - } - }) - .await - .expect("turn becomes idle"); - assert_eq!( - denials.record_denial(&turn.sub_id, turn.model_info()).await, - None - ); -} diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 3a03c7b24a..4c5cd6dfd5 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -306,12 +306,6 @@ impl Session { let cancellation_token = CancellationToken::new(); let done = Arc::new(Notify::new()); - codex_guardian_reviewer::ReviewDenials::clear_turn( - &self.services.thread_extension_data, - &turn_context.sub_id, - ) - .await; - let (pending_items, _) = self.input_queue.drain_mailbox_input_items().await; let turn_state = { let mut active = self.active_turn.lock().await; @@ -812,11 +806,6 @@ impl Session { }) }; self.send_event(turn_context.as_ref(), event).await; - codex_guardian_reviewer::ReviewDenials::clear_turn( - &self.services.thread_extension_data, - &turn_context.sub_id, - ) - .await; let cleared_active_turn = { let mut active = self.active_turn.lock().await; @@ -968,11 +957,6 @@ impl Session { duration_ms, }); self.send_event(task.turn_context.as_ref(), event).await; - codex_guardian_reviewer::ReviewDenials::clear_turn( - &self.services.thread_extension_data, - &task.turn_context.sub_id, - ) - .await; // Regular items were flushed before this terminal event was appended; buffering // thread writers may not flush it without another explicit barrier. if let Err(err) = self.flush_rollout().await { diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index e51b7adf52..43985f7f30 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -732,10 +732,7 @@ impl TestCodexBuilder { .or_else(|| codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").ok()); let thread_manager = Arc::new_cyclic(|manager| { let mut extensions = self.extensions.to_builder(); - let guardian = Arc::new(codex_guardian_v2::GuardianReviewerExtension::new( - manager.clone(), - )); - extensions.thread_lifecycle_contributor(guardian); + codex_guardian_v2::install_reviewer(&mut extensions, manager.clone()); let thread_manager = ThreadManager::new( &config, auth_manager.clone(), diff --git a/codex-rs/ext/guardian-reviewer/src/deadline.rs b/codex-rs/ext/guardian-reviewer/src/deadline.rs index b7b63b8012..c6f63f7180 100644 --- a/codex-rs/ext/guardian-reviewer/src/deadline.rs +++ b/codex-rs/ext/guardian-reviewer/src/deadline.rs @@ -21,20 +21,3 @@ pub async fn run_before_review_deadline( } => Err(GuardianReviewSessionOutcome::Aborted), } } - -pub(crate) async fn run_before_review_deadline_with_cancel( - deadline: tokio::time::Instant, - external_cancel: Option<&CancellationToken>, - cancel_token: &CancellationToken, - future: impl Future, -) -> Result { - let result = run_before_review_deadline(deadline, external_cancel, future).await; - if result.is_err() { - cancel_token.cancel(); - } - result -} - -#[cfg(test)] -#[path = "deadline_tests.rs"] -mod tests; diff --git a/codex-rs/ext/guardian-reviewer/src/deadline_tests.rs b/codex-rs/ext/guardian-reviewer/src/deadline_tests.rs deleted file mode 100644 index df3a5d1b22..0000000000 --- a/codex-rs/ext/guardian-reviewer/src/deadline_tests.rs +++ /dev/null @@ -1,69 +0,0 @@ -use std::time::Duration; - -use pretty_assertions::assert_eq; -use tokio_util::sync::CancellationToken; - -use super::run_before_review_deadline_with_cancel; -use crate::GuardianReviewSessionOutcome; - -#[tokio::test(flavor = "current_thread")] -async fn run_before_review_deadline_with_cancel_cancels_token_on_timeout() { - let cancel_token = CancellationToken::new(); - - let outcome = run_before_review_deadline_with_cancel( - tokio::time::Instant::now() + Duration::from_millis(10), - /*external_cancel*/ None, - &cancel_token, - async { - tokio::time::sleep(Duration::from_millis(50)).await; - }, - ) - .await; - - assert!(matches!( - outcome, - Err(GuardianReviewSessionOutcome::TimedOut) - )); - assert!(cancel_token.is_cancelled()); -} - -#[tokio::test(flavor = "current_thread")] -async fn run_before_review_deadline_with_cancel_cancels_token_on_abort() { - let external_cancel = CancellationToken::new(); - let external_canceller = external_cancel.clone(); - let cancel_token = CancellationToken::new(); - drop(tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(10)).await; - external_canceller.cancel(); - })); - - let outcome = run_before_review_deadline_with_cancel( - tokio::time::Instant::now() + Duration::from_secs(1), - Some(&external_cancel), - &cancel_token, - std::future::pending::<()>(), - ) - .await; - - assert!(matches!( - outcome, - Err(GuardianReviewSessionOutcome::Aborted) - )); - assert!(cancel_token.is_cancelled()); -} - -#[tokio::test(flavor = "current_thread")] -async fn run_before_review_deadline_with_cancel_preserves_token_on_success() { - let cancel_token = CancellationToken::new(); - - let outcome = run_before_review_deadline_with_cancel( - tokio::time::Instant::now() + Duration::from_secs(1), - /*external_cancel*/ None, - &cancel_token, - async { 42usize }, - ) - .await; - - assert_eq!(outcome.unwrap(), 42); - assert!(!cancel_token.is_cancelled()); -} diff --git a/codex-rs/ext/guardian-reviewer/src/lib.rs b/codex-rs/ext/guardian-reviewer/src/lib.rs index 2ac07168c8..0e9ad66cc6 100644 --- a/codex-rs/ext/guardian-reviewer/src/lib.rs +++ b/codex-rs/ext/guardian-reviewer/src/lib.rs @@ -40,7 +40,6 @@ pub const MAX_REVIEW_ATTEMPTS: i64 = 3; pub const REVIEW_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(90); pub use deadline::run_before_review_deadline; -pub(crate) use deadline::run_before_review_deadline_with_cancel; pub use pool::ReviewerPool; pub use pool::ReviewerRequest; pub use pool::ReviewerSession; diff --git a/codex-rs/ext/guardian-reviewer/src/pool.rs b/codex-rs/ext/guardian-reviewer/src/pool.rs index 8447bb97e3..fc9971d44e 100644 --- a/codex-rs/ext/guardian-reviewer/src/pool.rs +++ b/codex-rs/ext/guardian-reviewer/src/pool.rs @@ -1,6 +1,7 @@ //! Owns the reusable reviewer and temporary forks for one parent thread. //! Guardian supplies agent startup; the host supplies captured context. Selection stays serialized; -//! concurrent reviews fork committed context and shutdown joins every tracked session. +//! concurrent reviews fork committed context. Lifetime guards cancel agents; ThreadManager +//! performs cleanup and tracks its completion. The pool never runs a second shutdown protocol. //! Startup and fork futures stay boxed to bound the orchestration stack frames. use std::future::Future; @@ -16,7 +17,6 @@ use tokio_util::sync::CancellationToken; 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. @@ -26,18 +26,16 @@ pub struct ReviewerTasks { 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. +/// Context bookkeeping for a reviewer. Lifetime and cleanup belong to ThreadManager. +/// Context and snapshots remain opaque while the host context builder is being A/B tested. pub trait ReviewerSession: Send + Sync + 'static { type Setup: Send + Sync + 'static; type Context: Clone + PartialEq + Send + Sync; type Snapshot: Send + Sync; fn context(&self) -> &Self::Context; - fn cancel(&self); fn snapshot(&self) -> impl Future> + Send; fn commit_snapshot(&self) -> impl Future + Send; - fn shutdown(&self) -> impl Future + Send; } /// Executes one approval on a selected session. The host must drain the submitted @@ -74,8 +72,8 @@ pub enum SessionDisposition { /// Per-parent reviewer state. The same pool serves prewarm, review, invalidation and shutdown. pub struct ReviewerPool { - state: Arc>>, - cancellation: CancellationToken, + trunk: Mutex>>>, + runtime: Arc, spawn: Box>, } @@ -89,20 +87,23 @@ type SpawnReviewer = dyn Fn( + Send + Sync; -struct PoolState { - trunk: Option>>, - ephemeral_reviews: Vec>, -} - struct Trunk { session: Arc, review_lock: Semaphore, + cancellation: CancellationToken, +} + +impl Drop for Trunk { + fn drop(&mut self) { + self.cancellation.cancel(); + } } 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( + runtime: Arc, spawn: impl Fn( Arc, S::Context, @@ -115,11 +116,8 @@ impl ReviewerPool { + 'static, ) -> Self { Self { - state: Arc::new(Mutex::new(PoolState { - trunk: None, - ephemeral_reviews: Vec::new(), - })), - cancellation: CancellationToken::new(), + trunk: Mutex::new(None), + runtime, spawn: Box::new(spawn), } } @@ -128,17 +126,16 @@ impl ReviewerPool { impl ReviewerPool { /// Returns the current reviewer handle for host inspection and feedback collection. pub async fn trunk(&self) -> Option> { - self.state + self.trunk .lock() .await - .trunk .as_ref() .map(|trunk| Arc::clone(&trunk.session)) } /// Prepares the first reviewer without replacing a review that won the startup race. pub async fn prewarm(&self, setup: Arc, context: S::Context) -> anyhow::Result<()> { - let cancellation = self.cancellation.child_token(); + let cancellation = self.runtime.cancellation.child_token(); let guard = cancellation.clone().drop_guard(); let session = (self.spawn)( setup, @@ -148,44 +145,23 @@ impl ReviewerPool { cancellation.clone(), ) .await?; - let mut state = self.state.lock().await; - if !cancellation.is_cancelled() && state.trunk.is_none() { - state.trunk = Some(Arc::new(Trunk { + let mut trunk = self.trunk.lock().await; + if !cancellation.is_cancelled() && trunk.is_none() { + *trunk = Some(Arc::new(Trunk { session: Arc::new(session), review_lock: Semaphore::new(/*permits*/ 1), + cancellation: guard.disarm(), })); - drop(guard.disarm()); } Ok(()) } /// Permanently stops this parent's reviewer pool and waits for tracked runtimes. pub async fn shutdown(&self) { - self.cancellation.cancel(); - self.invalidate().await; - } - - /// Drops reusable context after parent history rollback or another host invalidation. - pub async fn invalidate(&self) { - let (trunk, ephemeral) = { - let mut state = self.state.lock().await; - ( - state.trunk.take(), - std::mem::take(&mut state.ephemeral_reviews), - ) - }; - for session in trunk - .into_iter() - .map(|trunk| Arc::clone(&trunk.session)) - .chain(ephemeral) - { - if self.cancellation.is_cancelled() { - session.shutdown().await; - } else { - session.cancel(); - shutdown_in_background(session); - } - } + self.runtime.cancellation.cancel(); + self.trunk.lock().await.take(); + self.runtime.tasks.close(); + self.runtime.tasks.wait().await; } /// Selects one reviewer; busy or incompatible trunks use an isolated temporary session. @@ -204,26 +180,24 @@ impl ReviewerPool { let (trunk, context) = match run_before_review_deadline( request.deadline(), request.cancellation(), - self.state.lock(), + self.trunk.lock(), ) .await { Ok(mut state) => { - 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 + let context = request.context(state.as_ref().map(|trunk| trunk.session.as_ref())); + if let Some(trunk) = state.as_ref() + && (trunk.cancellation.is_cancelled() || trunk.session.context() != &context) && trunk.review_lock.try_acquire().is_ok() - && let Some(stale) = state.trunk.take() { - shutdown_in_background(Arc::clone(&stale.session)); + state.take(); } - if state.trunk.is_none() { - let cancellation = self.cancellation.child_token(); - let session = match run_before_review_deadline_with_cancel( + if state.is_none() { + let cancellation = self.runtime.cancellation.child_token(); + let lifetime = cancellation.clone().drop_guard(); + let session = match run_before_review_deadline( request.deadline(), request.cancellation(), - &cancellation, (self.spawn)( request.setup(), context.clone(), @@ -245,13 +219,14 @@ impl ReviewerPool { return (outcome, GuardianReviewAnalyticsResult::without_session()); } }; - state.trunk = Some(Arc::new(Trunk { + *state = Some(Arc::new(Trunk { session, review_lock: Semaphore::new(/*permits*/ 1), + cancellation: lifetime.disarm(), })); spawned_trunk = true; } - (state.trunk.as_ref().cloned(), context) + (state.as_ref().cloned(), context) } Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), }; @@ -282,22 +257,27 @@ impl ReviewerPool { } else { GuardianReviewSessionKind::TrunkReused }; + // Dropping a review before it drains its turn must not leave a reusable agent. + let review_lifetime = trunk.cancellation.clone().drop_guard(); let (outcome, disposition, analytics) = request.run(&trunk.session, kind).await; if disposition == SessionDisposition::Reusable && matches!(outcome, GuardianReviewSessionOutcome::Completed(_)) { trunk.session.commit_snapshot().await; } + if disposition == SessionDisposition::Reusable { + review_lifetime.disarm(); + } else { + drop(review_lifetime); + } drop(guard); if disposition == SessionDisposition::Discard { - let mut state = self.state.lock().await; + let mut state = self.trunk.lock().await; if state - .trunk .as_ref() .is_some_and(|current| Arc::ptr_eq(current, &trunk)) - && let Some(removed) = state.trunk.take() { - shutdown_in_background(Arc::clone(&removed.session)); + state.take(); } } (outcome, analytics) @@ -312,11 +292,11 @@ impl ReviewerPool { where R: ReviewerRequest, { - let cancellation = self.cancellation.child_token(); - let session = match run_before_review_deadline_with_cancel( + let cancellation = self.runtime.cancellation.child_token(); + let _lifetime = cancellation.clone().drop_guard(); + let session = match run_before_review_deadline( request.deadline(), request.cancellation(), - &cancellation, (self.spawn)( request.setup(), context, @@ -336,63 +316,9 @@ impl ReviewerPool { } Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), }; - self.state - .lock() - .await - .ephemeral_reviews - .push(Arc::clone(&session)); - let mut cleanup = EphemeralCleanup { - state: Arc::clone(&self.state), - session: Some(Arc::clone(&session)), - }; let (outcome, _, analytics) = request .run(&session, GuardianReviewSessionKind::EphemeralForked) .await; - let removed = { - let mut state = self.state.lock().await; - state - .ephemeral_reviews - .iter() - .position(|active| Arc::ptr_eq(active, &session)) - .map(|index| state.ephemeral_reviews.swap_remove(index)) - }; - if let Some(removed) = removed { - cleanup.session = None; - shutdown_in_background(removed); - } (outcome, analytics) } } - -fn shutdown_in_background(session: Arc) { - drop(tokio::spawn(async move { - session.shutdown().await; - })); -} - -struct EphemeralCleanup { - state: Arc>>, - session: Option>, -} - -impl Drop for EphemeralCleanup { - fn drop(&mut self) { - let Some(session) = self.session.take() else { - return; - }; - let state = Arc::clone(&self.state); - drop(tokio::spawn(async move { - let removed = { - let mut state = state.lock().await; - state - .ephemeral_reviews - .iter() - .position(|active| Arc::ptr_eq(active, &session)) - .map(|index| state.ephemeral_reviews.swap_remove(index)) - }; - if let Some(removed) = removed { - removed.shutdown().await; - } - })); - } -} diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index 1cdb168392..35033d756b 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -38,4 +38,5 @@ uuid = { workspace = true } codex-model-provider-info = { workspace = true } core_test_support = { workspace = true } pretty_assertions = { workspace = true } +tempfile = { workspace = true } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/codex-rs/ext/guardian-v2/src/lib.rs b/codex-rs/ext/guardian-v2/src/lib.rs index 4c6e3feaa0..445eae6ae5 100644 --- a/codex-rs/ext/guardian-v2/src/lib.rs +++ b/codex-rs/ext/guardian-v2/src/lib.rs @@ -9,7 +9,7 @@ use codex_login::AuthManager; mod async_scorer; mod sync_reviewer; -pub use sync_reviewer::GuardianExtension as GuardianReviewerExtension; +pub use sync_reviewer::install as install_reviewer; /// Installs the guardian contributors into the extension registry. pub fn install( @@ -18,5 +18,5 @@ pub fn install( thread_manager: Weak, ) { async_scorer::install(registry, auth_manager, thread_manager.clone()); - sync_reviewer::install(registry, thread_manager); + install_reviewer(registry, thread_manager); } diff --git a/codex-rs/ext/guardian-v2/src/sync_reviewer/lifecycle_tests.rs b/codex-rs/ext/guardian-v2/src/sync_reviewer/lifecycle_tests.rs new file mode 100644 index 0000000000..da3728d1a3 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/sync_reviewer/lifecycle_tests.rs @@ -0,0 +1,125 @@ +//! Exercises Guardian registration and denial cleanup through real turns. + +use super::*; +use codex_core::TurnInputRequest; +use codex_protocol::openai_models::ModelInfo; +use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::Op; +use codex_protocol::protocol::TurnAbortReason; +use codex_protocol::user_input::UserInput; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::responses::start_mock_server; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use core_test_support::wait_for_event_match; +use pretty_assertions::assert_eq; +use std::time::Duration; + +struct StaleDenials(ModelInfo); + +impl TurnLifecycleContributor for StaleDenials { + fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(async move { + let denials = ReviewDenials::for_thread(input.thread_store); + for _ in 0..2 { + assert_eq!(denials.record_denial(input.turn_id, &self.0).await, None); + } + }) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn terminal_turn_clears_extension_owned_denials() -> anyhow::Result<()> { + let home = tempfile::tempdir()?; + let config = core_test_support::load_default_config_for_test(&home).await; + let model = codex_core::test_support::construct_model_info_offline("gpt-5", &config); + let mut extensions = ExtensionRegistryBuilder::::new(); + // Seed stale state before the production Guardian contributor runs its start hook. + extensions.turn_lifecycle_contributor(Arc::new(StaleDenials(model.clone()))); + let mut releases = Vec::new(); + let streams = (0..3) + .map(|_| { + let (release, gate) = tokio::sync::oneshot::channel(); + releases.push(release); + vec![ + StreamingSseChunk { + gate: None, + body: sse(vec![ev_response_created("response")]), + }, + StreamingSseChunk { + gate: Some(gate), + body: sse(vec![ev_completed("response")]), + }, + ] + }) + .collect(); + let (streaming, _) = start_streaming_sse_server(streams).await; + let server = start_mock_server().await; + let base_url = format!("{}/v1", streaming.uri()); + let test = test_codex() + .with_extensions(Arc::new(extensions.build())) + .with_config(move |config| { + config.model_provider.base_url = Some(base_url); + config.model_provider.supports_websockets = false; + }) + .build_with_auto_env(&server) + .await?; + let denials = ReviewDenials::for_thread(test.codex.thread_extension_data()); + + let mut previous_turn = None; + // The third turn also checks that an interrupted turn leaves the next turn usable. + for (index, release) in releases.into_iter().enumerate() { + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Continue.".into(), + text_elements: Vec::new(), + }])) + .await?; + let turn_id = wait_for_event_match(&test.codex, |event| match event { + EventMsg::TurnStarted(event) => Some(event.turn_id.clone()), + _ => None, + }) + .await; + // An interruption event can precede its cleanup. The next turn starts after it. + if let Some(previous) = previous_turn.replace(turn_id.clone()) { + assert_eq!(denials.record_denial(&previous, &model).await, None); + } + tokio::time::timeout( + Duration::from_secs(/*secs*/ 10), + streaming.wait_for_request_count(index + 1), + ) + .await?; + // Without start cleanup, these denials reach the three-denial limit. + for _ in 0..2 { + assert_eq!(denials.record_denial(&turn_id, &model).await, None); + } + if index == 1 { + test.codex.submit(Op::Interrupt).await?; + wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::TurnAborted(event) + if event.reason == TurnAbortReason::Interrupted + && event.turn_id.as_deref() == Some(turn_id.as_str())) + }) + .await; + drop(release); + } else { + release.send(()).unwrap(); + wait_for_event( + &test.codex, + |event| matches!(event, EventMsg::TurnComplete(event) if event.turn_id == turn_id), + ) + .await; + } + } + test.codex.shutdown_and_wait().await?; + assert_eq!( + denials.record_denial(&previous_turn.unwrap(), &model).await, + None + ); + streaming.shutdown().await; + Ok(()) +} 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 84cec63902..bc08398c6e 100644 --- a/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs +++ b/codex-rs/ext/guardian-v2/src/sync_reviewer/mod.rs @@ -16,6 +16,11 @@ use codex_extension_api::ThreadLifecycleContributor; use codex_extension_api::ThreadReadyInput; use codex_extension_api::ThreadStartInput; use codex_extension_api::ThreadStopInput; +use codex_extension_api::TurnAbortInput; +use codex_extension_api::TurnLifecycleContributor; +use codex_extension_api::TurnStartInput; +use codex_extension_api::TurnStopInput; +use codex_guardian_reviewer::ReviewDenials; use codex_guardian_reviewer::ReviewerPool; use codex_guardian_reviewer::ReviewerTasks; use codex_protocol::ThreadId; @@ -26,17 +31,11 @@ use codex_protocol::protocol::SessionSource; use codex_protocol::protocol::ThreadSource; /// Owns reviewer agents through the same thread manager as the parent conversation. -#[derive(Clone, Debug)] -pub struct GuardianExtension { +#[derive(Debug)] +struct GuardianExtension { thread_manager: Weak, } -impl GuardianExtension { - pub fn new(thread_manager: Weak) -> Self { - Self { thread_manager } - } -} - impl ThreadLifecycleContributor for GuardianExtension { fn on_thread_start<'a>( &'a self, @@ -50,6 +49,7 @@ impl ThreadLifecycleContributor for GuardianExtension { let runtime = input.thread_store.get_or_init(ReviewerTasks::default); input.thread_store.get_or_init(|| { ReviewerPool::::new( + Arc::clone(&runtime), move |context, key, kind, snapshot, cancel| { let manager = manager.clone(); let runtime = Arc::clone(&runtime); @@ -82,14 +82,12 @@ impl ThreadLifecycleContributor for GuardianExtension { 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 @@ -164,24 +162,44 @@ impl ThreadLifecycleContributor for GuardianExtension { 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; + if let Some(pool) = input + .thread_store + .get::>() + { + pool.shutdown().await; } }) } } +impl TurnLifecycleContributor for GuardianExtension { + fn on_turn_start<'a>(&'a self, input: TurnStartInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(ReviewDenials::clear_turn(input.thread_store, input.turn_id)) + } + fn on_turn_stop<'a>(&'a self, input: TurnStopInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(ReviewDenials::clear_turn( + input.thread_store, + input.turn_store.level_id(), + )) + } + fn on_turn_abort<'a>(&'a self, input: TurnAbortInput<'a>) -> ExtensionFuture<'a, ()> { + Box::pin(ReviewDenials::clear_turn( + input.thread_store, + input.turn_store.level_id(), + )) + } +} + +/// Registers the synchronous reviewer and its thread and turn cleanup. pub fn install( registry: &mut ExtensionRegistryBuilder, thread_manager: Weak, ) { - registry.thread_lifecycle_contributor(Arc::new(GuardianExtension::new(thread_manager))); + let extension = Arc::new(GuardianExtension { thread_manager }); + registry.thread_lifecycle_contributor(extension.clone()); + registry.turn_lifecycle_contributor(extension); } + +#[cfg(test)] +#[path = "lifecycle_tests.rs"] +mod tests;