mirror of
https://github.com/openai/codex.git
synced 2026-09-20 12:47:38 +00:00
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:
@@ -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(),
|
||||
|
||||
289
codex-rs/core/tests/suite/guardian_cached_score_tests.rs
Normal file
289
codex-rs/core/tests/suite/guardian_cached_score_tests.rs
Normal 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(())
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
@@ -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(¤t) {
|
||||
if cached.authorization.as_ref() != Some(¤t) {
|
||||
(
|
||||
GuardianReviewReason::AuthorizationChanged,
|
||||
"authorization_changed",
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
})
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
@@ -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(
|
||||
®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::<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(
|
||||
®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::<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(
|
||||
®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::<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(
|
||||
®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::<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(
|
||||
®istry,
|
||||
@@ -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(
|
||||
®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::<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(())
|
||||
}
|
||||
|
||||
@@ -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(),
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user