From 8630bb3caecaff6abc6add450a88035d9f6d3f8c Mon Sep 17 00:00:00 2001 From: jif Date: Fri, 14 Aug 2026 09:22:23 +0000 Subject: [PATCH] Record Guardian V2 risk scores on threads (#38540) ## What changed - Parse and validate the structured risk scores returned for each classified tool call. - Publish each score through thread extension data and persist it in rollout history for non-ephemeral threads. - Report parsing, thread lookup, validation, and persistence failures through the extension warning sink. ## Testing - Extend the Guardian V2 tool-classification test to verify that the sampled `action_risk` score is available from thread extension data and written to rollout history. GitOrigin-RevId: 46614fe0d686736dfb1f09398c6fcd848c7d8b2c --- codex-rs/Cargo.lock | 2 + codex-rs/ext/guardian-v2/Cargo.toml | 2 + codex-rs/ext/guardian-v2/src/extension.rs | 110 +++++++++++++----- .../ext/guardian-v2/src/extension_tests.rs | 53 ++++++++- 4 files changed, 132 insertions(+), 35 deletions(-) diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index 13524e63eb..63234ecaff 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -3304,7 +3304,9 @@ version = "0.0.0" dependencies = [ "anyhow", "codex-api", + "codex-core", "codex-extension-api", + "codex-history", "codex-http-client", "codex-login", "codex-model-provider", diff --git a/codex-rs/ext/guardian-v2/Cargo.toml b/codex-rs/ext/guardian-v2/Cargo.toml index a570952a58..2cbd2d9dca 100644 --- a/codex-rs/ext/guardian-v2/Cargo.toml +++ b/codex-rs/ext/guardian-v2/Cargo.toml @@ -14,7 +14,9 @@ workspace = true [dependencies] codex-api = { workspace = true } +codex-core = { workspace = true } codex-extension-api = { workspace = true } +codex-history = { workspace = true } codex-http-client = { workspace = true } codex-login = { workspace = true } codex-model-provider = { workspace = true } diff --git a/codex-rs/ext/guardian-v2/src/extension.rs b/codex-rs/ext/guardian-v2/src/extension.rs index 66fbbae0e8..49f1d0ae15 100644 --- a/codex-rs/ext/guardian-v2/src/extension.rs +++ b/codex-rs/ext/guardian-v2/src/extension.rs @@ -1,5 +1,7 @@ use std::sync::Arc; +use std::sync::Weak; +use codex_core::ThreadManager; use codex_extension_api::ExtensionEventSink; use codex_extension_api::ExtensionRegistryBuilder; use codex_extension_api::ExtensionWarning; @@ -7,7 +9,10 @@ use codex_extension_api::ToolLifecycleContributor; use codex_extension_api::ToolLifecycleFuture; use codex_extension_api::ToolPayload; use codex_extension_api::ToolStartInput; +use codex_history::RolloutItem; +use codex_protocol::ThreadId; use codex_protocol::openai_models::ReasoningEffort; +use codex_protocol::security_risk::SecurityRiskScore; use serde_json::json; use crate::LunaSampler; @@ -23,12 +28,14 @@ Return an action_risk score from 0.0 (safe and authorized) to 1.0 (dangerous or struct GuardianV2Extension { sampler: Arc, event_sink: Arc, + thread_manager: Weak, } impl ToolLifecycleContributor for GuardianV2Extension { fn on_tool_start<'a>(&'a self, input: ToolStartInput<'a>) -> ToolLifecycleFuture<'a> { let sampler = Arc::clone(&self.sampler); let event_sink = Arc::clone(&self.event_sink); + let thread_manager = self.thread_manager.clone(); let thread_id = input.thread_store.level_id().to_owned(); let turn_id = input.turn_id.to_owned(); let tool_name = input.tool_name.to_string(); @@ -69,38 +76,76 @@ impl ToolLifecycleContributor for GuardianV2Extension { {planned_action}\n\ >>> APPROVAL REQUEST END\n" ); - if let Err(error) = sampler - .sample(LunaSamplingRequest { - instructions: CLASSIFIER_INSTRUCTIONS.to_owned(), - input: classification_input, - output_schema: json!({ - "type": "object", - "properties": { - "scores": { - "type": "object", - "properties": { - "action_risk": { - "type": "number", - "minimum": 0.0, - "maximum": 1.0 - } - }, - "required": ["action_risk"], - "additionalProperties": false - } - }, - "required": ["scores"], - "additionalProperties": false - }), - reasoning_effort: ReasoningEffort::Low, - turn_id: turn_id.clone(), - }) - .await - { + let result: Result<(), String> = async { + let output = sampler + .sample(LunaSamplingRequest { + instructions: CLASSIFIER_INSTRUCTIONS.to_owned(), + input: classification_input, + output_schema: json!({ + "type": "object", + "properties": { + "scores": { + "type": "object", + "properties": { + "action_risk": { + "type": "number", + "minimum": 0.0, + "maximum": 1.0 + } + }, + "required": ["action_risk"], + "additionalProperties": false + } + }, + "required": ["scores"], + "additionalProperties": false + }), + reasoning_effort: ReasoningEffort::Low, + turn_id: turn_id.clone(), + }) + .await + .map_err(|error| error.to_string())?; + let output: serde_json::Value = + serde_json::from_str(&output).map_err(|error| error.to_string())?; + let scores = output + .get("scores") + .and_then(serde_json::Value::as_object) + .ok_or_else(|| "Luna returned no security risk scores".to_string())?; + 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 ephemeral = thread.config_snapshot().await.ephemeral; + for (category, value) in scores { + let score = value + .as_f64() + .filter(|score| (0.0..=1.0).contains(score)) + .ok_or_else(|| format!("invalid security risk score for {category}"))?; + let score = SecurityRiskScore { + category: category.clone(), + score, + }; + if !ephemeral { + thread + .append_rollout_items(&[RolloutItem::SecurityRiskScore(score.clone())]) + .await + .map_err(|error| error.to_string())?; + } + thread.thread_extension_data().insert(score); + } + Ok(()) + } + .await; + if let Err(error) = result { event_sink.emit_warning(ExtensionWarning { thread_id, turn_id: Some(turn_id), - message: format!("Guardian V2 Luna sampling failed: {error}"), + message: format!("Guardian V2 risk scoring failed: {error}"), }); } }); @@ -110,10 +155,15 @@ impl ToolLifecycleContributor for GuardianV2Extension { } /// Installs Guardian V2 tool classification over a caller-owned Luna sampler. -pub fn install(registry: &mut ExtensionRegistryBuilder, sampler: Arc) { +pub fn install( + registry: &mut ExtensionRegistryBuilder, + sampler: Arc, + thread_manager: Weak, +) { registry.tool_lifecycle_contributor(Arc::new(GuardianV2Extension { sampler, event_sink: registry.event_sink(), + thread_manager, })); } diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index 48a4adb023..9b0bbc675e 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/extension_tests.rs @@ -10,6 +10,7 @@ use codex_extension_api::ToolCallSource; use codex_extension_api::ToolName; use codex_extension_api::ToolPayload; use codex_extension_api::ToolStartInput; +use codex_history::RolloutItem; use codex_http_client::HttpClientFactory; use codex_http_client::OutboundProxyPolicy; use codex_login::AgentIdentityAuthPolicy; @@ -21,10 +22,12 @@ use codex_protocol::models::ContentItem; use codex_protocol::models::FunctionCallOutputPayload; use codex_protocol::models::ReasoningItemReasoningSummary; use codex_protocol::protocol::SessionSource; +use codex_protocol::security_risk::SecurityRiskScore; use core_test_support::responses; use core_test_support::responses::ev_assistant_message; use core_test_support::responses::ev_completed; use core_test_support::skip_if_no_network; +use core_test_support::test_codex::test_codex; use pretty_assertions::assert_eq; use serde_json::json; @@ -43,6 +46,9 @@ impl ConversationHistorySnapshot for TestConversationHistory { async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<()> { skip_if_no_network!(Ok(())); + let thread_server = responses::start_mock_server().await; + let test = test_codex().build_with_auto_env(&thread_server).await?; + let thread_id = test.session_configured.thread_id; let events = vec![ ev_assistant_message("sample", r#"{"scores":{"action_risk":0.25}}"#), ev_completed("response-1"), @@ -63,16 +69,20 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result< agent_identity_policy: AgentIdentityAuthPolicy::JwtOnly, session_source: SessionSource::Exec, session_id: "session-1".to_owned(), - thread_id: "thread-1".to_owned(), + thread_id: thread_id.to_string(), originator: None, service_tier: None, }) .await?; let mut builder = ExtensionRegistryBuilder::<()>::new(); - crate::install(&mut builder, Arc::new(sampler)); + crate::install( + &mut builder, + Arc::new(sampler), + Arc::downgrade(&test.thread_manager), + ); let registry = builder.build(); let session_store = ExtensionData::new("session-1"); - let thread_store = ExtensionData::new("thread-1"); + let thread_store = test.codex.thread_extension_data(); let turn_store = ExtensionData::new("turn-1"); let tool_name = ToolName::plain("read_file"); let tool_payload = ToolPayload::Function { @@ -126,7 +136,7 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result< registry.tool_lifecycle_contributors()[0] .on_tool_start(ToolStartInput { session_store: &session_store, - thread_store: &thread_store, + thread_store, turn_store: &turn_store, turn_id: "turn-1", call_id: "call-1", @@ -144,7 +154,10 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result< .await?; let request = request.body_json(); assert_eq!(request["model"], "gpt-5.6-luna"); - assert_eq!(request["client_metadata"]["thread_id"], "thread-1"); + assert_eq!( + request["client_metadata"]["thread_id"], + thread_id.to_string() + ); assert_eq!(request["client_metadata"]["turn_id"], "turn-1"); assert_eq!(request["reasoning"]["effort"], "low"); assert_eq!(request["reasoning"]["context"], "all_turns"); @@ -173,6 +186,36 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result< ">>> APPROVAL REQUEST END\n", ) ); + let score = tokio::time::timeout(Duration::from_secs(5), async { + loop { + if let Some(score) = thread_store.get::() { + return score; + } + tokio::task::yield_now().await; + } + }) + .await?; + assert_eq!( + score.as_ref(), + &SecurityRiskScore { + category: "action_risk".to_string(), + score: 0.25, + } + ); + test.codex.ensure_rollout_materialized().await; + test.codex.flush_rollout().await?; + let persisted_scores = test + .codex + .load_history(/*include_archived*/ false) + .await? + .items + .into_iter() + .filter_map(|item| match item { + RolloutItem::SecurityRiskScore(score) => Some(score), + _ => None, + }) + .collect::>(); + assert_eq!(persisted_scores, vec![score.as_ref().clone()]); Ok(()) }