openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-11-05 20:10:33 +00:00
Code Issues Packages Projects Releases Wiki Activity

Record Guardian V2 risk scores on threads (#38540)

Browse Source 
## 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
This commit is contained in:
jif
2026-08-14 09:27:41 +00:00
committed by copyberry
parent 4eff3b788b
commit 8630bb3cae
4 changed files with 132 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/Cargo.lock
Generated
+2
View File
@@ -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
+2
View File
@@ -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
+80 -30
View File
@@ -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
.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<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
+48 -5
View File
@@ -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(())
}