From ac776a36f1ec26dcfcedb218d3e3bb1bf9dff368 Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Fri, 12 Jun 2026 00:09:58 -0700 Subject: [PATCH] guardian: move rejection rationales into turn extension state --- codex-rs/core/src/guardian/mod.rs | 9 +---- codex-rs/core/src/guardian/review.rs | 14 ++++--- codex-rs/core/src/guardian/tests.rs | 13 +++++-- codex-rs/core/src/session/session.rs | 1 - codex-rs/core/src/session/tests.rs | 2 - codex-rs/core/src/state/service.rs | 3 -- codex-rs/core/src/tools/approval_dispatch.rs | 2 +- codex-rs/core/src/tools/orchestrator.rs | 2 +- .../tools/runtimes/shell/unix_escalation.rs | 21 +++++------ codex-rs/ext/guardian/src/lib.rs | 3 ++ codex-rs/ext/guardian/src/rejection_store.rs | 37 +++++++++++++++++++ .../ext/guardian/src/rejection_store_tests.rs | 30 +++++++++++++++ 12 files changed, 102 insertions(+), 35 deletions(-) create mode 100644 codex-rs/ext/guardian/src/rejection_store.rs create mode 100644 codex-rs/ext/guardian/src/rejection_store_tests.rs diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index 67573273f0..ebe8c67698 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -19,7 +19,6 @@ mod review_session; use std::time::Duration; -use codex_protocol::protocol::GuardianAssessmentDecisionSource; use codex_protocol::protocol::GuardianAssessmentOutcome; use serde::Deserialize; use serde::Serialize; @@ -35,8 +34,10 @@ pub(crate) use codex_extension_api::ApprovalReviewMcpAnnotations as GuardianMcpA 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::GuardianRejection; pub(crate) use codex_guardian::GuardianRejectionCircuitBreaker; pub(crate) use codex_guardian::GuardianRejectionCircuitBreakerAction; +pub(crate) use codex_guardian::GuardianRejectionStore; pub(crate) use review::guardian_rejection_message; pub(crate) use review::guardian_timeout_message; pub(crate) use review::is_guardian_reviewer_source; @@ -71,12 +72,6 @@ pub(crate) struct GuardianAssessment { pub(crate) rationale: String, } -#[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) struct GuardianRejection { - pub(crate) rationale: String, - pub(crate) source: GuardianAssessmentDecisionSource, -} - #[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 fce67278a5..b1b4691571 100644 --- a/codex-rs/core/src/guardian/review.rs +++ b/codex-rs/core/src/guardian/review.rs @@ -69,10 +69,13 @@ pub(crate) fn new_guardian_review_id() -> String { uuid::Uuid::new_v4().to_string() } -pub(crate) async fn guardian_rejection_message(session: &Session, review_id: &str) -> String { - let rejection = session - .services - .guardian_rejections +fn guardian_rejection_store(turn: &TurnContext) -> Arc> { + turn.extension_data + .get_or_init(|| Mutex::new(super::GuardianRejectionStore::default())) +} + +pub(crate) async fn guardian_rejection_message(turn: &TurnContext, review_id: &str) -> String { + let rejection = guardian_rejection_store(turn) .lock() .await .remove(review_id) @@ -551,7 +554,8 @@ async fn run_guardian_review( GuardianAssessmentStatus::Denied }; { - let mut rationales = session.services.guardian_rejections.lock().await; + let rejection_store = guardian_rejection_store(&turn); + let mut rationales = rejection_store.lock().await; if approved { rationales.remove(&review_id); } else { diff --git a/codex-rs/core/src/guardian/tests.rs b/codex-rs/core/src/guardian/tests.rs index 352e0ce4e3..68d2a72dc4 100644 --- a/codex-rs/core/src/guardian/tests.rs +++ b/codex-rs/core/src/guardian/tests.rs @@ -73,6 +73,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; use tempfile::TempDir; +use tokio::sync::Mutex; use tokio_util::sync::CancellationToken; fn fixed_guardian_parent_session_id() -> ThreadId { @@ -2031,12 +2032,16 @@ async fn guardian_review_surfaces_responses_api_errors_in_rejection_reason() -> "denial rationale should not fall back to the generic missing payload error" ); { - let rationales = session.services.guardian_rejections.lock().await; - assert!(rationales.contains_key("review-shell-guardian-error")); - assert!(!rationales.contains_key("shell-guardian-error")); + let rationales = turn + .extension_data + .get::>() + .expect("guardian rejection store should exist"); + let rationales = rationales.lock().await; + assert!(rationales.contains("review-shell-guardian-error")); + assert!(!rationales.contains("shell-guardian-error")); } let rejection_message = - guardian_rejection_message(session.as_ref(), "review-shell-guardian-error").await; + guardian_rejection_message(turn.as_ref(), "review-shell-guardian-error").await; assert!( rejection_message.contains("Reason: Automatic approval review failed:") && rejection_message.contains(error_message), diff --git a/codex-rs/core/src/session/session.rs b/codex-rs/core/src/session/session.rs index c5c740e33e..09d361c5f3 100644 --- a/codex-rs/core/src/session/session.rs +++ b/codex-rs/core/src/session/session.rs @@ -994,7 +994,6 @@ impl Session { session_telemetry, models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), - guardian_rejections: Mutex::new(HashMap::new()), 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 59db9017e7..61f52e772f 100644 --- a/codex-rs/core/src/session/tests.rs +++ b/codex-rs/core/src/session/tests.rs @@ -4968,7 +4968,6 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) { session_telemetry: session_telemetry.clone(), models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), - guardian_rejections: Mutex::new(std::collections::HashMap::new()), runtime_handle: tokio::runtime::Handle::current(), skills_manager, plugins_manager, @@ -7042,7 +7041,6 @@ where session_telemetry: session_telemetry.clone(), models_manager: Arc::clone(&models_manager), tool_approvals: Mutex::new(ApprovalStore::default()), - guardian_rejections: Mutex::new(std::collections::HashMap::new()), 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 dd3a03913c..5116aa3a0f 100644 --- a/codex-rs/core/src/state/service.rs +++ b/codex-rs/core/src/state/service.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::sync::Arc; use crate::SkillsManager; @@ -8,7 +7,6 @@ use crate::client::ModelClient; use crate::config::NetworkProxyAuditMetadata; use crate::config::StartedNetworkProxy; use crate::exec_policy::ExecPolicyManager; -use crate::guardian::GuardianRejection; use crate::mcp::McpManager; use crate::tools::code_mode::CodeModeService; use crate::tools::network_approval::NetworkApprovalService; @@ -57,7 +55,6 @@ pub(crate) struct SessionServices { pub(crate) models_manager: SharedModelsManager, pub(crate) session_telemetry: SessionTelemetry, pub(crate) tool_approvals: Mutex, - pub(crate) guardian_rejections: Mutex>, pub(crate) runtime_handle: Handle, pub(crate) skills_manager: Arc, pub(crate) plugins_manager: Arc, diff --git a/codex-rs/core/src/tools/approval_dispatch.rs b/codex-rs/core/src/tools/approval_dispatch.rs index 4b058aa031..a4ccc73b1c 100644 --- a/codex-rs/core/src/tools/approval_dispatch.rs +++ b/codex-rs/core/src/tools/approval_dispatch.rs @@ -103,7 +103,7 @@ impl ApprovalReviewRunner for CoreApprovalReviewRunner<'_> { }; let denial_message = match decision { ReviewDecision::Denied | ReviewDecision::Abort => { - Some(guardian_rejection_message(self.session, self.review_id).await) + Some(guardian_rejection_message(self.turn, self.review_id).await) } ReviewDecision::TimedOut => Some(guardian_timeout_message()), _ => None, diff --git a/codex-rs/core/src/tools/orchestrator.rs b/codex-rs/core/src/tools/orchestrator.rs index eeafa8823f..7cbbbba69d 100644 --- a/codex-rs/core/src/tools/orchestrator.rs +++ b/codex-rs/core/src/tools/orchestrator.rs @@ -486,7 +486,7 @@ impl ToolOrchestrator { match decision { ReviewDecision::Denied | ReviewDecision::Abort => { let reason = if let Some(review_id) = guardian_review_id { - guardian_rejection_message(tool_ctx.session.as_ref(), review_id).await + guardian_rejection_message(tool_ctx.turn.as_ref(), review_id).await } else { "rejected by user".to_string() }; diff --git a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs index 65a655a79f..f74e4f7f9a 100644 --- a/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs +++ b/codex-rs/core/src/tools/runtimes/shell/unix_escalation.rs @@ -570,17 +570,16 @@ impl CoreShellActionProvider { } }, ReviewDecision::Denied => { - let message = if let Some(message) = - prompt_decision.rejection_message.clone() - { - message - } else if let Some(review_id) = - prompt_decision.guardian_review_id.as_deref() - { - guardian_rejection_message(self.session.as_ref(), review_id).await - } else { - "User denied execution".to_string() - }; + let message = + if let Some(message) = prompt_decision.rejection_message.clone() { + message + } else if let Some(review_id) = + prompt_decision.guardian_review_id.as_deref() + { + guardian_rejection_message(self.turn.as_ref(), review_id).await + } else { + "User denied execution".to_string() + }; EscalationDecision::deny(Some(message)) } ReviewDecision::TimedOut => { diff --git a/codex-rs/ext/guardian/src/lib.rs b/codex-rs/ext/guardian/src/lib.rs index 09f22003e9..45489eb7bf 100644 --- a/codex-rs/ext/guardian/src/lib.rs +++ b/codex-rs/ext/guardian/src/lib.rs @@ -15,6 +15,7 @@ use codex_protocol::protocol::AskForApproval; mod approval_request; mod circuit_breaker; +mod rejection_store; mod truncation; pub use approval_request::FormattedGuardianAction; @@ -29,6 +30,8 @@ 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 rejection_store::GuardianRejection; +pub use rejection_store::GuardianRejectionStore; pub use truncation::guardian_truncate_text; /// Guardian extension dependencies supplied by the host at construction time. diff --git a/codex-rs/ext/guardian/src/rejection_store.rs b/codex-rs/ext/guardian/src/rejection_store.rs new file mode 100644 index 0000000000..e144aa66a3 --- /dev/null +++ b/codex-rs/ext/guardian/src/rejection_store.rs @@ -0,0 +1,37 @@ +use std::collections::HashMap; + +use codex_protocol::protocol::GuardianAssessmentDecisionSource; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GuardianRejection { + pub rationale: String, + pub source: GuardianAssessmentDecisionSource, +} + +/// Rejection rationales awaiting delivery to callers for one turn. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct GuardianRejectionStore { + rejections: HashMap, +} + +impl GuardianRejectionStore { + pub fn insert( + &mut self, + review_id: String, + rejection: GuardianRejection, + ) -> Option { + self.rejections.insert(review_id, rejection) + } + + pub fn remove(&mut self, review_id: &str) -> Option { + self.rejections.remove(review_id) + } + + pub fn contains(&self, review_id: &str) -> bool { + self.rejections.contains_key(review_id) + } +} + +#[cfg(test)] +#[path = "rejection_store_tests.rs"] +mod tests; diff --git a/codex-rs/ext/guardian/src/rejection_store_tests.rs b/codex-rs/ext/guardian/src/rejection_store_tests.rs new file mode 100644 index 0000000000..fa6f7ea37f --- /dev/null +++ b/codex-rs/ext/guardian/src/rejection_store_tests.rs @@ -0,0 +1,30 @@ +use std::collections::HashMap; + +use codex_protocol::protocol::GuardianAssessmentDecisionSource; +use pretty_assertions::assert_eq; + +use super::*; + +#[test] +fn stores_rejections_by_review_id_and_removes_them_once() { + let rejection = GuardianRejection { + rationale: "unsafe command".to_string(), + source: GuardianAssessmentDecisionSource::Agent, + }; + let mut store = GuardianRejectionStore::default(); + + assert_eq!( + store.insert("review-1".to_string(), rejection.clone()), + None + ); + assert!(store.contains("review-1")); + assert_eq!( + store, + GuardianRejectionStore { + rejections: HashMap::from([("review-1".to_string(), rejection.clone())]), + } + ); + assert_eq!(store.remove("review-1"), Some(rejection)); + assert_eq!(store.remove("review-1"), None); + assert_eq!(store, GuardianRejectionStore::default()); +}