From a2d506d9fbb871426a8a08247f7e740c554ea67c Mon Sep 17 00:00:00 2001 From: Ray Morgan Date: Tue, 9 Jun 2026 08:19:37 -0700 Subject: [PATCH] core: own Guardian reviews by parent turn --- codex-rs/core/src/guardian/mod.rs | 4 + codex-rs/core/src/guardian/review.rs | 56 ++++++- codex-rs/core/src/guardian/task_owner.rs | 143 ++++++++++++++++++ .../core/src/guardian/task_owner_tests.rs | 49 ++++++ codex-rs/core/src/session/review.rs | 1 + codex-rs/core/src/session/turn_context.rs | 5 + codex-rs/core/src/tasks/mod.rs | 16 +- 7 files changed, 265 insertions(+), 9 deletions(-) create mode 100644 codex-rs/core/src/guardian/task_owner.rs create mode 100644 codex-rs/core/src/guardian/task_owner_tests.rs diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index b4920f1ff6..88407b4629 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -16,6 +16,7 @@ mod metrics; mod prompt; mod review; mod review_session; +mod task_owner; use std::time::Duration; @@ -43,6 +44,9 @@ pub(crate) use review::routes_approval_to_guardian_with_reviewer; pub(crate) use review::spawn_approval_request_review; pub(crate) use review_session::GuardianReviewSessionManager; pub(crate) use review_session::prompt_cache_key_override_for_review_session; +pub(crate) use task_owner::GuardianReviewDrain; +pub(crate) use task_owner::GuardianReviewDrainOutcome; +pub(crate) use task_owner::GuardianReviewTaskOwner; pub(crate) const GUARDIAN_REVIEW_TIMEOUT: Duration = Duration::from_secs(90); pub(crate) const GUARDIAN_REVIEWER_NAME: &str = "guardian"; diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 934e68042f..722ba8d674 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -568,6 +568,50 @@ async fn run_guardian_review( } } +async fn run_guardian_review_in_task( + session: Arc, + turn: Arc, + review_id: String, + request: GuardianApprovalRequest, + retry_reason: Option, + approval_request_source: GuardianApprovalRequestSource, + owner_cancel: CancellationToken, +) -> ReviewDecision { + let Some(review_activity) = turn.guardian_reviews.begin() else { + return ReviewDecision::Abort; + }; + let cancel_activity = review_activity.clone(); + let request_cancel = owner_cancel.clone(); + let runtime_handle = session.services.runtime_handle.clone(); + let review = run_guardian_review( + session, + Arc::clone(&turn), + review_id, + request, + retry_reason, + approval_request_source, + Some(review_activity.cancellation_token()), + ); + let Some(review) = turn.guardian_reviews.spawn(&runtime_handle, async move { + tokio::pin!(review); + tokio::select! { + biased; + _ = request_cancel.cancelled() => { + cancel_activity.cancel(); + review.await + } + decision = &mut review => decision, + } + }) else { + return ReviewDecision::Abort; + }; + match review.await { + Ok(decision) => decision, + Err(err) if err.is_cancelled() || owner_cancel.is_cancelled() => ReviewDecision::Abort, + Err(_) => ReviewDecision::Denied, + } +} + /// Public entrypoint for approval requests that should be reviewed by guardian. pub(crate) async fn review_approval_request( session: &Arc, @@ -576,17 +620,15 @@ pub(crate) async fn review_approval_request( request: GuardianApprovalRequest, retry_reason: Option, ) -> ReviewDecision { - // Box the delegated review future so callers do not inline the entire - // guardian session state machine into their own async stack. - Box::pin(run_guardian_review( + run_guardian_review_in_task( Arc::clone(session), Arc::clone(turn), review_id, request, retry_reason, GuardianApprovalRequestSource::MainTurn, - /*external_cancel*/ None, - )) + turn.guardian_reviews.cancellation_token(), + ) .await } @@ -599,14 +641,14 @@ pub(crate) async fn review_approval_request_with_cancel( approval_request_source: GuardianApprovalRequestSource, cancel_token: CancellationToken, ) -> ReviewDecision { - run_guardian_review( + run_guardian_review_in_task( Arc::clone(session), Arc::clone(turn), review_id, request, retry_reason, approval_request_source, - Some(cancel_token), + cancel_token, ) .await } diff --git a/codex-rs/core/src/guardian/task_owner.rs b/codex-rs/core/src/guardian/task_owner.rs new file mode 100644 index 0000000000..90f36169cd --- /dev/null +++ b/codex-rs/core/src/guardian/task_owner.rs @@ -0,0 +1,143 @@ +use std::future::Future; +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; + +use tokio::runtime::Handle; +use tokio::task::AbortHandle; +use tokio::task::JoinHandle; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; +use tracing::warn; + +const GUARDIAN_REVIEW_DRAIN_TIMEOUT: Duration = Duration::from_secs(6); + +#[derive(Debug, PartialEq)] +pub(crate) enum GuardianReviewDrainOutcome { + Drained, + Forced, +} + +#[derive(Debug)] +struct TrackedGuardianReview { + abort_handle: AbortHandle, +} + +#[derive(Debug, Default)] +struct GuardianReviewTaskOwnerState { + closed_at: Option, + reviews: Vec, +} + +#[derive(Debug, Default)] +pub(crate) struct GuardianReviewTaskOwner { + cancellation_token: CancellationToken, + tasks: TaskTracker, + state: Mutex, +} + +impl GuardianReviewTaskOwner { + fn lock_state(&self) -> std::sync::MutexGuard<'_, GuardianReviewTaskOwnerState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancellation_token.child_token() + } + + pub(crate) fn begin(self: &Arc) -> Option { + if self.lock_state().closed_at.is_some() { + return None; + } + Some(GuardianReviewActivity { + cancellation_token: self.cancellation_token.child_token(), + }) + } + + pub(crate) fn spawn( + self: &Arc, + runtime_handle: &Handle, + future: F, + ) -> Option> + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + let state = self.lock_state(); + if state.closed_at.is_some() { + return None; + } + let task = self.tasks.spawn_on(future, runtime_handle); + let mut state = state; + state.reviews.push(TrackedGuardianReview { + abort_handle: task.abort_handle(), + }); + drop(state); + Some(task) + } + + pub(crate) fn close(self: &Arc) -> GuardianReviewDrain { + let closed_at = { + let mut state = self.lock_state(); + *state.closed_at.get_or_insert_with(|| { + self.cancellation_token.cancel(); + self.tasks.close(); + Instant::now() + }) + }; + GuardianReviewDrain { + owner: Arc::clone(self), + deadline: closed_at + GUARDIAN_REVIEW_DRAIN_TIMEOUT, + } + } +} + +#[derive(Clone)] +pub(crate) struct GuardianReviewActivity { + cancellation_token: CancellationToken, +} + +impl GuardianReviewActivity { + pub(crate) fn cancellation_token(&self) -> CancellationToken { + self.cancellation_token.clone() + } + + pub(crate) fn cancel(&self) { + self.cancellation_token.cancel(); + } +} + +#[must_use = "Guardian reviews must be drained after the parent turn is cancelled"] +pub(crate) struct GuardianReviewDrain { + owner: Arc, + deadline: Instant, +} + +impl GuardianReviewDrain { + pub(crate) async fn drain(self) -> GuardianReviewDrainOutcome { + let timed_out = tokio::time::timeout_at(self.deadline, self.owner.tasks.wait()) + .await + .is_err(); + let reviews = { + let mut state = self.owner.lock_state(); + std::mem::take(&mut state.reviews) + }; + if timed_out { + for review in &reviews { + review.abort_handle.abort(); + } + self.owner.tasks.wait().await; + warn!("timed out waiting for Guardian reviews to stop"); + GuardianReviewDrainOutcome::Forced + } else { + GuardianReviewDrainOutcome::Drained + } + } +} + +#[cfg(test)] +#[path = "task_owner_tests.rs"] +mod tests; diff --git a/codex-rs/core/src/guardian/task_owner_tests.rs b/codex-rs/core/src/guardian/task_owner_tests.rs new file mode 100644 index 0000000000..893ca2b4b3 --- /dev/null +++ b/codex-rs/core/src/guardian/task_owner_tests.rs @@ -0,0 +1,49 @@ +use pretty_assertions::assert_eq; +use std::sync::Arc; +use tokio::sync::oneshot; +use tokio::time::Instant; + +use super::GuardianReviewDrainOutcome; +use super::GuardianReviewTaskOwner; + +#[tokio::test] +async fn dropping_review_handle_leaves_cleanup_to_owner() { + let owner = Arc::new(GuardianReviewTaskOwner::default()); + let cancellation_token = owner.cancellation_token(); + let (completed_tx, completed_rx) = oneshot::channel(); + let task = owner + .spawn(&tokio::runtime::Handle::current(), async move { + cancellation_token.cancelled().await; + let _ = completed_tx.send(()); + }) + .expect("review task should start"); + + drop(task); + + assert_eq!( + owner.close().drain().await, + GuardianReviewDrainOutcome::Drained + ); + assert_eq!(completed_rx.await, Ok(())); +} + +#[tokio::test] +async fn forced_drain_aborts_review_task() { + let owner = Arc::new(GuardianReviewTaskOwner::default()); + let (drop_tx, mut drop_rx) = oneshot::channel::<()>(); + let task = owner + .spawn(&tokio::runtime::Handle::current(), async move { + let _drop_tx = drop_tx; + std::future::pending::<()>().await; + }) + .expect("review task should start"); + drop(task); + + let mut drain = owner.close(); + drain.deadline = Instant::now(); + assert_eq!(drain.drain().await, GuardianReviewDrainOutcome::Forced); + assert_eq!( + drop_rx.try_recv(), + Err(oneshot::error::TryRecvError::Closed) + ); +} diff --git a/codex-rs/core/src/session/review.rs b/codex-rs/core/src/session/review.rs index a3b23c6a78..b794a4d2f8 100644 --- a/codex-rs/core/src/session/review.rs +++ b/codex-rs/core/src/session/review.rs @@ -109,6 +109,7 @@ pub(super) async fn spawn_review_thread( let review_turn_context = TurnContext { sub_id: review_turn_id.clone(), + guardian_reviews: Arc::new(crate::guardian::GuardianReviewTaskOwner::default()), trace_id: current_span_trace_id(), realtime_active: parent_turn_context.realtime_active, config: per_turn_config, diff --git a/codex-rs/core/src/session/turn_context.rs b/codex-rs/core/src/session/turn_context.rs index be6e316a00..ecc7dda56d 100644 --- a/codex-rs/core/src/session/turn_context.rs +++ b/codex-rs/core/src/session/turn_context.rs @@ -20,6 +20,8 @@ use codex_sandboxing::policy_transforms::effective_network_sandbox_policy; use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use crate::guardian::GuardianReviewTaskOwner; + #[derive(Clone, Debug)] pub(crate) struct TurnSkillsContext { pub(crate) outcome: Arc, @@ -56,6 +58,7 @@ impl TurnEnvironment { #[derive(Debug)] pub struct TurnContext { pub(crate) sub_id: String, + pub(crate) guardian_reviews: Arc, pub(crate) trace_id: Option, pub(crate) realtime_active: bool, pub config: Arc, @@ -224,6 +227,7 @@ impl TurnContext { Self { sub_id: self.sub_id.clone(), + guardian_reviews: Arc::clone(&self.guardian_reviews), trace_id: self.trace_id.clone(), realtime_active: self.realtime_active, config: Arc::new(config), @@ -526,6 +530,7 @@ impl Session { extension_data.insert(HostLoadedSkills::new(Arc::clone(&skills_outcome))); TurnContext { sub_id, + guardian_reviews: Arc::new(GuardianReviewTaskOwner::default()), trace_id: current_span_trace_id(), realtime_active: false, config: per_turn_config.clone(), diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 911f30b319..bc111bd647 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -550,6 +550,12 @@ impl Session { true } + async fn drain_guardian_reviews(&self, drain: crate::guardian::GuardianReviewDrain) { + if drain.drain().await == crate::guardian::GuardianReviewDrainOutcome::Forced { + self.guardian_review_session.shutdown().await; + } + } + pub async fn on_task_finished( self: &Arc, turn_context: Arc, @@ -570,6 +576,8 @@ impl Session { let Some(turn_state) = turn_state else { return; }; + self.drain_guardian_reviews(turn_context.guardian_reviews.close()) + .await; let pending_input = self .input_queue .take_pending_input_for_turn_state(turn_state.as_ref()) @@ -784,12 +792,15 @@ impl Session { async fn handle_task_abort(self: &Arc, task: RunningTask, reason: TurnAbortReason) { let sub_id = task.turn_context.sub_id.clone(); - if task.cancellation_token.is_cancelled() { + let guardian_drain = task.turn_context.guardian_reviews.close(); + let task_already_cancelled = task.cancellation_token.is_cancelled(); + task.cancellation_token.cancel(); + if task_already_cancelled { + self.drain_guardian_reviews(guardian_drain).await; return; } trace!(task_kind = ?task.kind, sub_id, "aborting running task"); - task.cancellation_token.cancel(); task.turn_context .turn_metadata_state .cancel_git_enrichment_task(); @@ -812,6 +823,7 @@ impl Session { session_task .abort(session_ctx, Arc::clone(&task.turn_context)) .await; + self.drain_guardian_reviews(guardian_drain).await; if reason == TurnAbortReason::Interrupted && let Some(marker) = interrupted_turn_history_marker(