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

codex-rs/ext/extension-api/src/state.rs

@@ -104,6 +104,25 @@ impl ExtensionData {
             .map(downcast_data)
     }
 
+    /// Stores `value` only when `should_insert` accepts the current attachment.
+    ///
+    /// The predicate and insertion happen while this map is locked, so concurrent
+    /// callers cannot replace a value after checking a stale attachment.
+    pub fn insert_if<T>(&self, value: T, should_insert: impl FnOnce(Option<&T>) -> bool) -> bool
+    where
+        T: Any + Send + Sync,
+    {
+        let mut entries = self.entries();
+        let existing = entries
+            .get(&TypeId::of::<T>())
+            .map(|value| downcast_data::<T>(Arc::clone(value)));
+        if !should_insert(existing.as_deref()) {
+            return false;
+        }
+        entries.insert(TypeId::of::<T>(), Arc::new(value));
+        true
+    }
+
     /// Removes and returns the attached value of type `T`, if one exists.
     pub fn remove<T>(&self) -> Option<Arc<T>>
     where

codex-rs/ext/extension-api/tests/state.rs

@@ -1,5 +1,6 @@
 use std::panic::AssertUnwindSafe;
 use std::sync::Arc;
+use std::sync::Barrier;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 
@@ -29,6 +30,36 @@ fn typed_values_can_be_inserted_replaced_and_removed() {
     assert_eq!(data.get::<u64>().as_deref(), Some(&42));
 }
 
+#[test]
+fn conditional_insert_keeps_the_newest_concurrent_value() {
+    const CALLER_COUNT: u64 = 16;
+
+    let data = Arc::new(ExtensionData::new("thread-1"));
+    let callers_ready = Arc::new(Barrier::new(CALLER_COUNT as usize));
+    let handles = (0..CALLER_COUNT)
+        .map(|value| {
+            let data = Arc::clone(&data);
+            let callers_ready = Arc::clone(&callers_ready);
+            std::thread::spawn(move || {
+                callers_ready.wait();
+                data.insert_if(value, |existing| {
+                    existing.is_none_or(|existing| value > *existing)
+                });
+            })
+        })
+        .collect::<Vec<_>>();
+
+    for handle in handles {
+        handle.join().expect("insertion thread should succeed");
+    }
+
+    assert_eq!(data.get::<u64>().as_deref(), Some(&(CALLER_COUNT - 1)));
+    assert!(!data.insert_if(/*value*/ 0_u64, |existing| {
+        existing.is_none_or(|existing| *existing == 0)
+    }));
+    assert_eq!(data.get::<u64>().as_deref(), Some(&(CALLER_COUNT - 1)));
+}
+
 #[test]
 fn get_or_init_initializes_once_and_returns_shared_value() {
     const CALLER_COUNT: usize = 8;

codex-rs/ext/guardian-v2/src/extension.rs

@@ -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)])

codex-rs/ext/guardian-v2/src/extension_tests.rs

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