core: own Guardian reviews by parent turn

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";

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<Session>,
+    turn: Arc<TurnContext>,
+    review_id: String,
+    request: GuardianApprovalRequest,
+    retry_reason: Option<String>,
+    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<Session>,
@@ -576,17 +620,15 @@ pub(crate) async fn review_approval_request(
     request: GuardianApprovalRequest,
     retry_reason: Option<String>,
 ) -> 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
 }

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<Instant>,
+    reviews: Vec<TrackedGuardianReview>,
+}
+
+#[derive(Debug, Default)]
+pub(crate) struct GuardianReviewTaskOwner {
+    cancellation_token: CancellationToken,
+    tasks: TaskTracker,
+    state: Mutex<GuardianReviewTaskOwnerState>,
+}
+
+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<Self>) -> Option<GuardianReviewActivity> {
+        if self.lock_state().closed_at.is_some() {
+            return None;
+        }
+        Some(GuardianReviewActivity {
+            cancellation_token: self.cancellation_token.child_token(),
+        })
+    }
+
+    pub(crate) fn spawn<F>(
+        self: &Arc<Self>,
+        runtime_handle: &Handle,
+        future: F,
+    ) -> Option<JoinHandle<F::Output>>
+    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<Self>) -> 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<GuardianReviewTaskOwner>,
+    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;

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)
+    );
+}

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,

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<SkillLoadOutcome>,
@@ -56,6 +58,7 @@ impl TurnEnvironment {
 #[derive(Debug)]
 pub struct TurnContext {
     pub(crate) sub_id: String,
+    pub(crate) guardian_reviews: Arc<GuardianReviewTaskOwner>,
     pub(crate) trace_id: Option<String>,
     pub(crate) realtime_active: bool,
     pub config: Arc<Config>,
@@ -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(),

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<Self>,
         turn_context: Arc<TurnContext>,
@@ -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<Self>, 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(