diff --git a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst index 5e41454f63..a78af9ad45 100644 Binary files a/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst and b/codex-rs/app-server-protocol/schema/precomputed/app-server-exports-stable.json.zst differ diff --git a/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs b/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs index 3c3d4e2ef7..8f7e8e466a 100644 --- a/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs +++ b/codex-rs/app-server-protocol/src/protocol/thread_history_projection_tests.rs @@ -204,6 +204,7 @@ fn ignores_legacy_abort_without_turn_id_and_context_only_records() { })); let security_risk = project(RolloutItem::SecurityRiskScore(SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_string(), 0.92)]), + sampled_at: None, })); assert!(aborted.is_empty()); diff --git a/codex-rs/core/src/session/rollout_reconstruction_tests.rs b/codex-rs/core/src/session/rollout_reconstruction_tests.rs index a1cbeafee9..e581569dfe 100644 --- a/codex-rs/core/src/session/rollout_reconstruction_tests.rs +++ b/codex-rs/core/src/session/rollout_reconstruction_tests.rs @@ -154,6 +154,7 @@ async fn record_initial_history_ignores_security_risk_scores() { let user_item = user_message("visible user input"); let security_risk = SecurityRiskScore { scores: BTreeMap::from([("credential_access".to_string(), 0.92)]), + sampled_at: None, }; session diff --git a/codex-rs/ext/extension-api/src/state.rs b/codex-rs/ext/extension-api/src/state.rs index d2a55534c4..7d6e79dc58 100644 --- a/codex-rs/ext/extension-api/src/state.rs +++ b/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(&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::()) + .map(|value| downcast_data::(Arc::clone(value))); + if !should_insert(existing.as_deref()) { + return false; + } + entries.insert(TypeId::of::(), Arc::new(value)); + true + } + /// Removes and returns the attached value of type `T`, if one exists. pub fn remove(&self) -> Option> where diff --git a/codex-rs/ext/extension-api/tests/state.rs b/codex-rs/ext/extension-api/tests/state.rs index 31b643c29b..6e7720c913 100644 --- a/codex-rs/ext/extension-api/tests/state.rs +++ b/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::().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::>(); + + for handle in handles { + handle.join().expect("insertion thread should succeed"); + } + + assert_eq!(data.get::().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::().as_deref(), Some(&(CALLER_COUNT - 1))); +} + #[test] fn get_or_init_initializes_once_and_returns_shared_value() { const CALLER_COUNT: usize = 8; diff --git a/codex-rs/ext/guardian-v2/src/extension.rs b/codex-rs/ext/guardian-v2/src/extension.rs index 915ab266f2..836db2bd1b 100644 --- a/codex-rs/ext/guardian-v2/src/extension.rs +++ b/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::() 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::>()?; - 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)]) diff --git a/codex-rs/ext/guardian-v2/src/extension_tests.rs b/codex-rs/ext/guardian-v2/src/extension_tests.rs index a1eaa03745..c3eaee3e8f 100644 --- a/codex-rs/ext/guardian-v2/src/extension_tests.rs +++ b/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), diff --git a/codex-rs/external-agent-migration/src/sessions/append_tests.rs b/codex-rs/external-agent-migration/src/sessions/append_tests.rs index 23de57961d..17d1c48b05 100644 --- a/codex-rs/external-agent-migration/src/sessions/append_tests.rs +++ b/codex-rs/external-agent-migration/src/sessions/append_tests.rs @@ -66,6 +66,7 @@ fn requires_a_strict_nonempty_model_prefix() { } let security_risk = RolloutItem::SecurityRiskScore(SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_string(), 0.92)]), + sampled_at: None, }); metadata_changed.push(security_risk.clone()); assert!(model_transcripts_match(&history, &metadata_changed)); diff --git a/codex-rs/memories/write/src/phase1.rs b/codex-rs/memories/write/src/phase1.rs index 954eb6b3a8..820b59030f 100644 --- a/codex-rs/memories/write/src/phase1.rs +++ b/codex-rs/memories/write/src/phase1.rs @@ -727,6 +727,7 @@ mod tests { RolloutItem::ResponseItem(skill_message.into()), RolloutItem::SecurityRiskScore(SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_string(), 0.92)]), + sampled_at: None, }), RolloutItem::ResponseItem(subagent_message.clone().into()), ]) diff --git a/codex-rs/protocol/src/security_risk.rs b/codex-rs/protocol/src/security_risk.rs index 70c9a83d64..7d1dde8aab 100644 --- a/codex-rs/protocol/src/security_risk.rs +++ b/codex-rs/protocol/src/security_risk.rs @@ -1,3 +1,5 @@ +use chrono::DateTime; +use chrono::Utc; use schemars::JsonSchema; use serde::Deserialize; use serde::Serialize; @@ -10,4 +12,8 @@ use std::collections::BTreeMap; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct SecurityRiskScore { pub scores: BTreeMap, + /// When sampling started, if this snapshot was written by a timestamp-aware client. + #[serde(default, skip_serializing_if = "Option::is_none")] + #[schemars(with = "Option")] + pub sampled_at: Option>, } diff --git a/codex-rs/rollout/src/recorder_tests.rs b/codex-rs/rollout/src/recorder_tests.rs index a6f1237a89..a69919dec5 100644 --- a/codex-rs/rollout/src/recorder_tests.rs +++ b/codex-rs/rollout/src/recorder_tests.rs @@ -443,6 +443,7 @@ async fn load_rollout_items_preserves_security_risk_scores() -> std::io::Result< ("action_risk".to_string(), 0.76), ("data_exfiltration".to_string(), 0.31), ]), + sampled_at: None, }; let security_risk_item = RolloutItem::SecurityRiskScore(security_risk.clone()); for history_mode in [ThreadHistoryMode::Legacy, ThreadHistoryMode::Paginated] {