From c7645ca96e595dae80f7a8512e69046ffd09c551 Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Fri, 12 Jun 2026 00:01:37 -0700 Subject: [PATCH] guardian: move denial tracking into turn extension state --- codex-rs/Cargo.lock | 1 + codex-rs/core/src/guardian/mod.rs | 65 +----------- codex-rs/core/src/guardian/review.rs | 30 +++--- codex-rs/core/src/guardian/tests.rs | 87 --------------- codex-rs/core/src/session/session.rs | 1 - codex-rs/core/src/session/tests.rs | 2 - codex-rs/core/src/state/service.rs | 2 - codex-rs/core/src/tasks/mod.rs | 17 --- codex-rs/ext/guardian/Cargo.toml | 4 +- codex-rs/ext/guardian/src/circuit_breaker.rs | 58 ++++++++++ .../ext/guardian/src/circuit_breaker_tests.rs | 100 ++++++++++++++++++ codex-rs/ext/guardian/src/lib.rs | 6 ++ 12 files changed, 188 insertions(+), 185 deletions(-) create mode 100644 codex-rs/ext/guardian/src/circuit_breaker.rs create mode 100644 codex-rs/ext/guardian/src/circuit_breaker_tests.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 692ed5ad12..0792658925 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3078,6 +3078,7 @@ dependencies = [ "codex-shell-command", "codex-utils-absolute-path", "codex-utils-output-truncation", + "pretty_assertions", "serde", "serde_json", ] diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 1896caa6ad..67573273f0 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -34,6 +34,9 @@ pub(crate) use approval_request::guardian_request_turn_id; pub(crate) use codex_extension_api::ApprovalReviewMcpAnnotations as GuardianMcpAnnotations; pub(crate) use codex_extension_api::ApprovalReviewNetworkAccessTrigger as GuardianNetworkAccessTrigger; pub(crate) use codex_extension_api::ApprovalReviewRequest as GuardianApprovalRequest; +pub(crate) use codex_guardian::AUTO_REVIEW_DENIAL_WINDOW_SIZE; +pub(crate) use codex_guardian::GuardianRejectionCircuitBreaker; +pub(crate) use codex_guardian::GuardianRejectionCircuitBreakerAction; pub(crate) use review::guardian_rejection_message; pub(crate) use review::guardian_timeout_message; pub(crate) use review::is_guardian_reviewer_source; @@ -51,9 +54,6 @@ pub(crate) use review_session::prompt_cache_key_override_for_review_session; pub(crate) const GUARDIAN_REVIEW_TIMEOUT: Duration = Duration::from_secs(90); pub(crate) const GUARDIAN_REVIEWER_NAME: &str = "guardian"; -pub(crate) const MAX_CONSECUTIVE_GUARDIAN_DENIALS_PER_TURN: u32 = 3; -pub(crate) const MAX_RECENT_AUTO_REVIEW_DENIALS_PER_TURN: u32 = 10; -pub(crate) const AUTO_REVIEW_DENIAL_WINDOW_SIZE: usize = 50; pub(crate) const AUTO_REVIEW_DENIED_ACTION_APPROVAL_DEVELOPER_PREFIX: &str = "The user has manually approved a specific action that was previously `Rejected`."; const GUARDIAN_MAX_MESSAGE_TRANSCRIPT_TOKENS: usize = 10_000; @@ -77,65 +77,6 @@ pub(crate) struct GuardianRejection { pub(crate) source: GuardianAssessmentDecisionSource, } -#[derive(Debug, Default)] -pub(crate) struct GuardianRejectionCircuitBreaker { - turns: std::collections::HashMap, -} - -#[derive(Debug, Default)] -struct GuardianRejectionCircuitBreakerTurn { - consecutive_denials: u32, - recent_denials: std::collections::VecDeque, - interrupt_triggered: bool, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum GuardianRejectionCircuitBreakerAction { - Continue, - InterruptTurn { - consecutive_denials: u32, - recent_denials: u32, - }, -} - -impl GuardianRejectionCircuitBreaker { - pub(crate) fn clear_turn(&mut self, turn_id: &str) { - self.turns.remove(turn_id); - } - - pub(crate) fn record_denial(&mut self, turn_id: &str) -> GuardianRejectionCircuitBreakerAction { - let turn = self.turns.entry(turn_id.to_string()).or_default(); - turn.consecutive_denials = turn.consecutive_denials.saturating_add(1); - Self::record_recent_review(turn, /*denied*/ true); - let recent_denials = turn.recent_denials.iter().filter(|denied| **denied).count() as u32; - if !turn.interrupt_triggered - && (turn.consecutive_denials >= MAX_CONSECUTIVE_GUARDIAN_DENIALS_PER_TURN - || recent_denials >= MAX_RECENT_AUTO_REVIEW_DENIALS_PER_TURN) - { - turn.interrupt_triggered = true; - GuardianRejectionCircuitBreakerAction::InterruptTurn { - consecutive_denials: turn.consecutive_denials, - recent_denials, - } - } else { - GuardianRejectionCircuitBreakerAction::Continue - } - } - - pub(crate) fn record_non_denial(&mut self, turn_id: &str) { - let turn = self.turns.entry(turn_id.to_string()).or_default(); - turn.consecutive_denials = 0; - Self::record_recent_review(turn, /*denied*/ false); - } - - fn record_recent_review(turn: &mut GuardianRejectionCircuitBreakerTurn, denied: bool) { - turn.recent_denials.push_back(denied); - if turn.recent_denials.len() > AUTO_REVIEW_DENIAL_WINDOW_SIZE { - turn.recent_denials.pop_front(); - } - } -} - #[cfg(test)] use prompt::GuardianPromptMode; #[cfg(test)] diff --git a/codex-rs/core/src/guardian/review.rs b/codex-rs/core/src/guardian/review.rs index 99f9c5d999..fce67278a5 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -19,6 +19,7 @@ use codex_protocol::protocol::SubAgentSource; use codex_protocol::protocol::TurnAbortReason; use codex_protocol::protocol::WarningEvent; use std::sync::Arc; +use tokio::sync::Mutex; use tokio::sync::oneshot; use tokio::time::Instant; use tokio::time::sleep_until; @@ -211,22 +212,25 @@ fn track_guardian_review( .track_guardian_review(tracking, result, completed_at_ms); } -async fn record_guardian_non_denial(session: &Arc, turn_id: &str) { - session - .services - .guardian_rejection_circuit_breaker +fn guardian_rejection_circuit_breaker( + turn: &TurnContext, +) -> Arc> { + turn.extension_data + .get_or_init(|| Mutex::new(super::GuardianRejectionCircuitBreaker::default())) +} + +async fn record_guardian_non_denial(turn: &TurnContext) { + guardian_rejection_circuit_breaker(turn) .lock() .await - .record_non_denial(turn_id); + .record_non_denial(); } async fn record_guardian_denial(session: &Arc, turn: &Arc, turn_id: &str) { - let action = session - .services - .guardian_rejection_circuit_breaker + let action = guardian_rejection_circuit_breaker(turn) .lock() .await - .record_denial(turn_id); + .record_denial(); let GuardianRejectionCircuitBreakerAction::InterruptTurn { consecutive_denials, recent_denials, @@ -350,7 +354,7 @@ async fn run_guardian_review( }), ) .await; - record_guardian_non_denial(&session, &assessment_turn_id).await; + record_guardian_non_denial(&turn).await; return ReviewDecision::Abort; } @@ -443,7 +447,7 @@ async fn run_guardian_review( }), ) .await; - record_guardian_non_denial(&session, &assessment_turn_id).await; + record_guardian_non_denial(&turn).await; return ReviewDecision::TimedOut; } GuardianReviewError::Cancelled => { @@ -478,7 +482,7 @@ async fn run_guardian_review( }), ) .await; - record_guardian_non_denial(&session, &assessment_turn_id).await; + record_guardian_non_denial(&turn).await; return ReviewDecision::Abort; } GuardianReviewError::PromptBuild { .. } @@ -580,7 +584,7 @@ async fn run_guardian_review( if count_denial_for_circuit_breaker { record_guardian_denial(&session, &turn, &assessment_turn_id).await; } else { - record_guardian_non_denial(&session, &assessment_turn_id).await; + record_guardian_non_denial(&turn).await; } if approved { diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 1698ea0c68..352e0ce4e3 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -80,93 +80,6 @@ fn fixed_guardian_parent_session_id() -> ThreadId { .expect("fixed parent session id should be a valid UUID") } -#[test] -fn guardian_rejection_circuit_breaker_interrupts_after_three_consecutive_denials() { - let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::InterruptTurn { - consecutive_denials: 3, - recent_denials: 3, - } - ); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); -} - -#[test] -fn guardian_rejection_circuit_breaker_resets_consecutive_denials_on_non_denial() { - let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - circuit_breaker.record_non_denial("turn-1"); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::InterruptTurn { - consecutive_denials: 3, - recent_denials: 4, - } - ); -} - -#[test] -fn auto_review_rejection_circuit_breaker_interrupts_after_ten_recent_denials() { - let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); - for _ in 0..9 { - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - circuit_breaker.record_non_denial("turn-1"); - } - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::InterruptTurn { - consecutive_denials: 1, - recent_denials: 10, - } - ); -} - -#[test] -fn auto_review_rejection_circuit_breaker_forgets_denials_outside_recent_review_window() { - let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); - for _ in 0..9 { - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); - circuit_breaker.record_non_denial("turn-1"); - } - for _ in 0..(AUTO_REVIEW_DENIAL_WINDOW_SIZE - 18) { - circuit_breaker.record_non_denial("turn-1"); - } - assert_eq!( - circuit_breaker.record_denial("turn-1"), - GuardianRejectionCircuitBreakerAction::Continue - ); -} - async fn guardian_test_session_and_turn( server: &wiremock::MockServer, ) -> (Arc, Arc) { diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index b8d6d7a1bc..c5c740e33e 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -995,7 +995,6 @@ impl Session { models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), guardian_rejections: Mutex::new(HashMap::new()), - guardian_rejection_circuit_breaker: Mutex::new(Default::default()), runtime_handle: tokio::runtime::Handle::current(), skills_manager, plugins_manager: Arc::clone(&plugins_manager), diff --git a/codex-rs/core/src/session/tests.rs b/codex-rs/core/src/session/tests.rs index 0dbb13fcd2..59db9017e7 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4969,7 +4969,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), guardian_rejections: Mutex::new(std::collections::HashMap::new()), - guardian_rejection_circuit_breaker: Mutex::new(Default::default()), runtime_handle: tokio::runtime::Handle::current(), skills_manager, plugins_manager, @@ -7044,7 +7043,6 @@ where models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), guardian_rejections: Mutex::new(std::collections::HashMap::new()), - guardian_rejection_circuit_breaker: Mutex::new(Default::default()), runtime_handle: tokio::runtime::Handle::current(), skills_manager, plugins_manager, diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index d8d91e85a8..dd3a03913c 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -9,7 +9,6 @@ use crate::config::NetworkProxyAuditMetadata; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; use crate::guardian::GuardianRejection; -use crate::guardian::GuardianRejectionCircuitBreaker; use crate::mcp::McpManager; use crate::tools::code_mode::CodeModeService; use crate::tools::network_approval::NetworkApprovalService; @@ -59,7 +58,6 @@ pub(crate) struct SessionServices { pub(crate) session_telemetry: SessionTelemetry, pub(crate) tool_approvals: Mutex, pub(crate) guardian_rejections: Mutex>, - pub(crate) guardian_rejection_circuit_breaker: Mutex, pub(crate) runtime_handle: Handle, pub(crate) skills_manager: Arc, pub(crate) plugins_manager: Arc, diff --git a/codex-rs/core/src/tasks/mod.rs b/codex-rs/core/src/tasks/mod.rs index 61f6980156..f11f86f1c1 100644 --- a/codex-rs/core/src/tasks/mod.rs +++ b/codex-rs/core/src/tasks/mod.rs @@ -337,12 +337,6 @@ impl Session { let cancellation_token = CancellationToken::new(); let done = Arc::new(Notify::new()); - self.services - .guardian_rejection_circuit_breaker - .lock() - .await - .clear_turn(&turn_context.sub_id); - let pending_items = self.input_queue.get_pending_input(&self.active_turn).await; let turn_state = { let mut active = self.active_turn.lock().await; @@ -749,12 +743,6 @@ impl Session { time_to_first_token_ms, }); self.send_event(turn_context.as_ref(), event).await; - self.services - .guardian_rejection_circuit_breaker - .lock() - .await - .clear_turn(&turn_context.sub_id); - let cleared_active_turn = { let mut active = self.active_turn.lock().await; if let Some(active_turn) = active.as_ref() @@ -865,11 +853,6 @@ impl Session { duration_ms, }); self.send_event(task.turn_context.as_ref(), event).await; - self.services - .guardian_rejection_circuit_breaker - .lock() - .await - .clear_turn(&task.turn_context.sub_id); } } diff --git a/codex-rs/ext/guardian/Cargo.toml b/codex-rs/ext/guardian/Cargo.toml index 460c5ecd96..1467df7c9f 100644 --- a/codex-rs/ext/guardian/Cargo.toml +++ b/codex-rs/ext/guardian/Cargo.toml @@ -7,7 +7,6 @@ version.workspace = true [lib] name = "codex_guardian" path = "src/lib.rs" -test = false doctest = false [lints] @@ -22,3 +21,6 @@ codex-utils-absolute-path = { workspace = true } codex-utils-output-truncation = { workspace = true } serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } + +[dev-dependencies] +pretty_assertions = { workspace = true } diff --git a/codex-rs/ext/guardian/src/circuit_breaker.rs b/codex-rs/ext/guardian/src/circuit_breaker.rs new file mode 100644 index 0000000000..07a4be7f29 --- /dev/null +++ b/codex-rs/ext/guardian/src/circuit_breaker.rs @@ -0,0 +1,58 @@ +use std::collections::VecDeque; + +pub const MAX_CONSECUTIVE_GUARDIAN_DENIALS_PER_TURN: u32 = 3; +pub const MAX_RECENT_AUTO_REVIEW_DENIALS_PER_TURN: u32 = 10; +pub const AUTO_REVIEW_DENIAL_WINDOW_SIZE: usize = 50; + +/// Per-turn denial history used to stop repeated unsafe approval requests. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct GuardianRejectionCircuitBreaker { + consecutive_denials: u32, + recent_denials: VecDeque, + interrupt_triggered: bool, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum GuardianRejectionCircuitBreakerAction { + Continue, + InterruptTurn { + consecutive_denials: u32, + recent_denials: u32, + }, +} + +impl GuardianRejectionCircuitBreaker { + pub fn record_denial(&mut self) -> GuardianRejectionCircuitBreakerAction { + self.consecutive_denials = self.consecutive_denials.saturating_add(1); + self.record_recent_review(/*denied*/ true); + let recent_denials = self.recent_denials.iter().filter(|denied| **denied).count() as u32; + if !self.interrupt_triggered + && (self.consecutive_denials >= MAX_CONSECUTIVE_GUARDIAN_DENIALS_PER_TURN + || recent_denials >= MAX_RECENT_AUTO_REVIEW_DENIALS_PER_TURN) + { + self.interrupt_triggered = true; + GuardianRejectionCircuitBreakerAction::InterruptTurn { + consecutive_denials: self.consecutive_denials, + recent_denials, + } + } else { + GuardianRejectionCircuitBreakerAction::Continue + } + } + + pub fn record_non_denial(&mut self) { + self.consecutive_denials = 0; + self.record_recent_review(/*denied*/ false); + } + + fn record_recent_review(&mut self, denied: bool) { + self.recent_denials.push_back(denied); + if self.recent_denials.len() > AUTO_REVIEW_DENIAL_WINDOW_SIZE { + self.recent_denials.pop_front(); + } + } +} + +#[cfg(test)] +#[path = "circuit_breaker_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian/src/circuit_breaker_tests.rs b/codex-rs/ext/guardian/src/circuit_breaker_tests.rs new file mode 100644 index 0000000000..e092bda7b2 --- /dev/null +++ b/codex-rs/ext/guardian/src/circuit_breaker_tests.rs @@ -0,0 +1,100 @@ +use std::collections::VecDeque; + +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn interrupts_once_after_three_consecutive_denials() { + let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::InterruptTurn { + consecutive_denials: 3, + recent_denials: 3, + } + ); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + assert_eq!( + circuit_breaker, + GuardianRejectionCircuitBreaker { + consecutive_denials: 4, + recent_denials: VecDeque::from([true, true, true, true]), + interrupt_triggered: true, + } + ); +} + +#[test] +fn non_denial_resets_consecutive_denials() { + let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + circuit_breaker.record_non_denial(); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::InterruptTurn { + consecutive_denials: 3, + recent_denials: 4, + } + ); +} + +#[test] +fn interrupts_after_ten_recent_denials() { + let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); + for _ in 0..9 { + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + circuit_breaker.record_non_denial(); + } + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::InterruptTurn { + consecutive_denials: 1, + recent_denials: 10, + } + ); +} + +#[test] +fn forgets_denials_outside_recent_review_window() { + let mut circuit_breaker = GuardianRejectionCircuitBreaker::default(); + for _ in 0..9 { + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); + circuit_breaker.record_non_denial(); + } + for _ in 0..(AUTO_REVIEW_DENIAL_WINDOW_SIZE - 18) { + circuit_breaker.record_non_denial(); + } + assert_eq!( + circuit_breaker.record_denial(), + GuardianRejectionCircuitBreakerAction::Continue + ); +} diff --git a/codex-rs/ext/guardian/src/lib.rs b/codex-rs/ext/guardian/src/lib.rs index c51bdc9532..09f22003e9 100644 --- a/codex-rs/ext/guardian/src/lib.rs +++ b/codex-rs/ext/guardian/src/lib.rs @@ -14,6 +14,7 @@ use codex_protocol::config_types::ApprovalsReviewer; use codex_protocol::protocol::AskForApproval; mod approval_request; +mod circuit_breaker; mod truncation; pub use approval_request::FormattedGuardianAction; @@ -23,6 +24,11 @@ pub use approval_request::guardian_assessment_action; pub use approval_request::guardian_request_target_item_id; pub use approval_request::guardian_request_turn_id; pub use approval_request::guardian_reviewed_action; +pub use circuit_breaker::AUTO_REVIEW_DENIAL_WINDOW_SIZE; +pub use circuit_breaker::GuardianRejectionCircuitBreaker; +pub use circuit_breaker::GuardianRejectionCircuitBreakerAction; +pub use circuit_breaker::MAX_CONSECUTIVE_GUARDIAN_DENIALS_PER_TURN; +pub use circuit_breaker::MAX_RECENT_AUTO_REVIEW_DENIALS_PER_TURN; pub use truncation::guardian_truncate_text; /// Guardian extension dependencies supplied by the host at construction time.