diff --git a/codex-rs/ext/guardian-v2/src/extension.rs b/codex-rs/ext/guardian-v2/src/extension.rs index d0393d9bce..9006415251 100644 --- a/codex-rs/ext/guardian-v2/src/extension.rs +++ b/codex-rs/ext/guardian-v2/src/extension.rs @@ -202,6 +202,7 @@ const CLASSIFICATION_DURATION_METRIC: &str = "codex.guardian_v2.classification.d struct GuardianV2ScoreProgress { latest_tool_call: AtomicUsize, latest_scored_tool_call: AtomicUsize, + latest_failed_tool_call: AtomicUsize, metrics: Option>, } @@ -324,14 +325,13 @@ impl ApprovalReviewContributor for GuardianV2Extension { thread_store.get::()?; let guardian_config = thread_store.get::()?; let score_progress = thread_store.get::()?; + let latest_scored_tool_call = score_progress + .latest_scored_tool_call + .load(Ordering::Acquire); let tool_call_lag = score_progress .latest_tool_call .load(Ordering::Acquire) - .saturating_sub( - score_progress - .latest_scored_tool_call - .load(Ordering::Acquire), - ); + .saturating_sub(latest_scored_tool_call); if let Some(metrics) = &extension_metrics { metrics.histogram( TOOL_CALL_LAG_METRIC, @@ -350,6 +350,14 @@ impl ApprovalReviewContributor for GuardianV2Extension { } return None; } + if score_progress + .latest_failed_tool_call + .load(Ordering::Acquire) + > latest_scored_tool_call + { + thread_store.insert(StrictReviewReason::ElevatedRisk); + return None; + } let score = thread_store .get::() @@ -367,17 +375,34 @@ impl ApprovalReviewContributor for GuardianV2Extension { impl ToolLifecycleContributor for GuardianV2Extension { fn on_tool_start<'a>(&'a self, input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> { + Box::pin(self.score_tool(input)) + } +} + +impl GuardianV2Extension { + fn record_fail_closed_score(thread_store: &ExtensionData, sampled_at: SystemTime) { + let score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), + sampled_at: Some(sampled_at.into()), + }; + thread_store.insert_if(score.clone(), |previous| { + previous.is_none_or(|previous| previous.sampled_at <= score.sampled_at) + }); + } + + async fn score_tool(&self, input: ToolStartInput<'_>) { let classification_started_at = Instant::now(); let Some(sampler) = input.thread_store.get::() else { - return Box::pin(std::future::ready(())); + return; }; let Some(guardian_config) = input.thread_store.get::() else { - return Box::pin(std::future::ready(())); + return; }; let Some(score_progress) = input.thread_store.get::() else { - return Box::pin(std::future::ready(())); + return; }; let metrics = score_progress.metrics.clone(); + let sampled_at = SystemTime::now(); let parent_model = input.thread_store.get::(); let model_defaults = parent_model .as_ref() @@ -386,6 +411,7 @@ impl ToolLifecycleContributor for GuardianV2Extension { let guardian_config = match guardian_config.with_model_defaults(model_defaults) { Ok(config) => config, Err(error) => { + Self::record_fail_closed_score(input.thread_store, sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -396,7 +422,7 @@ impl ToolLifecycleContributor for GuardianV2Extension { turn_id: Some(input.turn_id.to_owned()), message: error, }); - return Box::pin(std::future::ready(())); + return; } }; input.thread_store.insert(guardian_config.clone()); @@ -404,7 +430,6 @@ impl ToolLifecycleContributor for GuardianV2Extension { .latest_tool_call .fetch_add(/*val*/ 1, Ordering::Relaxed) .saturating_add(1); - let sampled_at = SystemTime::now(); let latest_parent_compaction = input .conversation_history .items() @@ -436,19 +461,13 @@ impl ToolLifecycleContributor for GuardianV2Extension { _ => false, }) { - let score = SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), - sampled_at: Some(sampled_at.into()), - }; - input.thread_store.insert_if(score.clone(), |previous| { - previous.is_none_or(|previous| previous.sampled_at <= score.sampled_at) - }); + Self::record_fail_closed_score(input.thread_store, sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), "failure", ); - return Box::pin(std::future::ready(())); + return; } let event_sink = Arc::clone(&self.event_sink); let thread_manager = self.thread_manager.clone(); @@ -476,6 +495,38 @@ impl ToolLifecycleContributor for GuardianV2Extension { }; tokio::spawn(async move { + let thread_context: Result<_, String> = async { + let parsed_thread_id = + ThreadId::from_string(&thread_id).map_err(|error| error.to_string())?; + let manager = thread_manager + .upgrade() + .ok_or_else(|| "thread manager is unavailable".to_string())?; + let thread = manager + .get_thread(parsed_thread_id) + .await + .map_err(|error| error.to_string())?; + Ok((manager, thread)) + } + .await; + let (manager, thread) = match thread_context { + Ok(context) => context, + Err(error) => { + score_progress + .latest_failed_tool_call + .fetch_max(tool_call_index, Ordering::Release); + record_classification( + metrics.as_deref(), + classification_started_at.elapsed(), + "failure", + ); + event_sink.emit_warning(ExtensionWarning { + thread_id, + turn_id: Some(turn_id), + message: format!("Guardian V2 risk scoring failed: {error}"), + }); + return; + } + }; let transcript = guardian_config .transcript .build(conversation_history.items()); @@ -486,6 +537,7 @@ impl ToolLifecycleContributor for GuardianV2Extension { let planned_action = match action.render(guardian_config.max_action_tokens) { Ok(planned_action) => planned_action, Err(error) => { + Self::record_fail_closed_score(thread.thread_extension_data(), sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -511,15 +563,6 @@ impl ToolLifecycleContributor for GuardianV2Extension { ]); let mut classification_finished_at = None; let result: Result<&str, String> = async { - let parsed_thread_id = - ThreadId::from_string(&thread_id).map_err(|error| error.to_string())?; - let manager = thread_manager - .upgrade() - .ok_or_else(|| "thread manager is unavailable".to_string())?; - let thread = manager - .get_thread(parsed_thread_id) - .await - .map_err(|error| error.to_string())?; let config = thread.config().await; let review_model_messages = if config.guardian_policy_config.is_none() { let review_model_id = review_model_override.as_deref().unwrap_or_else(|| { @@ -616,6 +659,9 @@ impl ToolLifecycleContributor for GuardianV2Extension { Ok("success") } .await; + if result.is_err() { + Self::record_fail_closed_score(thread.thread_extension_data(), sampled_at); + } record_classification( metrics.as_deref(), classification_finished_at @@ -633,8 +679,6 @@ impl ToolLifecycleContributor for GuardianV2Extension { }); } }); - - Box::pin(std::future::ready(())) } } diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index 8e529c1f37..bd5e29ec40 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/extension_tests.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use std::sync::Mutex; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::SystemTime; use anyhow::Result; use codex_core::config::Config; @@ -50,6 +51,7 @@ use serde_json::json; use super::CLASSIFICATION_DURATION_METRIC; use super::CLASSIFICATION_METRIC; +use super::GuardianV2Extension; use super::GuardianV2ScoreProgress; use super::REVIEW_FALLBACK_METRIC; use super::StrictReviewReason; @@ -223,6 +225,37 @@ impl ConversationHistorySnapshot for TestConversationHistory { } } +#[test] +fn fail_closed_score_preserves_classification_order() { + let thread_store = ExtensionData::new("thread-1"); + let newer_sampled_at = SystemTime::UNIX_EPOCH + Duration::from_secs(1); + let newest_sampled_at = newer_sampled_at + Duration::from_secs(1); + let newer_score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + sampled_at: Some(newer_sampled_at.into()), + }; + thread_store.insert(newer_score.clone()); + + GuardianV2Extension::record_fail_closed_score(&thread_store, SystemTime::UNIX_EPOCH); + assert_eq!( + thread_store.get::().as_deref(), + Some(&newer_score) + ); + + GuardianV2Extension::record_fail_closed_score(&thread_store, newest_sampled_at); + let fail_closed_score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), + sampled_at: Some(newest_sampled_at.into()), + }; + assert!(!thread_store.insert_if(newer_score.clone(), |previous| { + previous.is_none_or(|previous| previous.sampled_at < newer_score.sampled_at) + })); + assert_eq!( + thread_store.get::().as_deref(), + Some(&fail_closed_score) + ); +} + #[test] fn encrypted_parent_compaction_preserves_the_latest_valid_item() { let older = ResponseItem::Compaction { @@ -533,6 +566,186 @@ async fn sample_configured_conversation_history( Ok((request.body_json(), test, registry)) } +struct GuardianFailureFixture { + test: TestCodex, + registry: ExtensionRegistry, + session_store: ExtensionData, +} + +impl GuardianFailureFixture { + async fn new() -> Result { + let (_, test, registry) = sample_conversation_history( + Vec::new(), + r#"{"path":"README.md"}"#, + Some(TEST_GUARDIAN_POLICY), + ) + .await?; + let thread_store = test.codex.thread_extension_data(); + tokio::time::timeout(Duration::from_secs(5), async { + while thread_store.get::().is_none() { + tokio::task::yield_now().await; + } + }) + .await?; + + Ok(Self { + test, + registry, + session_store: ExtensionData::new("session-1"), + }) + } + + async fn score_tool(&self, tool_name: ToolName) { + let thread_store = self.test.codex.thread_extension_data(); + thread_store.insert(SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + sampled_at: None, + }); + let turn_store = ExtensionData::new("turn-1"); + let payload = ToolPayload::Function { + arguments: r#"{"path":"README.md"}"#.to_owned(), + }; + self.registry.tool_lifecycle_contributors()[0] + .on_tool_start(ToolStartInput { + session_store: &self.session_store, + thread_store, + turn_store: &turn_store, + turn_id: "turn-1", + call_id: "call-1", + tool_name: &tool_name, + payload: &payload, + conversation_history: Arc::new(TestConversationHistory(Vec::new())), + source: ToolCallSource::Direct, + }) + .await; + } + + async fn assert_fails_closed(&self) -> Result<()> { + let thread_store = self.test.codex.thread_extension_data(); + let score_progress = thread_store + .get::() + .expect("Guardian v2 should track score progress per thread"); + tokio::time::timeout(Duration::from_secs(5), async { + while thread_store + .get::() + .is_none_or(|score| score.scores.get("action_risk") != Some(&1.0)) + && score_progress + .latest_failed_tool_call + .load(Ordering::Acquire) + <= score_progress + .latest_scored_tool_call + .load(Ordering::Acquire) + { + tokio::task::yield_now().await; + } + }) + .await?; + assert_eq!( + self.registry + .approval_review( + &self.session_store, + thread_store, + "review action", + /*extension_metrics*/ None, + ) + .await, + None + ); + Ok(()) + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_fails_closed_when_thread_lookup_fails() -> Result<()> { + skip_if_no_network!(Ok(())); + + let fixture = GuardianFailureFixture::new().await?; + fixture + .test + .thread_manager + .remove_thread(&fixture.test.session_configured.thread_id) + .await + .expect("the test thread should exist before simulating a failed lookup"); + + fixture.score_tool(ToolName::plain("read_file")).await; + fixture.assert_fails_closed().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_fails_closed_when_model_configuration_is_invalid() -> Result<()> { + skip_if_no_network!(Ok(())); + + let fixture = GuardianFailureFixture::new().await?; + let mut parent_model = fixture + .test + .thread_manager + .get_models_manager() + .get_model_info("gpt-5.5", &fixture.test.config.to_models_manager_config()) + .await; + parent_model + .model_messages + .as_mut() + .expect("test model should expose model messages") + .guardian_v2 = Some(GuardianV2ModelConfig { + max_action_tokens: Some(1), + ..Default::default() + }); + fixture + .test + .codex + .thread_extension_data() + .insert(parent_model); + + fixture.score_tool(ToolName::plain("read_file")).await; + fixture.assert_fails_closed().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_fails_closed_when_action_serialization_fails() -> Result<()> { + skip_if_no_network!(Ok(())); + + let fixture = GuardianFailureFixture::new().await?; + let oversized_tool_name = ToolName::plain( + "x".repeat(TruncationPolicy::Tokens(DEFAULT_MODEL_CONTEXT_ITEM_TOKENS).byte_budget()), + ); + + fixture.score_tool(oversized_tool_name).await; + fixture.assert_fails_closed().await +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn contributor_fails_closed_when_luna_classification_fails() -> Result<()> { + skip_if_no_network!(Ok(())); + + let fixture = GuardianFailureFixture::new().await?; + let invalid_score = vec![ + ev_assistant_message("sample", r#"{"scores":{"action_risk":"invalid"}}"#), + ev_completed("response-invalid"), + ]; + let server = responses::start_websocket_server(vec![Vec::new(), vec![invalid_score]]).await; + let mut config = fixture.test.config.clone(); + config.model_provider = ModelProviderInfo::create_openai_provider(Some(format!( + "http://{}/v1", + server.uri().trim_start_matches("ws://") + ))); + config.features.enable(Feature::GuardianV2)?; + fixture.registry.thread_lifecycle_contributors()[0] + .on_thread_start(ThreadStartInput { + config: &config, + session_source: &SessionSource::Exec, + persistent_thread_state_available: false, + environments: &[], + mcp_resource_client: None, + extension_metrics: None, + session_store: &fixture.session_store, + thread_store: fixture.test.codex.thread_extension_data(), + }) + .await; + + fixture.score_tool(ToolName::plain("read_file")).await; + fixture.assert_fails_closed().await +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn contributor_renders_policy_inside_a_configured_prompt() -> Result<()> { skip_if_no_network!(Ok(()));