diff --git a/codex-rs/core/src/guardian/review_session.rs b/codex-rs/core/src/guardian/review_session.rs index c139a910a7..f6efdb02c6 100644 --- a/codex-rs/core/src/guardian/review_session.rs +++ b/codex-rs/core/src/guardian/review_session.rs @@ -303,6 +303,32 @@ impl GuardianReviewSessionManager { } } + pub(crate) async fn interrupt_active_reviews(&self) { + let review_sessions = { + let state = self.state.lock().await; + state + .trunk + .iter() + .chain(&state.ephemeral_reviews) + .cloned() + .collect::>() + }; + let active_review_sessions = review_sessions + .into_iter() + .filter(|review_session| review_session.review_lock.try_acquire().is_err()) + .collect::>(); + for review_session in &active_review_sessions { + let _ = review_session.codex.submit(Op::Interrupt).await; + } + for review_session in active_review_sessions { + let _ = tokio::time::timeout( + GUARDIAN_INTERRUPT_DRAIN_TIMEOUT, + review_session.review_lock.acquire(), + ) + .await; + } + } + #[expect( clippy::await_holding_invalid_type, reason = "review session selection and trunk spawning must stay serialized" @@ -561,6 +587,17 @@ impl GuardianReviewSessionManager { } Err(outcome) => return (outcome, GuardianReviewAnalyticsResult::without_session()), }; + let review_guard = match review_session.review_lock.acquire().await { + Ok(review_guard) => review_guard, + Err(err) => { + return ( + GuardianReviewSessionOutcome::SessionFailed(anyhow!( + "guardian review lock closed: {err}" + )), + GuardianReviewAnalyticsResult::without_session(), + ); + } + }; self.register_active_ephemeral(Arc::clone(&review_session)) .await; let mut cleanup = @@ -573,6 +610,7 @@ impl GuardianReviewSessionManager { deadline, )) .await; + drop(review_guard); if let Some(review_session) = self.take_active_ephemeral(&review_session).await { cleanup.disarm(); review_session.shutdown_in_background(); diff --git a/codex-rs/core/src/session/mod.rs b/codex-rs/core/src/session/mod.rs index 8ce16de740..b0d634e051 100644 --- a/codex-rs/core/src/session/mod.rs +++ b/codex-rs/core/src/session/mod.rs @@ -3298,6 +3298,9 @@ impl Session { pub async fn interrupt_task(self: &Arc) { info!("interrupt received: abort current task, if any"); let had_active_turn = self.active_turn.lock().await.is_some(); + self.guardian_review_session + .interrupt_active_reviews() + .await; self.abort_all_tasks(TurnAbortReason::Interrupted).await; if !had_active_turn { self.cancel_mcp_startup().await; diff --git a/codex-rs/core/tests/suite/guardian_review.rs b/codex-rs/core/tests/suite/guardian_review.rs index df38fb4353..7c363a5654 100644 --- a/codex-rs/core/tests/suite/guardian_review.rs +++ b/codex-rs/core/tests/suite/guardian_review.rs @@ -6,6 +6,7 @@ use codex_core::sandboxing::SandboxPermissions; use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::EventMsg; +use codex_protocol::protocol::GuardianAssessmentStatus; use codex_protocol::protocol::Op; use codex_protocol::protocol::SandboxPolicy; use codex_protocol::user_input::UserInput; @@ -14,8 +15,10 @@ use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::responses::ev_function_call; use core_test_support::responses::ev_response_created; +use core_test_support::responses::mount_response_sequence; use core_test_support::responses::mount_sse_sequence; use core_test_support::responses::sse; +use core_test_support::responses::sse_response; use core_test_support::responses::start_mock_server; use core_test_support::skip_if_no_network; use core_test_support::skip_if_sandbox; @@ -29,6 +32,105 @@ use std::os::unix::fs::PermissionsExt; use std::time::Duration; use tempfile::TempDir; +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn interrupting_parent_turn_shuts_down_active_guardian_review() -> Result<()> { + skip_if_no_network!(Ok(())); + skip_if_sandbox!(Ok(())); + + let server = start_mock_server().await; + let approval_policy = AskForApproval::OnRequest; + let sandbox_policy = SandboxPolicy::WorkspaceWrite { + writable_roots: vec![], + network_access: false, + exclude_tmpdir_env_var: true, + exclude_slash_tmp: true, + }; + let sandbox_policy_for_config = sandbox_policy.clone(); + let mut builder = test_codex().with_config(move |config| { + config.permissions.approval_policy = Constrained::allow_any(approval_policy); + config + .set_legacy_sandbox_policy(sandbox_policy_for_config) + .expect("set sandbox policy"); + }); + let test = builder.build(&server).await?; + + let tool_args = json!({ + "cmd": "echo guardian-review", + "yield_time_ms": 1_000_u64, + "sandbox_permissions": SandboxPermissions::RequireEscalated, + "justification": "Exercise Guardian cancellation.", + }); + let parent_response = sse_response(sse(vec![ + ev_response_created("resp-parent-tool"), + ev_function_call( + "exec-call", + "exec_command", + &serde_json::to_string(&tool_args)?, + ), + ev_completed("resp-parent-tool"), + ])); + let guardian_response = sse_response(sse(vec![ev_response_created("resp-guardian-delayed")])) + .set_delay(Duration::from_secs(60)); + let responses = + mount_response_sequence(&server, vec![parent_response, guardian_response]).await; + + test.codex + .submit(Op::UserInput { + environments: None, + items: vec![UserInput::Text { + text: "run a command that requires Guardian review".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + responsesapi_client_metadata: None, + additional_context: Default::default(), + thread_settings: codex_protocol::protocol::ThreadSettingsOverrides { + cwd: Some(test.cwd.path().to_path_buf()), + approval_policy: Some(approval_policy), + approvals_reviewer: Some(ApprovalsReviewer::AutoReview), + sandbox_policy: Some(sandbox_policy), + ..Default::default() + }, + }) + .await?; + let guardian_started = wait_for_event(&test.codex, |event| { + matches!(event, EventMsg::GuardianAssessment(_)) + }) + .await; + let EventMsg::GuardianAssessment(guardian_started) = guardian_started else { + unreachable!("wait predicate only accepts guardian assessments"); + }; + tokio::time::timeout(Duration::from_secs(5), async { + while responses.requests().len() < 2 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("guardian review request should be in flight before interrupt"); + + test.codex.submit(Op::Interrupt).await?; + let guardian_statuses = tokio::time::timeout(Duration::from_secs(5), async { + let mut statuses = vec![guardian_started.status]; + loop { + match test.codex.next_event().await?.msg { + EventMsg::GuardianAssessment(event) => statuses.push(event.status), + EventMsg::TurnAborted(_) => return Ok::<_, anyhow::Error>(statuses), + _ => {} + } + } + }) + .await??; + assert_eq!( + guardian_statuses, + vec![ + GuardianAssessmentStatus::InProgress, + GuardianAssessmentStatus::Aborted, + ] + ); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn guardian_review_session_does_not_inherit_legacy_notify() -> Result<()> { skip_if_no_network!(Ok(()));