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

@@ -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",

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 }

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<LunaSampler>,
     event_sink: Arc<dyn ExtensionEventSink>,
+    thread_manager: Weak<ThreadManager>,
 }
 
 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
+            let result: Result<(), String> = async {
-                .sample(LunaSamplingRequest {
+                let output = sampler
-                    instructions: CLASSIFIER_INSTRUCTIONS.to_owned(),
+                    .sample(LunaSamplingRequest {
-                    input: classification_input,
+                        instructions: CLASSIFIER_INSTRUCTIONS.to_owned(),
-                    output_schema: json!({
+                        input: classification_input,
-                        "type": "object",
+                        output_schema: json!({
-                        "properties": {
+                            "type": "object",
-                            "scores": {
+                            "properties": {
-                                "type": "object",
+                                "scores": {
-                                "properties": {
+                                    "type": "object",
-                                    "action_risk": {
+                                    "properties": {
-                                        "type": "number",
+                                        "action_risk": {
-                                        "minimum": 0.0,
+                                            "type": "number",
-                                        "maximum": 1.0
+                                            "minimum": 0.0,
-                                    }
+                                            "maximum": 1.0
-                                },
+                                        }
-                                "required": ["action_risk"],
+                                    },
-                                "additionalProperties": false
+                                    "required": ["action_risk"],
-                            }
+                                    "additionalProperties": false
-                        },
+                                }
-                        "required": ["scores"],
+                            },
-                        "additionalProperties": false
+                            "required": ["scores"],
-                    }),
+                            "additionalProperties": false
-                    reasoning_effort: ReasoningEffort::Low,
+                        }),
-                    turn_id: turn_id.clone(),
+                        reasoning_effort: ReasoningEffort::Low,
-                })
+                        turn_id: turn_id.clone(),
-                .await
+                    })
-            {
+                    .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<C: Sync>(registry: &mut ExtensionRegistryBuilder<C>, sampler: Arc<LunaSampler>) {
+pub fn install<C: Sync>(
+    registry: &mut ExtensionRegistryBuilder<C>,
+    sampler: Arc<LunaSampler>,
+    thread_manager: Weak<ThreadManager>,
+) {
     registry.tool_lifecycle_contributor(Arc::new(GuardianV2Extension {
         sampler,
         event_sink: registry.event_sink(),
+        thread_manager,
     }));
 }
 

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::<SecurityRiskScore>() {
+                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::<Vec<_>>();
+    assert_eq!(persisted_scores, vec![score.as_ref().clone()]);
 
     Ok(())
 }