Publish Guardian cached scores and coverage atomically (#46245)

## Why

Asynchronous score publication updated the risk score, authorization, and tool-call coverage separately. Approval checks could therefore read inconsistent evidence when deciding whether to reuse a cached score.

## What changed

Keep Guardian's cached score and observation state under one lock, and use consistent snapshots for approval checks. Publish successful scores together with their authorization and classified tool-call index. Preserve timestamp ordering, fail-closed precedence on timestamp ties, and per-call oversized-action tracking.

## Testing

Add regression coverage for rejected delayed results and timestamp ties. Add a gated integration test showing that a delayed score becomes stale after intervening tool calls, while a sufficiently fresh score permits cached approval.

GitOrigin-RevId: 8df6fafd5e0ee59aa0f78f541eac562efce3d4d0
This commit is contained in:
felixxia-oai
2026-09-17 15:57:29 +00:00
committed by copyberry
parent 2833985d88
commit fcf05456bb
12 changed files with 833 additions and 401 deletions

View File

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

View File

@@ -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::<RequestPermissionsResponse>(&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(())
}

View File

@@ -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.

View File

@@ -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<ThreadManager>,
@@ -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::<SecurityRiskScore>()
.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(&current) {
if cached.authorization.as_ref() != Some(&current) {
(
GuardianReviewReason::AuthorizationChanged,
"authorization_changed",

View File

@@ -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::<GuardianV2ScoreProgress>().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::<GuardianV2ScoreProgress>().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;
}
})

View File

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

View File

@@ -88,10 +88,9 @@ impl ThreadLifecycleContributor<Config> 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::<GuardianV2ScoreProgress>() {
progress
.oversized_tool_calls
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.remove(input.call_id);
progress.finish(input.call_id);
}
})
}

View File

@@ -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<_>>(),
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::<SecurityRiskScore>()
.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::<GuardianV2ScoreProgress>()
.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::<GuardianV2ScoreProgress>()
@@ -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::<GuardianV2ScoreProgress>()
.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::<GuardianV2ScoreProgress>()
.expect("Guardian v2 should track score progress per thread");
tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
while thread_store.get::<SecurityRiskScore>().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::<GuardianV2ScoreProgress>()
.expect("Guardian v2 should track score progress per thread");
tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
while thread_store
.get::<SecurityRiskScore>()
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::<RecordingMetrics>().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(
&registry,
@@ -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(
&registry,
@@ -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(
&registry,
@@ -1565,9 +1535,7 @@ max_recent_non_user_entries = 8
let initial_metrics = thread_store.get::<RecordingMetrics>().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(
&registry,
@@ -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(
&registry,
@@ -1879,7 +1850,7 @@ async fn contributor_uses_model_defaults_and_preserves_local_overrides() -> Resu
assert!(thread_store.get::<NodeReplReviewEvidence>().is_some());
let score = tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
loop {
if let Some(score) = thread_store.get::<SecurityRiskScore>() {
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::<serde_json::Value>(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(
&registry,
@@ -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::<SecurityRiskScore>() {
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(
&registry,
@@ -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(
&registry,
@@ -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(
&registry,
@@ -2286,13 +2269,16 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> {
let progress = thread_store
.get::<GuardianV2ScoreProgress>()
.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(
&registry,
@@ -2333,7 +2319,7 @@ async fn contributor_skips_required_models_in_standard_scope() -> Result<()> {
.await;
assert!(
thread_store.get::<SecurityRiskScore>().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::<GuardianV2ScoreProgress>().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::<GuardianV2ScoreProgress>()
.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::<GuardianV2ScoreProgress>()
.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(
&registry,
@@ -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(
&registry,
@@ -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::<SecurityRiskScore>() {
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::<SecurityRiskScore>()
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::<SecurityRiskScore>() {
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::<GuardianV2ScoreProgress>().unwrap();
let mut score = store.get::<SecurityRiskScore>().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<GuardianScope> {
#[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::<GuardianV2ScoreProgress>().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(())
}

View File

@@ -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::<AnalyticsEventsClient>();
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::<SecurityRiskScore>();
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(),

View File

@@ -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<BTreeSet<String>>,
// Serialize successful score publication with its authorization metadata.
pub(super) authorization: Mutex<Option<ScoreAuthorization>>,
state: Mutex<ScoreState>,
pub(super) metrics: Option<Arc<dyn ExtensionMetrics>>,
}
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<SecurityRiskScore>,
wrapper_lag: WrapperLag,
latest_tool_call: usize,
js_executions: usize,
latest_scored_tool_call: usize,
latest_failed_tool_call: usize,
oversized_tool_calls: BTreeSet<String>,
authorization: Option<ScoreAuthorization>,
}
/// 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<f64>,
pub(super) authorization: Option<ScoreAuthorization>,
}
impl GuardianV2ScoreProgress {
pub(super) fn new(metrics: Option<Arc<dyn ExtensionMetrics>>) -> 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;

View File

@@ -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<SecurityRiskScore> {
let progress = store.get::<GuardianV2ScoreProgress>()?;
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::<SecurityRiskScore>().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::<SecurityRiskScore>().as_deref(),
cached_score(&thread_store).as_ref(),
Some(&fail_closed_score)
);
}

View File

@@ -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<VecDeque<ToolStart>>);
pub(super) struct WrapperLag {
starts: VecDeque<ToolStart>,
}
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)