Keep the latest Guardian risk score during concurrent sampling (#38580)

## Why

Concurrent Guardian samples can finish out of order, allowing an older sample to
replace the thread's newer in-memory security risk score.

## What changed

- Record when each `SecurityRiskScore` sample starts while preserving compatibility
  with records that have no timestamp.
- Add `ExtensionData::insert_if` so checking and replacing a typed attachment happen
  under the same lock.
- Update Guardian V2 to replace the thread's risk score only when the completed sample
  is newer than the stored one.

## Testing

- Verify concurrent conditional inserts retain the newest value.
- Verify Guardian-generated risk scores include a sampling timestamp.

GitOrigin-RevId: 6543a4e5497da4c623331050789e91613f931c4b
This commit is contained in:
jif authored and copyberry committed 2026-08-14 13:44:51 +00:00
1 parent 42d842a91a
commit 2a452d7dc1
11 files changed
+76 -2

No files matched your search

+11 -2
View File
@@ -1,5 +1,6 @@
use std::sync::Arc;
use std::sync::Weak;
use std::time::SystemTime;
use codex_core::ThreadManager;
use codex_core::config::Config;
@@ -133,6 +134,7 @@ impl ToolLifecycleContributor for GuardianV2Extension {
let Some(sampler) = input.thread_store.get::<LunaSampler>() else {
return Box::pin(std::future::ready(()));
};
let sampled_at = SystemTime::now();
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();
@@ -239,8 +241,15 @@ impl ToolLifecycleContributor for GuardianV2Extension {
.ok_or_else(|| format!("invalid security risk score for {category}"))
})
.collect::<Result<_, _>>()?;
let score = SecurityRiskScore { scores };
thread.thread_extension_data().insert(score.clone());
let score = SecurityRiskScore {
scores,
sampled_at: Some(sampled_at.into()),
};
thread
.thread_extension_data()
.insert_if(score.clone(), |previous| {
previous.is_none_or(|previous| previous.sampled_at < score.sampled_at)
});
if !ephemeral {
thread
.append_rollout_items(&[RolloutItem::SecurityRiskScore(score)])
@@ -284,8 +284,10 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
score.as_ref(),
&SecurityRiskScore {
scores: BTreeMap::from([("action_risk".to_string(), 0.8)]),
sampled_at: score.sampled_at,
}
);
assert!(score.sampled_at.is_some());
assert_eq!(
registry.approval_requirement(thread_store),
ApprovalRequirement::RequireAutomaticReview
@@ -307,6 +309,7 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
thread_store.insert(SecurityRiskScore {
scores: BTreeMap::from([("action_risk".to_string(), 0.25)]),
sampled_at: None,
});
assert_eq!(
registry.approval_requirement(thread_store),
@@ -316,6 +319,7 @@ async fn contributor_samples_tool_calls_with_the_existing_luna_pool() -> Result<
let disabled_thread_store = ExtensionData::new("disabled-thread");
disabled_thread_store.insert(SecurityRiskScore {
scores: BTreeMap::from([("action_risk".to_string(), 0.8)]),
sampled_at: None,
});
assert_eq!(
registry.approval_requirement(&disabled_thread_store),