diff --git a/codex-rs/core/tests/common/test_codex.rs b/codex-rs/core/tests/common/test_codex.rs index 3c837e96d8..2880e21b6d 100644 --- a/codex-rs/core/tests/common/test_codex.rs +++ b/codex-rs/core/tests/common/test_codex.rs @@ -742,7 +742,11 @@ impl TestCodexBuilder { .or_else(|| codex_utils_cargo_bin::cargo_bin("codex-code-mode-host").ok()); let thread_manager = Arc::new_cyclic(|manager| { let mut extensions = self.extensions.to_builder(); - codex_guardian_v2::install_reviewer(&mut extensions, manager.clone()); + if config.features.enabled(Feature::GuardianV2) { + codex_guardian_v2::install(&mut extensions, auth_manager.clone(), manager.clone()); + } else { + codex_guardian_v2::install_reviewer(&mut extensions, manager.clone()); + } let thread_manager = ThreadManager::new( &config, auth_manager.clone(), diff --git a/codex-rs/core/tests/suite/guardian_cached_score_tests.rs b/codex-rs/core/tests/suite/guardian_cached_score_tests.rs new file mode 100644 index 0000000000..83b470aba2 --- /dev/null +++ b/codex-rs/core/tests/suite/guardian_cached_score_tests.rs @@ -0,0 +1,289 @@ +//! Exercises asynchronous score publication through real tool calls and approval results. + +use std::time::Duration; + +use anyhow::Result; +use codex_core::TurnInputRequest; +use codex_core::config::Constrained; +use codex_features::Feature; +use codex_history::RolloutItem; +use codex_protocol::approvals::GuardianAssessmentStatus; +use codex_protocol::approvals::GuardianReviewReason; +use codex_protocol::config_types::ApprovalsReviewer; +use codex_protocol::models::NetworkPermissions; +use codex_protocol::models::PermissionProfile; +use codex_protocol::protocol::AskForApproval; +use codex_protocol::protocol::EventMsg; +use codex_protocol::request_permissions::PermissionGrantScope; +use codex_protocol::request_permissions::RequestPermissionProfile; +use codex_protocol::request_permissions::RequestPermissionsResponse; +use codex_protocol::user_input::UserInput; +use core_test_support::responses; +use core_test_support::responses::ev_assistant_message; +use core_test_support::responses::ev_completed; +use core_test_support::responses::ev_function_call; +use core_test_support::responses::ev_response_created; +use core_test_support::responses::sse; +use core_test_support::skip_if_no_network; +use core_test_support::streaming_sse::StreamingSseChunk; +use core_test_support::streaming_sse::start_streaming_sse_server; +use core_test_support::test_codex::test_codex; +use core_test_support::wait_for_event; +use pretty_assertions::assert_eq; +use serde_json::json; +use tokio::sync::oneshot; +use tokio::time::timeout; +use wiremock::Mock; +use wiremock::ResponseTemplate; +use wiremock::matchers::body_partial_json; +use wiremock::matchers::method; +use wiremock::matchers::path; + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn delayed_score_only_covers_the_tool_call_it_classified() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let permission_args = json!({ + "reason": "request network access", + "permissions": { "network": { "enabled": true } }, + }) + .to_string(); + let mut parent_gates = Vec::new(); + let mut classifier_gates = Vec::new(); + let mut parent_responses = Vec::new(); + let mut classifier_responses = Vec::new(); + for (call_id, tool, args) in [ + ( + "before-publication", + "request_permissions", + permission_args.clone(), + ), + ( + "intervening-call", + "update_plan", + json!({ "plan": [{ "step": "Inspect permissions", "status": "in_progress" }] }) + .to_string(), + ), + ( + "stale-publication", + "request_permissions", + permission_args.clone(), + ), + ("fresh-publication", "request_permissions", permission_args), + ] { + let (parent_tx, parent_rx) = oneshot::channel(); + parent_gates.push(parent_tx); + parent_responses.push(vec![StreamingSseChunk { + gate: Some(parent_rx), + body: sse(vec![ + ev_response_created(call_id), + ev_function_call(call_id, tool, &args), + ev_completed(call_id), + ]), + }]); + let (classifier_tx, classifier_rx) = oneshot::channel(); + classifier_gates.push(classifier_tx); + classifier_responses.push(vec![StreamingSseChunk { + gate: Some(classifier_rx), + body: sse(vec![ + ev_response_created(call_id), + ev_assistant_message(call_id, "low"), + ev_completed(call_id), + ]), + }]); + } + let (parent_server, _) = start_streaming_sse_server(parent_responses).await; + let (classifier_server, _) = start_streaming_sse_server(classifier_responses).await; + + // Route the concurrent parent and classifier requests to independent gated streams. + for (model, destination) in [ + ("guardian-publication-parent", parent_server.uri()), + ("gpt-5.6-luna", classifier_server.uri()), + ] { + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(body_partial_json(json!({ "model": model }))) + .respond_with( + ResponseTemplate::new(/*s*/ 307) + .insert_header("location", format!("{destination}/v1/responses")), + ) + .with_priority(/*p*/ 1) + .up_to_n_times(/*n*/ 4) + .mount(&server) + .await; + } + let final_response = responses::mount_sse_once_match( + &server, + body_partial_json(json!({ "model": "guardian-publication-parent" })), + sse(vec![ + ev_response_created("done"), + ev_assistant_message("done", "done"), + ev_completed("done"), + ]), + ) + .await; + Mock::given(method("POST")) + .and(path("/v1/responses")) + .and(body_partial_json(json!({ + "model": "gpt-5.5", + "client_metadata": { "x-openai-subagent": "guardian" }, + }))) + .respond_with(responses::sse_response(sse(vec![ + ev_response_created("review"), + ev_assistant_message( + "review", + &json!({ + "risk_level": "high", + "user_authorization": "low", + "outcome": "deny", + "rationale": "Network access is not authorized.", + }) + .to_string(), + ), + ev_completed("review"), + ]))) + .with_priority(/*p*/ 2) + .expect(/*r*/ 2) + .mount(&server) + .await; + + let test = test_codex() + .with_pre_build_hook(|home| { + std::fs::write( + home.join("config.toml"), + "[features.guardianv2]\nenabled = true\npersist_scores = true\nmax_tool_call_lag = 1\n\n[features.guardianv2.review_scope]\ncomputer_use_only = false\n", + ) + .expect("write Guardian configuration"); + }) + .with_model_info_override("guardian-publication-parent", |model| { + model.guardian = None; + model.auto_review_model_override = Some("gpt-5.5".to_owned()); + }) + .with_config(|config| { + config.update_plan_enabled = true; + config.approvals_reviewer = ApprovalsReviewer::AutoReview; + config.permissions.approval_policy = Constrained::allow_any(AskForApproval::OnRequest); + config + .permissions + .set_permission_profile(PermissionProfile::read_only()) + .expect("set read-only permissions"); + config + .features + .enable(Feature::RequestPermissionsTool) + .expect("enable permission requests"); + }) + .build_with_auto_env(&server) + .await?; + test.codex.ensure_rollout_materialized().await; + test.codex + .start_or_steer_turn(TurnInputRequest::user_input(vec![UserInput::Text { + text: "Inspect the available permissions.".to_owned(), + text_elements: Vec::new(), + }])) + .await?; + + let mut parent_gates = parent_gates.into_iter(); + let mut classifier_gates = classifier_gates.into_iter(); + parent_gates.next().unwrap().send(()).unwrap(); + timeout(Duration::from_secs(30), async { + classifier_server.wait_for_request_count(/*count*/ 1).await; + parent_server.wait_for_request_count(/*count*/ 2).await; + }) + .await?; + parent_gates.next().unwrap().send(()).unwrap(); + timeout(Duration::from_secs(30), async { + classifier_server.wait_for_request_count(/*count*/ 2).await; + parent_server.wait_for_request_count(/*count*/ 3).await; + }) + .await?; + + // Publishing call 1 after call 2 starts must not claim coverage through call 2. + // At call 3 it is two calls behind, exceeding the configured lag of one. + let first_score = classifier_gates.next().unwrap(); + let _held_intervening_score = classifier_gates.next().unwrap(); + let current_score = classifier_gates.next().unwrap(); + let _held_final_score = classifier_gates.next().unwrap(); + let publish_score = async |call_id: &str, release_score: oneshot::Sender<()>| -> Result<()> { + release_score.send(()).unwrap(); + // Persisted rollout scores are emitted only after publication. Observe that + // public result instead of reading or modifying the private score state. + timeout(Duration::from_secs(30), async { + loop { + let history = test.codex.load_history(/*include_archived*/ false).await?; + if history.items.into_iter().any( + |item| matches!(item, RolloutItem::SecurityRiskScore(score) if score.call_id.as_deref() == Some(call_id)), + ) { + return Ok::<_, anyhow::Error>(()); + } + tokio::task::yield_now().await; + } + }) + .await??; + Ok(()) + }; + publish_score("before-publication", first_score).await?; + parent_gates.next().unwrap().send(()).unwrap(); + timeout(Duration::from_secs(30), async { + classifier_server.wait_for_request_count(/*count*/ 3).await; + parent_server.wait_for_request_count(/*count*/ 4).await; + }) + .await?; + + // Call 3's result is now fresh enough for call 4, whose own score is still held. + publish_score("stale-publication", current_score).await?; + parent_gates.next().unwrap().send(()).unwrap(); + + let mut review_reasons = Vec::new(); + loop { + match wait_for_event(&test.codex, |_| true).await { + EventMsg::GuardianAssessment(event) + if event.status == GuardianAssessmentStatus::Denied => + { + review_reasons.push(event.review_reason); + } + EventMsg::RequestPermissions(event) => panic!("unexpected user prompt: {event:?}"), + EventMsg::TurnComplete(_) => break, + _ => {} + } + } + assert_eq!( + review_reasons, + vec![ + Some(GuardianReviewReason::MissingScore), + Some(GuardianReviewReason::StaleScore) + ], + ); + let request = final_response.single_request(); + for (call_id, network) in [ + ("before-publication", None), + ("stale-publication", None), + ( + "fresh-publication", + Some(NetworkPermissions { + enabled: Some(true), + }), + ), + ] { + let output = request + .function_call_output_text(call_id) + .expect("permission result"); + assert_eq!( + serde_json::from_str::(&output)?, + RequestPermissionsResponse { + permissions: RequestPermissionProfile { + network, + file_system: None + }, + scope: PermissionGrantScope::Turn, + strict_auto_review: false, + }, + "{call_id}", + ); + } + + test.codex.shutdown_and_wait().await?; + parent_server.shutdown().await; + classifier_server.shutdown().await; + Ok(()) +} diff --git a/codex-rs/core/tests/suite/mod.rs b/codex-rs/core/tests/suite/mod.rs index acea3ca06a..5208956150 100644 --- a/codex-rs/core/tests/suite/mod.rs +++ b/codex-rs/core/tests/suite/mod.rs @@ -74,6 +74,8 @@ mod external_auth; mod fork_thread; mod git_enrichment; mod guardian_authorization; +#[path = "guardian_cached_score_tests.rs"] +mod guardian_cached_score; #[path = "guardian_checkpoint_migration_tests.rs"] mod guardian_checkpoint_migration; // Uses the same command-approval harness as guardian_review below. diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs b/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs index a557ce1845..faa3d73455 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/approval.rs @@ -22,9 +22,7 @@ use codex_protocol::openai_models::GuardianReviewMode; use codex_protocol::openai_models::GuardianScope; use codex_protocol::protocol::AskForApproval; use codex_protocol::protocol::TruncationPolicy; -use codex_protocol::security_risk::SecurityRiskScore; use std::sync::Weak; -use std::sync::atomic::Ordering; pub(super) struct GuardianApprovalReviewer { pub(super) thread_manager: Weak, @@ -144,19 +142,12 @@ async fn cached_evidence( let max_action_bytes = TruncationPolicy::Tokens(config.max_action_tokens).byte_budget(); let action_fits = serde_json::to_string_pretty(&action) .is_ok_and(|action| action.len().saturating_add(1) <= max_action_bytes); - if !action_fits - || input.tool_call_id.is_some_and(|call_id| { - progress - .oversized_tool_calls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .contains(call_id) - }) - { + let history = thread.conversation_history_snapshot().await; + let cached = progress.inspect(input.tool_call_id); + if !action_fits || cached.oversized { record_fast_decision(metrics, "deferred", "scoring_failure"); return Err(GuardianReviewReason::ScoringFailure); } - let history = thread.conversation_history_snapshot().await; let context_mode = GuardianContextMode::from_history(history.as_ref()); if context_mode == GuardianContextMode::ThreadOwned { let sampler = store @@ -183,12 +174,18 @@ async fn cached_evidence( .get("connector_id") .and_then(serde_json::Value::as_str) == Some("node_repl") - && progress.js_executions.load(Ordering::Acquire) == 1 + && cached.js_executions == 1 { record_fast_decision(metrics, "approved", "initial_cua_call"); return Ok(()); } let current = ScoreAuthorization::current(thread).await; + // Classification may publish or fail while authorization is collected. + let cached = progress.inspect(input.tool_call_id); + if cached.oversized { + record_fast_decision(metrics, "deferred", "scoring_failure"); + return Err(GuardianReviewReason::ScoringFailure); + } if !current.local.retained_context_complete || current .root @@ -197,20 +194,7 @@ async fn cached_evidence( record_fast_decision(metrics, "deferred", "incomplete_authorization"); return Err(GuardianReviewReason::Policy); } - let scored_authorization = progress - .authorization - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let latest_scored = progress.latest_scored_tool_call.load(Ordering::Acquire); - let lag = progress - .latest_tool_call - .load(Ordering::Acquire) - .saturating_sub(latest_scored) - .saturating_sub( - progress - .wrapper_lag - .discount(input.tool_call_id, latest_scored), - ); + let lag = cached.lag; if let Some(metrics) = metrics { metrics.histogram( TOOL_CALL_LAG_METRIC, @@ -220,10 +204,7 @@ async fn cached_evidence( } // Reuse the latest thread score within the lag limit, even across categories // and while the current action's async score is still in flight. - let score = store - .get::() - .and_then(|score| score.scores.get("action_risk").copied()); - let (reason, label) = match score { + let (reason, label) = match cached.action_risk { _ if lag > config.max_tool_call_lag => { if let Some(metrics) = metrics { metrics.counter( @@ -234,12 +215,12 @@ async fn cached_evidence( } (GuardianReviewReason::StaleScore, "stale_score") } - _ if progress.latest_failed_tool_call.load(Ordering::Acquire) > latest_scored => { + _ if cached.has_unscored_failure => { (GuardianReviewReason::ScoringFailure, "scoring_failure") } None => (GuardianReviewReason::MissingScore, "missing_score"), Some(score) if score < config.review_threshold => { - if scored_authorization.as_ref() != Some(¤t) { + if cached.authorization.as_ref() != Some(¤t) { ( GuardianReviewReason::AuthorizationChanged, "authorization_changed", diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs index db98c57dd2..7c18812439 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/budget_tests.rs @@ -186,15 +186,18 @@ async fn assert_catalog_budget(evidence: BudgetEvidence) -> Result<()> { }; let thread_store = fixture.test.codex.thread_extension_data(); if matches!(evidence, BudgetEvidence::UserInstructions) { - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); let progress = thread_store.get::().unwrap(); let authorization = ScoreAuthorization::current(&fixture.test.codex).await; - *progress.authorization.lock().unwrap() = Some(authorization); + seed_cached_score(&progress, thread_store, /*index*/ 0, authorization); assert_eq!( cached_approval( &fixture.registry, @@ -242,9 +245,7 @@ async fn assert_catalog_budget(evidence: BudgetEvidence) -> Result<()> { } let progress = thread_store.get::().unwrap(); tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { - while progress.latest_scored_tool_call.load(Ordering::Acquire) - < progress.latest_tool_call.load(Ordering::Acquire) - { + while progress.inspect(/*call_id*/ None).lag > 0 { tokio::task::yield_now().await; } }) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs index 2e65196314..e9bed2cfdd 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs @@ -3,7 +3,6 @@ use std::collections::BTreeMap; use std::sync::Arc; -use std::sync::atomic::Ordering; use std::time::Instant; use std::time::SystemTime; @@ -47,7 +46,6 @@ use super::sampler::LunaSampler; use super::sampler::LunaSamplerError; use super::sampler::LunaSamplingRequest; use super::score::GuardianV2ScoreProgress; -use super::score::record_fail_closed_score; use super::transcript::ContextInput; use super::truncation::ClassificationTruncations; use super::trusted_skills::TrustedSkillInvocations; @@ -201,7 +199,7 @@ impl Classification { let mut transcript = match transcript { Ok(transcript) => transcript, Err(error) => { - record_fail_closed_score(thread.thread_extension_data(), sampled_at); + score_progress.fail_closed(sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -306,22 +304,8 @@ impl Classification { if score_authorization != ScoreAuthorization::current(&thread).await { return Ok(ClassificationOutcome::Superseded); } - let accepted = { - let mut scored_authorization = score_progress - .authorization - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); - let accepted = - thread - .thread_extension_data() - .insert_if(score.clone(), |previous| { - previous.is_none_or(|previous| previous.sampled_at < score.sampled_at) - }); - if accepted { - *scored_authorization = Some(score_authorization); - } - accepted - }; + let accepted = + score_progress.publish(score.clone(), score_authorization, tool_call_index); tracing::info!( %thread_id, %turn_id, @@ -336,9 +320,6 @@ impl Classification { if !accepted { return Ok(ClassificationOutcome::Superseded); } - score_progress - .latest_scored_tool_call - .fetch_max(tool_call_index, Ordering::Release); classification_finished_at = Some(Instant::now()); record_classification_risk(metrics.as_deref(), output.as_str()); if guardian_config.persist_scores @@ -359,7 +340,7 @@ impl Classification { } .await; if result.is_err() { - record_fail_closed_score(thread.thread_extension_data(), sampled_at); + score_progress.fail_closed(sampled_at); } let duration = classification_finished_at .map(|finished_at: Instant| finished_at.duration_since(classification_started_at)) diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs index 4adf7057d9..f6af94c22d 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension.rs @@ -88,10 +88,9 @@ impl ThreadLifecycleContributor for GuardianV2Extension { .thread_store .get_or_init(|| LunaSampler::new(sampler_config)); input.thread_store.insert(guardian_config); - input.thread_store.insert(GuardianV2ScoreProgress { - metrics: input.extension_metrics.clone(), - ..Default::default() - }); + input.thread_store.insert(GuardianV2ScoreProgress::new( + input.extension_metrics.clone(), + )); // Preserve the answer path selected by the host for this thread. input .thread_store @@ -156,11 +155,7 @@ impl ToolLifecycleContributor for GuardianV2Extension { fn on_tool_finish<'a>(&'a self, input: ToolFinishInput<'a>) -> ToolLifecycleFuture<'a> { Box::pin(async move { if let Some(progress) = input.thread_store.get::() { - progress - .oversized_tool_calls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .remove(input.call_id); + progress.finish(input.call_id); } }) } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index c2454b8008..aed5cec823 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -1,7 +1,6 @@ use std::collections::BTreeMap; use std::sync::Arc; use std::sync::Mutex; -use std::sync::atomic::Ordering; use std::time::Duration; use anyhow::Result; @@ -87,6 +86,8 @@ use crate::async_scorer::sampler::MODEL; use crate::async_scorer::sampler::tests::ProxyPrewarmLimit; use crate::async_scorer::sampler::tests::proxy_websocket_servers_with_http; use crate::async_scorer::score::GuardianV2ScoreProgress; +use crate::async_scorer::score::tests::cached_score; +use crate::async_scorer::score::tests::set_cached_score; use crate::async_scorer::transcript::MAX_MESSAGE_ENTRY_TOKENS; use crate::async_scorer::transcript::MAX_TOOL_ENTRY_TOKENS; use codex_features::GuardianV2ReviewScopeConfigToml; @@ -300,7 +301,7 @@ async fn installed_extension_uses_http_after_warm_socket_auth_expires() -> Resul }) .await; tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { - while progress.latest_scored_tool_call.load(Ordering::Acquire) < call_index { + while progress.inspect(/*call_id*/ None).lag > 0 { tokio::task::yield_now().await; } }) @@ -325,7 +326,7 @@ async fn installed_extension_uses_http_after_warm_socket_auth_expires() -> Resul .collect::>(), vec![Some("Bearer original".to_owned()); INITIAL_WEBSOCKET_CONNECTIONS] ); - assert_eq!(progress.latest_failed_tool_call.load(Ordering::Acquire), 0); + assert!(!progress.inspect(/*call_id*/ None).has_unscored_failure); let http_request = http_mock.single_request(); assert_eq!( http_request.header("authorization"), @@ -540,19 +541,13 @@ async fn sandboxed_shell_classification_respects_review_scope() -> Result<()> { let fixture = GuardianFailureFixture::new().await?; let thread_store = fixture.test.codex.thread_extension_data(); - let mut score = thread_store - .get::() - .expect("fixture should publish a score") - .as_ref() - .clone(); + let mut score = cached_score(thread_store).expect("fixture should publish a score"); score.scores.insert("action_risk".to_owned(), 0.0); - thread_store.insert(score); + set_cached_score(thread_store, score); let score_progress = thread_store .get::() .expect("Guardian v2 should track score progress per thread"); - let latest_scored_tool_call = score_progress - .latest_scored_tool_call - .load(Ordering::Acquire); + let mut expected = score_progress.inspect(/*call_id*/ None); let turn_store = ExtensionData::new("turn-1"); let tool_name = ToolName::plain("exec_command"); for (call_id, payload, expected_decision) in [ @@ -588,16 +583,9 @@ async fn sandboxed_shell_classification_respects_review_scope() -> Result<()> { ); } - assert_eq!( - score_progress.latest_tool_call.load(Ordering::Acquire), - latest_scored_tool_call + 2 - ); - assert_eq!( - score_progress - .latest_scored_tool_call - .load(Ordering::Acquire), - latest_scored_tool_call - ); + expected.lag += 2; + expected.has_unscored_failure = true; + assert_eq!(score_progress.inspect(/*call_id*/ None), expected); fixture.assert_fails_closed("scoring_failure").await?; Ok(()) } @@ -642,12 +630,15 @@ async fn computer_use_only_scores_cannot_approve_other_actions() -> Result<()> { .clone(); config.policy = legacy_loader(/*scope*/ None); thread_store.insert(config); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); thread_store.insert(RecordingMetrics::default()); let progress = thread_store .get::() @@ -655,8 +646,8 @@ async fn computer_use_only_scores_cannot_approve_other_actions() -> Result<()> { // The seeded low score belongs to the model selected above. let authorization = super::super::authorization::ScoreAuthorization::current(&fixture.test.codex).await; - *progress.authorization.lock().unwrap() = Some(authorization); - let latest_tool_call = progress.latest_tool_call.load(Ordering::Acquire); + seed_cached_score(&progress, thread_store, /*index*/ 1, authorization); + let cached = progress.inspect(/*call_id*/ None); let turn_store = ExtensionData::new("turn-1"); let ordinary_tool = ToolName::namespaced("mcp__ordinary__", "write_record"); let payload = ToolPayload::Function { @@ -682,8 +673,8 @@ async fn computer_use_only_scores_cannot_approve_other_actions() -> Result<()> { }) .await; assert_eq!( - progress.latest_tool_call.load(Ordering::Acquire), - latest_tool_call, + progress.inspect(/*call_id*/ None), + cached, "unrelated code-mode calls must not age browser/CUA scores" ); @@ -785,20 +776,20 @@ async fn computer_use_only_scores_cannot_approve_other_actions() -> Result<()> { fixture .score_tool(ToolName::namespaced("mcp__node_repl__", "js")) .await; - assert!( - progress.latest_failed_tool_call.load(Ordering::Acquire) - > progress.latest_scored_tool_call.load(Ordering::Acquire) - ); + assert!(progress.inspect(/*call_id*/ None).has_unscored_failure); for required in [false, true] { model.node_repl_auto_review_required = required; thread_store.insert(model.clone()); // Also reject a low score published by an older, in-flight classifier. - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.0)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.0)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( &fixture.registry, @@ -1027,36 +1018,20 @@ async fn unscored_tools_invalidate_cached_scores() -> Result<()> { let progress = thread_store .get::() .expect("score progress"); - let before = ( - progress.latest_tool_call.load(Ordering::Acquire), - progress.latest_failed_tool_call.load(Ordering::Acquire), - ); + let mut expected = progress.inspect(/*call_id*/ None); + expected.action_risk = Some(0.25); fixture.score_tool(ToolName::plain("wait")).await; - assert_eq!( - ( - progress.latest_tool_call.load(Ordering::Acquire), - progress.latest_failed_tool_call.load(Ordering::Acquire), - ), - before, - ); + assert_eq!(progress.inspect(/*call_id*/ None), expected); // A dynamic function named exec is not a Code Mode wrapper. An MCP tool // named wait is not a Code Mode poll. Both invalidate earlier scores. - for (index, tool) in [ + for tool in [ ToolName::plain("exec"), ToolName::namespaced("mcp__ordinary", "wait"), - ] - .into_iter() - .enumerate() - { + ] { fixture.score_tool(tool).await; - let expected = before.0 + index + 1; - assert_eq!( - ( - progress.latest_tool_call.load(Ordering::Acquire), - progress.latest_failed_tool_call.load(Ordering::Acquire), - ), - (expected, expected), - ); + expected.lag += 1; + expected.has_unscored_failure = true; + assert_eq!(progress.inspect(/*call_id*/ None), expected); } Ok(()) } @@ -1080,11 +1055,8 @@ impl GuardianFailureFixture { .get::() .expect("Guardian v2 should track score progress per thread"); tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { - while thread_store.get::().is_none() - || score_progress - .latest_scored_tool_call - .load(Ordering::Acquire) - < score_progress.latest_tool_call.load(Ordering::Acquire) + while cached_score(thread_store).is_none() + || score_progress.inspect(/*call_id*/ None).lag > 0 { tokio::task::yield_now().await; } @@ -1100,12 +1072,15 @@ impl GuardianFailureFixture { async fn score_tool(&self, tool_name: ToolName) { let thread_store = self.test.codex.thread_extension_data(); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); let turn_store = ExtensionData::new("turn-1"); let payload = ToolPayload::Function { arguments: r#"{"path":"README.md"}"#.to_owned(), @@ -1134,15 +1109,11 @@ impl GuardianFailureFixture { .get::() .expect("Guardian v2 should track score progress per thread"); tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { - while thread_store - .get::() + while cached_score(thread_store) .is_none_or(|score| score.scores.get("action_risk") != Some(&1.0)) - && score_progress - .latest_failed_tool_call - .load(Ordering::Acquire) - <= score_progress - .latest_scored_tool_call - .load(Ordering::Acquire) + && !score_progress + .inspect(/*call_id*/ None) + .has_unscored_failure { tokio::task::yield_now().await; } @@ -1458,9 +1429,9 @@ max_recent_non_user_entries = 8 let metrics = thread_store.get::().unwrap(); tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { while score_progress - .latest_scored_tool_call - .load(Ordering::Acquire) - == 0 + .inspect(/*call_id*/ None) + .authorization + .is_none() || metrics.classification_samples().len() < 10 { tokio::task::yield_now().await; @@ -1477,12 +1448,15 @@ max_recent_non_user_entries = 8 ("measurement".to_owned(), "text_bytes".to_owned()), ]) })); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.65)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.65)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -1495,12 +1469,15 @@ max_recent_non_user_entries = 8 .await, None ); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.55)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.55)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -1541,15 +1518,8 @@ max_recent_non_user_entries = 8 } } - assert_eq!( - score_progress - .latest_scored_tool_call - .load(Ordering::Acquire), - 1 - ); - score_progress - .latest_tool_call - .store(/*val*/ 3, Ordering::Release); + let first_unscored = observe_unscored_call(&score_progress, thread_store); + observe_unscored_call(&score_progress, thread_store); assert_eq!( cached_approval( ®istry, @@ -1565,9 +1535,7 @@ max_recent_non_user_entries = 8 let initial_metrics = thread_store.get::().unwrap(); thread_store.insert(RecordingMetrics::default()); - score_progress - .latest_tool_call - .store(/*val*/ 4, Ordering::Release); + observe_unscored_call(&score_progress, thread_store); assert_eq!( cached_approval( ®istry, @@ -1581,9 +1549,12 @@ max_recent_non_user_entries = 8 None ); - score_progress - .latest_scored_tool_call - .store(/*val*/ 2, Ordering::Release); + seed_cached_score( + &score_progress, + thread_store, + first_unscored, + ScoreAuthorization::current(&test.codex).await, + ); assert_eq!( cached_approval( ®istry, @@ -1879,7 +1850,7 @@ async fn contributor_uses_model_defaults_and_preserves_local_overrides() -> Resu assert!(thread_store.get::().is_some()); let score = tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { loop { - if let Some(score) = thread_store.get::() { + if let Some(score) = cached_score(thread_store) { return score; } tokio::task::yield_now().await; @@ -1890,12 +1861,15 @@ async fn contributor_uses_model_defaults_and_preserves_local_overrides() -> Resu score.action, Some(serde_json::from_str::(action)?) ); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.65)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.65)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -2056,7 +2030,7 @@ async fn assert_luna_pool_context(thread_context_enabled: bool) -> Result<()> { assert_eq!(request["input"][2]["content"], expected_content); let score = tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { loop { - if let Some(score) = thread_store.get::() { + if let Some(score) = cached_score(thread_store) { return score; } tokio::task::yield_now().await; @@ -2064,7 +2038,7 @@ async fn assert_luna_pool_context(thread_context_enabled: bool) -> Result<()> { }) .await?; assert_eq!( - score.as_ref(), + &score, &SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_string(), 1.0)]), call_id: Some("call-1".to_owned()), @@ -2094,12 +2068,15 @@ async fn assert_luna_pool_context(thread_context_enabled: bool) -> Result<()> { .await, None ); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_string(), 0.5)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_string(), 0.5)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -2111,12 +2088,15 @@ async fn assert_luna_pool_context(thread_context_enabled: bool) -> Result<()> { None ); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_string(), 0.49)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_string(), 0.49)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -2129,12 +2109,15 @@ async fn assert_luna_pool_context(thread_context_enabled: bool) -> Result<()> { ); let disabled_thread_store = ExtensionData::new("disabled-thread"); - disabled_thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_string(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + &disabled_thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_string(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -2286,13 +2269,16 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> { let progress = thread_store .get::() .expect("Guardian v2 should track score progress per thread"); - *progress.authorization.lock().unwrap() = Some(authorization); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + seed_cached_score(&progress, thread_store, /*index*/ 0, authorization); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -2333,7 +2319,7 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> { .await; assert!( - thread_store.get::().is_none(), + cached_score(thread_store).is_none(), "protected models must not receive Guardian v2 fail-closed scores" ); assert!( @@ -2359,17 +2345,20 @@ async fn cached_score_survives_compaction_and_internal_context_but_not_user_inpu let thread_store = test.codex.thread_extension_data(); let progress = thread_store.get::().unwrap(); tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { - while progress.latest_scored_tool_call.load(Ordering::Acquire) == 0 { + while progress.inspect(/*call_id*/ None).authorization.is_none() { tokio::task::yield_now().await; } }) .await?; - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); test.codex .inject_response_items(vec![ContextualUserFragment::into( @@ -2438,12 +2427,15 @@ async fn assert_compaction_approval_policy(thread_context_enabled: bool) -> Resu )) .await?; let thread_store = fixture.test.codex.thread_extension_data(); - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.0)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.0)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( &fixture.registry, @@ -2478,12 +2470,15 @@ async fn assert_compaction_approval_policy(thread_context_enabled: bool) -> Resu authorization ); let score_authorization = ScoreAuthorization::current(&fixture.test.codex).await; - *thread_store + let progress = thread_store .get::() - .expect("score progress") - .authorization - .lock() - .unwrap() = Some(score_authorization); + .expect("score progress"); + seed_cached_score( + &progress, + thread_store, + /*index*/ 1, + score_authorization, + ); // No new sample runs: only the enabled path rejects cached and initial-call approvals. for (computer_use_only, prompt) in [ (false, "review action"), @@ -2501,11 +2496,9 @@ async fn assert_compaction_approval_policy(thread_context_enabled: bool) -> Resu sandboxed_exec_commands: Some(true), })); thread_store.insert(config); - thread_store - .get::() - .expect("score progress") - .js_executions - .store(/*val*/ 1, Ordering::Release); + if computer_use_only { + progress.observe_js_execution(); + } assert_eq!( cached_approval( &fixture.registry, @@ -2539,20 +2532,23 @@ async fn contributor_counts_failed_thread_lookups_toward_score_lag() -> Result<( .expect("Guardian v2 should track score progress per thread"); tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { while score_progress - .latest_scored_tool_call - .load(Ordering::Acquire) - == 0 + .inspect(/*call_id*/ None) + .authorization + .is_none() { tokio::task::yield_now().await; } }) .await?; - thread_store.insert(SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), - call_id: None, - action: None, - sampled_at: None, - }); + set_cached_score( + thread_store, + SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }, + ); assert_eq!( cached_approval( ®istry, @@ -2590,7 +2586,7 @@ async fn contributor_counts_failed_thread_lookups_toward_score_lag() -> Result<( }) .await; - assert_eq!(score_progress.latest_tool_call.load(Ordering::Acquire), 2); + assert_eq!(score_progress.inspect(/*call_id*/ None).lag, 1); assert_eq!( cached_approval( ®istry, @@ -3131,7 +3127,7 @@ async fn assert_parent_compaction_reuse(thread_context_enabled: bool) -> Result< let previous_score = tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { loop { - if let Some(score) = thread_store.get::() { + if let Some(score) = cached_score(thread_store) { return score; } tokio::task::yield_now().await; @@ -3182,11 +3178,10 @@ async fn assert_parent_compaction_reuse(thread_context_enabled: bool) -> Result< }) .await; - let fail_closed_score = thread_store - .get::() + let fail_closed_score = cached_score(thread_store) .expect("an oversized compaction should immediately receive the maximum risk score"); assert_eq!( - fail_closed_score.as_ref(), + &fail_closed_score, &SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), call_id: None, @@ -3258,7 +3253,7 @@ async fn legacy_contributor_can_disable_parent_compaction_reuse() -> Result<()> let thread_store = test.codex.thread_extension_data(); let score = tokio::time::timeout(ASYNC_TEST_TIMEOUT, async { loop { - if let Some(score) = thread_store.get::() { + if let Some(score) = cached_score(thread_store) { return score; } tokio::task::yield_now().await; @@ -3286,9 +3281,9 @@ async fn cached_approval_discounts_only_its_own_unscored_wrapper() -> Result<()> let fixture = GuardianFailureFixture::new().await?; let store = fixture.test.codex.thread_extension_data(); let progress = store.get::().unwrap(); - let mut score = store.get::().unwrap().as_ref().clone(); + let mut score = cached_score(store).unwrap(); score.scores.insert("action_risk".to_owned(), 0.0); - store.insert(score); + set_cached_score(store, score); // Hold the cached score fixed while advancing real tool-start metadata. let start = |call_id: &str, origin: &ResponseItemId, source: ToolCallSource| { let tool_name = match source { @@ -3303,28 +3298,20 @@ async fn cached_approval_discounts_only_its_own_unscored_wrapper() -> Result<()> arguments: "{}".to_owned(), }, }; - let index = progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - + 1; - progress.wrapper_lag.record( - &ToolStartInput { - session_store: &fixture.session_store, - thread_store: store, - turn_store: &fixture.session_store, - turn_id: "turn", - root_turn_id: None, - call_id, - originating_item_id: Some(origin), - tool_name: &tool_name, - mcp_tool: None, - payload: &payload, - conversation_history: Arc::new(TestConversationHistory(Vec::new())), - source, - }, - index, - ); - index + progress.observe(&ToolStartInput { + session_store: &fixture.session_store, + thread_store: store, + turn_store: &fixture.session_store, + turn_id: "turn", + root_turn_id: None, + call_id, + originating_item_id: Some(origin), + tool_name: &tool_name, + mcp_tool: None, + payload: &payload, + conversation_history: Arc::new(TestConversationHistory(Vec::new())), + source, + }) }; let approve = async |call_id: &str| { cached_approval( @@ -3350,13 +3337,19 @@ async fn cached_approval_discounts_only_its_own_unscored_wrapper() -> Result<()> assert_eq!(approve(call_id).await, expected); } // A score covering the wrapper must not receive another discount. - progress - .latest_scored_tool_call - .store(wrapper, Ordering::Release); + seed_cached_score( + &progress, + store, + wrapper, + ScoreAuthorization::current(&fixture.test.codex).await, + ); assert_eq!(approve("third").await, None); - progress - .latest_scored_tool_call - .store(wrapper + 3, Ordering::Release); + seed_cached_score( + &progress, + store, + wrapper + 3, + ScoreAuthorization::current(&fixture.test.codex).await, + ); let output = start("output-only", &origin, ToolCallSource::Direct); let other = ResponseItemId::from_server("other-wrapper".to_owned()); start("other-wrapper", &other, ToolCallSource::Direct); @@ -3366,9 +3359,12 @@ async fn cached_approval_discounts_only_its_own_unscored_wrapper() -> Result<()> start("other-second", &other, nested.clone()); assert_eq!(approve("other-second").await, None); // Covering a different wrapper still leaves this call's parent in its lag. - progress - .latest_scored_tool_call - .store(output, Ordering::Release); + seed_cached_score( + &progress, + store, + output, + ScoreAuthorization::current(&fixture.test.codex).await, + ); assert_eq!( approve("other-second").await, Some(ReviewDecision::Approved) @@ -3379,10 +3375,8 @@ async fn cached_approval_discounts_only_its_own_unscored_wrapper() -> Result<()> } start("late-child", &other, nested); assert_eq!( - progress - .wrapper_lag - .discount(Some("late-child"), /*latest_scored*/ 0), - 0 + progress.inspect(Some("late-child")).lag, + progress.inspect(/*call_id*/ None).lag ); Ok(()) } @@ -3448,3 +3442,88 @@ fn review_scope(action: &serde_json::Value) -> Option { #[path = "budget_tests.rs"] mod budget; + +fn seed_cached_score( + progress: &GuardianV2ScoreProgress, + store: &ExtensionData, + index: usize, + authorization: ScoreAuthorization, +) { + let mut score = cached_score(store).unwrap_or(SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 0.25)]), + call_id: None, + action: None, + sampled_at: None, + }); + score.sampled_at = Some(std::time::SystemTime::now().into()); + assert!(progress.publish(score, authorization, index)); +} + +fn observe_unscored_call(progress: &GuardianV2ScoreProgress, store: &ExtensionData) -> usize { + progress.observe(&ToolStartInput { + session_store: store, + thread_store: store, + turn_store: store, + turn_id: "turn-1", + root_turn_id: None, + call_id: "unscored", + originating_item_id: None, + tool_name: &ToolName::plain("read_file"), + mcp_tool: None, + payload: &ToolPayload::Function { + arguments: "{}".to_owned(), + }, + conversation_history: Arc::new(TestConversationHistory(Vec::new())), + source: ToolCallSource::Direct, + }) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cached_score_publication_rejects_delayed_results_without_changing_coverage() -> Result<()> +{ + skip_if_no_network!(Ok(())); + let fixture = GuardianFailureFixture::new().await?; + let store = fixture.test.codex.thread_extension_data(); + let progress = store.get::().unwrap(); + let authorization = ScoreAuthorization::current(&fixture.test.codex).await; + seed_cached_score(&progress, store, /*index*/ 1, authorization.clone()); + let score = cached_score(store).unwrap(); + observe_unscored_call(&progress, store); + observe_unscored_call(&progress, store); + let cached = progress.inspect(/*call_id*/ None); + let mut outdated_authorization = authorization.clone(); + outdated_authorization.local.user_message_revision += 1; + // Even a result for a higher call index cannot overwrite an equally old sample. + assert!(!progress.publish(score.clone(), outdated_authorization, /*index*/ 2)); + assert_eq!(progress.inspect(/*call_id*/ None), cached); + assert_eq!(cached_score(store).as_ref(), Some(&score)); + + let failed_at = std::time::SystemTime::now() + Duration::from_secs(1); + progress.fail_closed(failed_at); + let mut delayed_score = score; + delayed_score.sampled_at = Some(failed_at.into()); + assert!(!progress.publish( + delayed_score.clone(), + authorization.clone(), + /*index*/ 2 + )); + + delayed_score.sampled_at = Some((failed_at + Duration::from_secs(1)).into()); + delayed_score.scores.insert("action_risk".to_owned(), 0.0); + progress.mark_oversized("active-overflow", /*index*/ 2); + assert!(progress.publish( + delayed_score.clone(), + authorization.clone(), + /*index*/ 3 + )); + assert_eq!(cached_score(store).as_ref(), Some(&delayed_score)); + let mut expected = cached; + expected.lag = 0; + expected.action_risk = Some(0.0); + expected.authorization = Some(authorization); + assert_eq!(progress.inspect(/*call_id*/ None), expected); + assert!(progress.inspect(Some("active-overflow")).oversized); + progress.finish("active-overflow"); + assert!(!progress.inspect(Some("active-overflow")).oversized); + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs index b5a8158f71..f67539258b 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/observation.rs @@ -2,7 +2,6 @@ //! Keep snapshots on their existing side of the background task boundary. use std::sync::Arc; -use std::sync::atomic::Ordering; use std::time::Instant; use std::time::SystemTime; @@ -21,7 +20,6 @@ use codex_protocol::mcp::is_node_repl_backed_server; use codex_protocol::openai_models::GuardianReviewMode; use codex_protocol::openai_models::GuardianScope; use codex_protocol::openai_models::ModelInfo; -use codex_protocol::security_risk::SecurityRiskScore; use super::action::ActionRenderError; use super::action::GuardianAction; @@ -35,7 +33,6 @@ use super::parent_compaction::ParentCompactionError; use super::parent_compaction::select_parent_compaction; use super::sampler::LunaSampler; use super::score::GuardianV2ScoreProgress; -use super::score::record_fail_closed_score; use codex_protocol::openai_models::GuardianUnscoredAction as UnscoredAction; impl GuardianV2Extension { @@ -80,11 +77,7 @@ impl GuardianV2Extension { match policy.unscored_action { UnscoredAction::Ignore => {} UnscoredAction::AgeScore => { - let index = score_progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - .saturating_add(/*rhs*/ 1); - score_progress.wrapper_lag.record(&input, index); + let index = score_progress.observe(&input); // Unscored permission widening must not reuse an earlier approval score. if input.tool_name.is_default_namespace() && input.tool_name.name == "exec_command" @@ -95,20 +88,12 @@ impl GuardianV2Extension { .and_then(serde_json::Value::as_str) == Some("with_additional_permissions") { - score_progress - .latest_failed_tool_call - .fetch_max(index, Ordering::Release); + score_progress.invalidate(index); } } UnscoredAction::InvalidateScore => { - let index = score_progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - .saturating_add(/*rhs*/ 1); - score_progress.wrapper_lag.record(&input, index); - score_progress - .latest_failed_tool_call - .fetch_max(index, Ordering::Release); + let index = score_progress.observe(&input); + score_progress.invalidate(index); } } return; @@ -117,18 +102,12 @@ impl GuardianV2Extension { let info = tool.tool_info(); is_node_repl_backed_server(&info.server_name) && info.tool.name == "js" }) { - score_progress - .js_executions - .fetch_add(/*val*/ 1, Ordering::Relaxed); + score_progress.observe_js_execution(); } let metrics = score_progress.metrics.clone(); let analytics = input.session_store.get::(); let sampled_at = SystemTime::now(); - let tool_call_index = score_progress - .latest_tool_call - .fetch_add(/*val*/ 1, Ordering::Relaxed) - .saturating_add(/*rhs*/ 1); - score_progress.wrapper_lag.record(&input, tool_call_index); + let tool_call_index = score_progress.observe(&input); let event_sink = Arc::clone(&self.event_sink); let thread_id = input.thread_store.level_id().to_owned(); let turn_id = input.turn_id.to_owned(); @@ -155,9 +134,7 @@ impl GuardianV2Extension { let (manager, thread, config) = match thread_context { Ok(context) => context, Err(error) => { - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); + score_progress.invalidate(tool_call_index); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -178,9 +155,7 @@ impl GuardianV2Extension { || thread.approvals_reviewer_for_turn(input.turn_id).await == ApprovalsReviewer::User { // A skipped call invalidates older scores, including ones still in flight. - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); + score_progress.invalidate(tool_call_index); return; } // A required model keeps synchronous review outside its CUA allowance. @@ -192,7 +167,7 @@ impl GuardianV2Extension { .auto_review_required_for_model(&model.slug) }) { - input.thread_store.remove::(); + score_progress.clear_score(); return; } input.thread_store.insert(GuardianV2Enabled); @@ -203,7 +178,7 @@ impl GuardianV2Extension { let guardian_config = match guardian_config.with_model_defaults(model_defaults) { Ok(config) => config, Err(error) => { - record_fail_closed_score(input.thread_store, sampled_at); + score_progress.fail_closed(sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -241,14 +216,12 @@ impl GuardianV2Extension { Ok(compaction) => compaction, Err(error) => { let (outcome, failure_reason) = if error == ParentCompactionError::RequiresSync { - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); + score_progress.invalidate(tool_call_index); ("skipped", None) } else { ("failure", Some("parent_compaction_error")) }; - record_fail_closed_score(input.thread_store, sampled_at); + score_progress.fail_closed(sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -269,15 +242,8 @@ impl GuardianV2Extension { let planned_action = match action.render(guardian_config.max_action_tokens) { Ok(text) => text, Err(ActionRenderError::TooLarge { .. }) => { - score_progress - .oversized_tool_calls - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner) - .insert(input.call_id.to_owned()); - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - record_fail_closed_score(input.thread_store, sampled_at); + score_progress.mark_oversized(input.call_id, tool_call_index); + score_progress.fail_closed(sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), @@ -287,10 +253,8 @@ impl GuardianV2Extension { return; } Err(error) => { - score_progress - .latest_failed_tool_call - .fetch_max(tool_call_index, Ordering::Release); - record_fail_closed_score(input.thread_store, sampled_at); + score_progress.invalidate(tool_call_index); + score_progress.fail_closed(sampled_at); record_classification( metrics.as_deref(), classification_started_at.elapsed(), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/score.rs b/codex-rs/ext/guardian-v2/src/async_scorer/score.rs index 47729efd3c..ade8ece519 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/score.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/score.rs @@ -1,15 +1,15 @@ -//! Tracks observation coverage and the authorization attached to the latest score. -//! Failed samples may replace an equally old score, but never a newer one. +//! Owns cached-score publication, observation coverage, and reuse snapshots. +//! The cached score, authorization, and coverage are published under one lock; +//! failures never replace a newer sample. use std::collections::BTreeMap; use std::collections::BTreeSet; use std::sync::Arc; use std::sync::Mutex; -use std::sync::atomic::AtomicUsize; use std::time::SystemTime; -use codex_extension_api::ExtensionData; use codex_extension_api::ExtensionMetrics; +use codex_extension_api::ToolStartInput; use codex_protocol::security_risk::SecurityRiskScore; use super::authorization::ScoreAuthorization; @@ -17,32 +17,163 @@ use super::wrapper_lag::WrapperLag; #[derive(Default)] pub(super) struct GuardianV2ScoreProgress { - pub(super) wrapper_lag: WrapperLag, - pub(super) latest_tool_call: AtomicUsize, - // Setup and reset calls must not consume the first JS execution allowance. - pub(super) js_executions: AtomicUsize, - pub(super) latest_scored_tool_call: AtomicUsize, - pub(super) latest_failed_tool_call: AtomicUsize, - // Keep overflow attached to each active call even after a newer score succeeds. - // The host's finish callback removes entries on completion, failure, or cancellation. - pub(super) oversized_tool_calls: Mutex>, - // Serialize successful score publication with its authorization metadata. - pub(super) authorization: Mutex>, + state: Mutex, pub(super) metrics: Option>, } -pub(super) fn record_fail_closed_score(thread_store: &ExtensionData, sampled_at: SystemTime) { - let score = SecurityRiskScore { - scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), - call_id: None, - action: None, - sampled_at: Some(sampled_at.into()), - }; - thread_store.insert_if(score.clone(), |previous| { - previous.is_none_or(|previous| previous.sampled_at <= score.sampled_at) - }); +#[derive(Default)] +struct ScoreState { + score: Option, + wrapper_lag: WrapperLag, + latest_tool_call: usize, + js_executions: usize, + latest_scored_tool_call: usize, + latest_failed_tool_call: usize, + oversized_tool_calls: BTreeSet, + authorization: Option, +} + +/// A consistent view of the published score and the observations it covers. +#[derive(Debug, PartialEq)] +pub(super) struct CachedScore { + pub(super) lag: usize, + pub(super) has_unscored_failure: bool, + pub(super) js_executions: usize, + pub(super) oversized: bool, + pub(super) action_risk: Option, + pub(super) authorization: Option, +} + +impl GuardianV2ScoreProgress { + pub(super) fn new(metrics: Option>) -> Self { + Self { + metrics, + ..Default::default() + } + } + + pub(super) fn observe(&self, input: &ToolStartInput<'_>) -> usize { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.latest_tool_call = state.latest_tool_call.saturating_add(/*rhs*/ 1); + let index = state.latest_tool_call; + state.wrapper_lag.record(input, index); + index + } + + pub(super) fn observe_js_execution(&self) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Setup/reset calls must not consume the first JS execution allowance. + state.js_executions = state.js_executions.saturating_add(/*rhs*/ 1); + } + + pub(super) fn invalidate(&self, index: usize) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.latest_failed_tool_call = state.latest_failed_tool_call.max(index); + } + + pub(super) fn mark_oversized(&self, call_id: &str, index: usize) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + state.oversized_tool_calls.insert(call_id.to_owned()); + state.latest_failed_tool_call = state.latest_failed_tool_call.max(index); + } + + pub(super) fn finish(&self, call_id: &str) { + // Overflow belongs to each active call even after a newer score succeeds. + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .oversized_tool_calls + .remove(call_id); + } + + pub(super) fn publish( + &self, + score: SecurityRiskScore, + authorization: ScoreAuthorization, + index: usize, + ) -> bool { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let accepted = state + .score + .as_ref() + .is_none_or(|previous| previous.sampled_at < score.sampled_at); + if accepted { + state.score = Some(score); + state.authorization = Some(authorization); + state.latest_scored_tool_call = state.latest_scored_tool_call.max(index); + } + accepted + } + + pub(super) fn fail_closed(&self, sampled_at: SystemTime) { + let mut state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let score = SecurityRiskScore { + scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), + call_id: None, + action: None, + sampled_at: Some(sampled_at.into()), + }; + // Failures win timestamp ties; successful samples must be strictly newer. + if state + .score + .as_ref() + .is_none_or(|previous| previous.sampled_at <= score.sampled_at) + { + state.score = Some(score); + } + } + + pub(super) fn clear_score(&self) { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .score = None; + } + + pub(super) fn inspect(&self, call_id: Option<&str>) -> CachedScore { + let state = self + .state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + CachedScore { + lag: state + .latest_tool_call + .saturating_sub(state.latest_scored_tool_call) + .saturating_sub( + state + .wrapper_lag + .discount(call_id, state.latest_scored_tool_call), + ), + has_unscored_failure: state.latest_failed_tool_call > state.latest_scored_tool_call, + js_executions: state.js_executions, + oversized: call_id.is_some_and(|id| state.oversized_tool_calls.contains(id)), + action_risk: state + .score + .as_ref() + .and_then(|score| score.scores.get("action_risk").copied()), + authorization: state.authorization.clone(), + } + } } #[cfg(test)] #[path = "score_tests.rs"] -mod tests; +pub(super) mod tests; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs index 8fcf3bff9f..2489f13cc1 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/score_tests.rs @@ -6,11 +6,24 @@ use codex_extension_api::ExtensionData; use codex_protocol::security_risk::SecurityRiskScore; use pretty_assertions::assert_eq; -use super::record_fail_closed_score; +use super::GuardianV2ScoreProgress; + +pub(in crate::async_scorer) fn cached_score(store: &ExtensionData) -> Option { + let progress = store.get::()?; + let state = progress.state.lock().unwrap(); + state.score.clone() +} + +// Replace only the score so fixtures can exercise authorization and coverage independently. +pub(in crate::async_scorer) fn set_cached_score(store: &ExtensionData, score: SecurityRiskScore) { + let progress = store.get_or_init(GuardianV2ScoreProgress::default); + progress.state.lock().unwrap().score = Some(score); +} #[test] fn fail_closed_score_preserves_classification_order() { let thread_store = ExtensionData::new("thread-1"); + let progress = thread_store.get_or_init(GuardianV2ScoreProgress::default); let newer_sampled_at = SystemTime::UNIX_EPOCH + Duration::from_secs(1); let newest_sampled_at = newer_sampled_at + Duration::from_secs(1); let newer_score = SecurityRiskScore { @@ -19,17 +32,14 @@ fn fail_closed_score_preserves_classification_order() { action: None, sampled_at: Some(newer_sampled_at.into()), }; - thread_store.insert(newer_score.clone()); + set_cached_score(&thread_store, newer_score.clone()); - record_fail_closed_score(&thread_store, SystemTime::UNIX_EPOCH); - assert_eq!( - thread_store.get::().as_deref(), - Some(&newer_score) - ); + progress.fail_closed(SystemTime::UNIX_EPOCH); + assert_eq!(cached_score(&thread_store).as_ref(), Some(&newer_score)); for sampled_at in [newer_sampled_at, newest_sampled_at] { - thread_store.insert(newer_score.clone()); - record_fail_closed_score(&thread_store, sampled_at); + set_cached_score(&thread_store, newer_score.clone()); + progress.fail_closed(sampled_at); let fail_closed_score = SecurityRiskScore { scores: BTreeMap::from([("action_risk".to_owned(), 1.0)]), call_id: None, @@ -37,7 +47,7 @@ fn fail_closed_score_preserves_classification_order() { sampled_at: Some(sampled_at.into()), }; assert_eq!( - thread_store.get::().as_deref(), + cached_score(&thread_store).as_ref(), Some(&fail_closed_score) ); } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/wrapper_lag.rs b/codex-rs/ext/guardian-v2/src/async_scorer/wrapper_lag.rs index bde42597b3..b4e198b4f6 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/wrapper_lag.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/wrapper_lag.rs @@ -1,8 +1,8 @@ //! Discounts only an approval's own unscored code-mode wrapper. Missing provenance //! keeps the full lag; the bounded history never changes scoring or failure order. +//! The score owner's state lock protects this history. use std::collections::VecDeque; -use std::sync::Mutex; use codex_extension_api::ToolCallSource; use codex_extension_api::ToolPayload; @@ -17,14 +17,13 @@ struct ToolStart { } #[derive(Default)] -pub(super) struct WrapperLag(Mutex>); +pub(super) struct WrapperLag { + starts: VecDeque, +} impl WrapperLag { - pub(super) fn record(&self, input: &ToolStartInput<'_>, index: usize) { - let mut starts = self - .0 - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); + pub(super) fn record(&mut self, input: &ToolStartInput<'_>, index: usize) { + let starts = &mut self.starts; let parent_wrapper_index = match &input.source { ToolCallSource::CodeMode { .. } => input.originating_item_id.and_then(|item_id| { starts @@ -56,14 +55,10 @@ impl WrapperLag { } pub(super) fn discount(&self, call_id: Option<&str>, latest_scored: usize) -> usize { - let starts = self - .0 - .lock() - .unwrap_or_else(std::sync::PoisonError::into_inner); usize::from( call_id .and_then(|call_id| { - starts + self.starts .iter() .rev() .find(|start| start.call_id == call_id)