github/codex
4
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 00:18:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

Cancel Guardian V2 classifications that exceed the tool-call lag window (#53699)

Browse Source 
## Why

Pending classifications can outlive the `max_tool_call_lag` window in which their scores are reusable, keeping metadata waits or classifier transports active after the result becomes stale.

## What changed

- Expire pending classifications when tool-call progress exceeds the configured lag allowance, preserving the wrapper discount while nested code-mode actions remain active.
- Cancel metadata waits and sampling in snapshot and conversation modes, including snapshot response drains after an early score.
- Treat expiry as supersession, preserving cached scores and the last completed conversation checkpoint and transcript cursor.

## Testing

Add regression coverage for lag boundaries, active wrapper discounts, transport cancellation before and after early scores, expired requests skipping preparation, and conversation history and transcript deltas surviving cancellation.

GitOrigin-RevId: 892ff13a05fd2b1573cef9842b934d1fc865d951
This commit is contained in:
felixxia-oai
2026-10-13 15:25:12 +00:00
committed by copyberry
parent ffc1402633
commit f6aca7dc46
10 changed files with 423 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/ext/guardian-v2/src/async_scorer/classification.rs
+53 -13
View File
@@ -1,5 +1,7 @@
//! Builds context and samples a captured observation in the background.
//! Successful results require current authorization and review context and a newer sample timestamp.
//! Requests outside every active action's lag allowance stop as superseded, not failed.
//! Expiry also cancels metadata waits before transcript preparation, releasing reservations.
use std::collections::BTreeMap;
use std::sync::Arc;
@@ -49,6 +51,7 @@ use super::sampler::LunaSampler;
use super::sampler::LunaSamplerError;
use super::sampler::LunaSamplingRequest;
use super::score::GuardianV2ScoreProgress;
use super::score::PendingScore;
use super::transcript::CollectedTranscript;
use super::transcript::ContextInput;
use super::truncation::ClassificationTruncations;
@@ -162,18 +165,44 @@ impl Classification {
let generation = reservation
.as_ref()
.map(|reservation| reservation.generation.clone());
let pending_score = PendingScore {
progress: Arc::clone(&score_progress),
index: tool_call_index,
max_lag: guardian_config.max_tool_call_lag,
};
let metadata = tokio::select! {
biased;
() = pending_score.expired() => None,
metadata = async {
let trusted_tool_context = match mcp_tool.as_ref() {
Some(tool) => {
trusted_tool_context(tool.tool_info(), tool.source(), &manager, &config).await
}
None => None,
};
let root_snapshot = if context_mode != GuardianContextMode::Legacy {
root_snapshot
} else {
thread.guardian_root_snapshot().await
};
(trusted_tool_context, root_snapshot)
} => Some(metadata),
};
// Metadata can become ready in the same poll as expiry. Do not start
// synchronous transcript preparation for an observation that is already stale.
let Some((trusted_tool_context, root_snapshot)) =
metadata.filter(|_| !pending_score.is_expired())
else {
record_classification(
metrics.as_deref(),
context_mode,
classification_started_at.elapsed(),
"superseded",
/*failure_reason*/ None,
);
return;
};
let mut truncations = ClassificationTruncations::default();
let trusted_tool_context = match mcp_tool.as_ref() {
Some(tool) => {
trusted_tool_context(tool.tool_info(), tool.source(), &manager, &config).await
}
None => None,
};
let root_snapshot = if context_mode != GuardianContextMode::Legacy {
root_snapshot
} else {
thread.guardian_root_snapshot().await
};
let mut trusted_skills = TrustedSkillInvocations::default();
for path in local_trusted_skill_paths.iter().chain(
root_snapshot
@@ -324,9 +353,11 @@ impl Classification {
reasoning_effort: guardian_config.reasoning_effort.clone(),
parent_turn_id: turn_id.clone(),
root_turn_id,
pending_score: None,
};
let result = match transcript {
ClassificationContext::Snapshot(context) => {
sampling.pending_score = Some(pending_score.clone());
sampling.input = context.into_messages();
if let Some(decisions_sampler) = decisions_sampler.as_ref() {
decisions_task =
@@ -342,7 +373,11 @@ impl Classification {
);
}
let sampling_started = Instant::now();
let result = sampler.sample(sampling).await;
let result = tokio::select! {
biased;
() = pending_score.expired() => Err(LunaSamplerError::Superseded),
result = sampler.sample(sampling) => result,
};
responses_duration = Some(sampling_started.elapsed());
result
}
@@ -372,8 +407,13 @@ impl Classification {
authorization: score_authorization.clone(),
thread: Arc::clone(&thread),
metrics: metrics.clone(),
pending_score: pending_score.clone(),
}));
score.await.unwrap_or(Err(LunaSamplerError::Superseded))
tokio::select! {
biased;
() = pending_score.expired() => Err(LunaSamplerError::Superseded),
result = score => result.unwrap_or(Err(LunaSamplerError::Superseded)),
}
}
};
let output = match result {
codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs
+9 -1
View File
@@ -6,6 +6,7 @@
//! Oversized retained requests rebuild from fresh evidence before the model's hard limit.
//! Each request owns its fork; failed or cancelled requests leave committed history intact.
//! Host-only instruction delivery metadata is committed with its completed input messages.
//! Expired classifications release their transport without committing partial history.
use std::sync::Arc;
use std::sync::Mutex;
@@ -25,6 +26,7 @@ use super::authorization::ScoreAuthorization;
use super::sampler::LunaSampler;
use super::sampler::LunaSamplerError;
use super::sampler::LunaSamplingRequest;
use super::score::PendingScore;
use super::transcript::CollectedTranscript;
const MAX_OUTSTANDING: usize = 16;
@@ -58,6 +60,7 @@ pub(super) struct ConversationRequest {
pub(super) thread: Arc<CodexThread>,
pub(super) ready: oneshot::Sender<Result<String, LunaSamplerError>>,
pub(super) metrics: Option<Arc<dyn ExtensionMetrics>>,
pub(super) pending_score: PendingScore,
}
// Authorization covers model settings; local config changes replace the backend.
@@ -108,6 +111,7 @@ impl Reservation {
pub(super) async fn run(self, request: ConversationRequest) {
if self.generation.strong_count() == 0
|| !request.authorization.is_current(&request.thread).await
|| request.pending_score.is_expired()
{
let _ = request.ready.send(Err(LunaSamplerError::Superseded));
return;
@@ -171,7 +175,11 @@ impl Reservation {
drop(request.evidence);
let cursor = prepared.cursor;
let pending_truncations = std::mem::take(&mut prepared.truncations);
let history = self.sampler.sample_retained(prepared, request.ready).await;
let history = tokio::select! {
biased;
() = request.pending_score.expired() => None,
history = self.sampler.sample_retained(prepared, request.ready) => history,
};
if let Some(history) = history
&& self.generation.strong_count() > 0
&& request.authorization.is_current(&request.thread).await
codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs
+132 -2
View File
@@ -7,6 +7,7 @@ use crate::async_scorer::sampler::LunaSamplerError;
use crate::async_scorer::sampler::LunaSamplingRequest;
use crate::async_scorer::sampler::tests::sample_request;
use crate::async_scorer::sampler::tests::sampler_config;
use crate::async_scorer::score::PendingScore;
use crate::async_scorer::transcript::ContextInput;
use codex_extension_api::ThreadStopInput;
use codex_guardian_context::ContextTarget;
@@ -45,6 +46,14 @@ async fn pending_request(
})?;
let (ready, score) = tokio::sync::oneshot::channel();
let authorization = ScoreAuthorization::current(&thread, &Default::default()).await;
let pending_score = PendingScore {
progress: thread
.thread_extension_data()
.get::<GuardianV2ScoreProgress>()
.unwrap(),
index: 0,
max_lag: config.max_tool_call_lag,
};
Ok((
ConversationRequest {
evidence,
@@ -53,6 +62,7 @@ async fn pending_request(
thread,
ready,
metrics: None,
pending_score,
sampling: LunaSamplingRequest {
instructions,
input: Vec::new(),
@@ -242,14 +252,17 @@ async fn delayed_preparation_does_not_block_newer_requests_and_capacity_is_bound
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stale_authorization_and_stopped_generations_cannot_start() -> Result<()> {
async fn stale_authorization_expiry_and_stopped_generations_cannot_start() -> Result<()> {
skip_if_no_network!(Ok(()));
let server = responses::start_mock_server().await;
let fixture = GuardianFailureFixture::with_config(
"[features.guardianv2]\nasync_classifier_mode = 'conversation'",
)
.await?;
let sampler = Arc::new(LunaSampler::new(sampler_config(server.uri())));
// Preparation would fail with InputTooLarge if it ran before the expiry check.
let mut config = sampler_config(server.uri());
config.max_input_tokens = 0;
let sampler = Arc::new(LunaSampler::new(config));
fixture
.test
.codex
@@ -257,6 +270,13 @@ async fn stale_authorization_and_stopped_generations_cannot_start() -> Result<()
.insert(ConversationBackend::new(sampler));
let thread_store = fixture.test.codex.thread_extension_data();
let backend = thread_store.get::<ConversationBackend>().unwrap();
let expired = backend.reserve(|| 0).1.unwrap();
let (mut request, score) = pending_request(&fixture, "expired").await?;
request.pending_score.max_lag = 0;
observe_unscored_call(&request.pending_score.progress, thread_store);
expired.run(request).await;
assert!(matches!(score.await?, Err(LunaSamplerError::Superseded)));
let blocker = backend.reserve(|| 0).1.unwrap();
let stale = backend.reserve(|| 0).1.unwrap();
let (request, score) = pending_request(&fixture, "stale").await?;
@@ -458,3 +478,113 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul
}
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn expiry_preserves_the_completed_checkpoint_and_transcript_cursor() -> Result<()> {
skip_if_no_network!(Ok(()));
let fixture = GuardianFailureFixture::with_config(
"[features.guardianv2]\nasync_classifier_mode = 'conversation'",
)
.await?;
let store = fixture.test.codex.thread_extension_data();
let progress = store.get::<GuardianV2ScoreProgress>().unwrap();
let (release_expired, finish_expired) = tokio::sync::oneshot::channel();
let (server, _completions) = start_streaming_sse_server(vec![
vec![StreamingSseChunk {
gate: None,
body: responses::sse(vec![
ev_assistant_message("checkpoint", "low"),
ev_completed("checkpoint"),
]),
}],
vec![
StreamingSseChunk {
gate: None,
body: responses::sse(vec![
responses::ev_response_created("expired"),
responses::ev_output_text_delta("low"),
ev_assistant_message("uncommitted", "low"),
]),
},
StreamingSseChunk {
gate: Some(finish_expired),
body: responses::sse(vec![ev_completed("expired")]),
},
],
vec![StreamingSseChunk {
gate: None,
body: responses::sse(vec![
ev_assistant_message("resumed", "low"),
ev_completed("resumed"),
]),
}],
])
.await;
let backend = ConversationBackend::new(Arc::new(LunaSampler::new(sampler_config(format!(
"{}/v1",
server.uri()
)))));
let (mut request, score) = pending_request(&fixture, "checkpoint").await?;
let (index, reservation) = backend.reserve(|| observe_unscored_call(&progress, store));
request.pending_score.index = index;
reservation.unwrap().run(request).await;
assert_eq!(score.await??, "low");
// This evidence is beyond the committed cursor. A cancelled review must not
// advance that cursor or cause its evidence to disappear from the next delta.
fixture
.test
.codex
.inject_response_items(vec![serde_json::from_value(
responses::ev_function_call(
"observed-after-checkpoint",
"read_file",
r#"{"path":"evidence-after-checkpoint"}"#,
)["item"]
.clone(),
)?])
.await?;
let (mut request, early_score) = pending_request(&fixture, "expired").await?;
let (index, reservation) = backend.reserve(|| observe_unscored_call(&progress, store));
request.pending_score.index = index;
let expired = tokio::spawn(reservation.unwrap().run(request));
assert_eq!(early_score.await??, "low");
// Keep completion gated even after the early score, then expire the request.
for _ in 0..3 {
observe_unscored_call(&progress, store);
}
tokio::time::timeout(ASYNC_TEST_TIMEOUT, expired).await??;
let (mut request, score) = pending_request(&fixture, "resumed").await?;
let (index, reservation) = backend.reserve(|| observe_unscored_call(&progress, store));
request.pending_score.index = index;
reservation.unwrap().run(request).await;
assert_eq!(
tokio::time::timeout(ASYNC_TEST_TIMEOUT, score).await???,
"low"
);
let requests = server.requests().await;
assert_eq!(requests.len(), 3);
let checkpoint: serde_json::Value = serde_json::from_slice(&requests[0])?;
let resumed: serde_json::Value = serde_json::from_slice(&requests[2])?;
let prefix = checkpoint["input"].as_array().unwrap();
let input = resumed["input"].as_array().unwrap();
assert_eq!(&input[..prefix.len()], prefix);
assert_eq!(
input
.iter()
.filter(|item| item["role"] == "assistant")
.filter_map(|item| item["id"].as_str())
.collect::<Vec<_>>(),
["checkpoint"],
);
let delta = serde_json::to_string(&input[prefix.len()..])?;
assert!(delta.contains("TRANSCRIPT DELTA START"));
assert_eq!(delta.matches("evidence-after-checkpoint").count(), 1);
drop(release_expired);
drop(backend);
fixture.test.codex.shutdown_and_wait().await?;
server.shutdown().await;
Ok(())
}
codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs
+28
View File
@@ -74,6 +74,8 @@ use crate::async_scorer::authorization::ScoreAuthorization;
#[path = "conversation_tests.rs"]
mod conversation;
#[path = "pending_score_tests.rs"]
mod pending_score;
use crate::async_scorer::config::CLASSIFICATION_OUTPUT_INSTRUCTIONS;
use crate::async_scorer::config::DEFAULT_PARENT_COMPACTION_TOKENS;
use crate::async_scorer::config::GuardianV2Config;
@@ -3543,6 +3545,11 @@ async fn cached_approval_respects_action_order_and_wrapper_lag() -> Result<()> {
};
let origin = ResponseItemId::from_server("wrapper".to_owned());
let wrapper = start("wrapper", &origin, ToolCallSource::Direct);
let pending = crate::async_scorer::score::PendingScore {
progress: Arc::clone(&progress),
index: wrapper - 1,
max_lag: 2,
};
for (call_id, expected) in [
("first", Some(ReviewDecision::Approved)),
("second", Some(ReviewDecision::Approved)),
@@ -3550,6 +3557,27 @@ async fn cached_approval_respects_action_order_and_wrapper_lag() -> Result<()> {
] {
start(call_id, &origin, nested.clone());
assert_eq!(approve(call_id).await, expected);
let expiration = pending.expired();
tokio::pin!(expiration);
if expected.is_some() {
tokio::select! {
biased;
() = &mut expiration => panic!("a usable wrapper-discounted score expired"),
() = std::future::ready(()) => {},
}
} else {
tokio::time::timeout(ASYNC_TEST_TIMEOUT, &mut expiration).await?;
}
if call_id == "second" {
progress.finish("first");
tokio::select! {
biased;
() = &mut expiration => panic!("the second nested action can still reuse the score"),
() = std::future::ready(()) => {},
}
progress.finish("second");
tokio::time::timeout(ASYNC_TEST_TIMEOUT, &mut expiration).await?;
}
}
// A score covering the wrapper must not receive another discount.
seed_cached_score(
codex-rs/ext/guardian-v2/src/async_scorer/pending_score_tests.rs
+103
View File
@@ -0,0 +1,103 @@
//! Exercises lag expiry through tool admission and real classifier transports.
use super::*;
use crate::async_scorer::conversation::ConversationBackend;
use crate::async_scorer::sampler::tests::sampler_config;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn stale_classifications_close_the_transport_without_invalidating_scores() -> Result<()> {
skip_if_no_network!(Ok(()));
for mode in ["snapshot", "conversation"] {
for (max_lag, early_score) in [(0, false), (2, false), (4, false), (2, true)] {
let fixture = GuardianFailureFixture::with_config(&format!(
"[features.guardianv2]\nasync_classifier_mode = '{mode}'\nmax_tool_call_lag = {max_lag}"
))
.await?;
let mut events = vec![responses::ev_response_created("stale")];
if early_score {
events.push(responses::ev_output_text_delta("low"));
}
let server =
responses::start_websocket_server_with_headers(vec![WebSocketConnectionConfig {
// Rejected interruption exercises the close fallback.
requests: vec![
events,
vec![json!({
"type": "response.interrupt.failed",
"response_id": "stale",
"error": {"code": "interrupt_not_supported"},
})],
// Keep reading until the client closes; no third request is expected.
Vec::new(),
],
response_headers: Vec::new(),
accept_delay: None,
close_after_requests: true,
}])
.await;
let store = fixture.test.codex.thread_extension_data();
store.insert(LunaSampler::new(sampler_config(
proxy_websocket_servers_with_http(
&[&server],
ProxyPrewarmLimit::AllConnections,
/*http_url*/ None,
)
.await?,
)));
let sampler = store.get::<LunaSampler>().unwrap();
sampler.prewarm().await;
if mode == "conversation" {
store.insert(ConversationBackend::new(Arc::clone(&sampler)));
}
let progress = store.get::<GuardianV2ScoreProgress>().unwrap();
let mut updates = progress.updates.subscribe();
fixture.score_tool(ToolName::plain("read_file")).await;
if early_score {
tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
while progress.inspect(/*call_id*/ None).action_risk != Some(0.0) {
updates
.changed()
.await
.expect("score progress remains live");
}
})
.await?;
}
tokio::time::timeout(
ASYNC_TEST_TIMEOUT,
server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0),
)
.await?;
let score = cached_score(store);
let mut expected = progress.inspect(/*call_id*/ None);
for _ in 0..max_lag {
observe_unscored_call(&progress, store);
}
// The running classifier must not send an interrupt at the reusable boundary.
assert!(
tokio::time::timeout(
Duration::from_millis(250),
server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 1),
)
.await
.is_err(),
"classifier interrupted inside its configured lag window"
);
observe_unscored_call(&progress, store);
assert!(
server
.wait_for_closed_connections(/*expected*/ 1, ASYNC_TEST_TIMEOUT)
.await
);
assert_eq!(server.single_connection().len(), 2);
expected.lag += max_lag + 1;
assert_eq!(progress.inspect(/*call_id*/ None), expected);
assert_eq!(cached_score(store), score);
fixture.test.codex.shutdown_and_wait().await?;
server.shutdown().await;
}
}
Ok(())
}
codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs
+24 -1
View File
@@ -6,6 +6,7 @@ mod execution;
mod response;
use super::request::PreparedRequest;
use super::score::PendingScore;
use connection_pool::ConnectionPool;
use std::collections::VecDeque;
use std::sync::Arc;
@@ -93,6 +94,8 @@ pub struct LunaSamplingRequest {
pub parent_turn_id: String,
/// Trusted causal root of the owning turn, absent when unknown or ambiguous.
pub root_turn_id: Option<String>,
/// Snapshot expiry remains active until its response finishes draining.
pub(super) pending_score: Option<PendingScore>,
}
/// Failures returned while connecting or sampling the Luna model.
@@ -113,7 +116,7 @@ pub enum LunaSamplerError {
/// The response exceeded the bounded output limit.
#[error("Luna response exceeded the output limit")]
OutputTooLarge,
/// A newer classification replaced this request when the pool was full.
/// The classification expired, lost its owning generation, or was replaced at capacity.
#[error("Luna request was superseded by a newer classification")]
Superseded,
/// The supplied parent checkpoint cannot be consumed by this Luna configuration.
@@ -272,6 +275,14 @@ impl LunaSampler {
&self,
mut request: LunaSamplingRequest,
) -> Result<String, LunaSamplerError> {
if request
.pending_score
.as_ref()
.is_some_and(PendingScore::is_expired)
{
return Err(LunaSamplerError::Superseded);
}
let pending_score = request.pending_score.take();
let auth_owner_generation = self.auth_owner_generation();
if request.parent_compaction.is_some()
&& !self.supports_parent_compaction(request.parent_compaction_hash.as_deref())
@@ -318,6 +329,18 @@ impl LunaSampler {
scored: Arc::clone(&scored),
});
}
// The same cancellation future moves into the detached drain after an early score.
let superseded = Box::pin(async move {
tokio::select! {
_ = superseded => {},
() = async {
match pending_score {
Some(pending) => pending.expired().await,
None => std::future::pending().await,
}
} => {},
}
});
self.execution(&request, input, auth_owner_generation)
.run(execution::SamplingMode::Snapshot { superseded, scored })
.await
codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs
+6 -4
View File
@@ -1,6 +1,6 @@
//! Executes one prepared classifier request with the existing retry and streaming rules.
//! Snapshot output returns immediately with a detached drain. Retained requests publish
//! the same early result but keep ownership until completion to capture reusable output.
//! Snapshot output returns immediately with a detached drain that retains cancellation.
//! Retained requests publish the same early result but keep ownership through completion.
//! Requests and retries stop when the account owner that started the classification changes.
//! Dropped responses interrupt and drain before their transport can return to the pool.
@@ -25,6 +25,8 @@ use codex_protocol::turn_input::CyberAccessProgram;
use http::StatusCode;
use serde_json::json;
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
@@ -57,7 +59,7 @@ pub(super) struct RetainedCompletion {
/// Snapshot admission owns supersession; retained admission owns the response through completion.
pub(super) enum SamplingMode<'a> {
Snapshot {
superseded: oneshot::Receiver<()>,
superseded: Pin<Box<dyn Future<Output = ()> + Send>>,
scored: Arc<AtomicBool>,
},
Retained(&'a mut RetainedCompletion),
@@ -67,7 +69,7 @@ impl SamplingMode<'_> {
async fn superseded(&mut self) {
match self {
Self::Snapshot { superseded, .. } => {
let _ = superseded.await;
superseded.await;
}
Self::Retained(_) => std::future::pending().await,
}
codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs
+4
View File
@@ -281,6 +281,7 @@ pub(in crate::async_scorer) fn sample_request(parent_turn_id: &str) -> LunaSampl
reasoning_effort: ReasoningEffort::None,
parent_turn_id: parent_turn_id.to_owned(),
root_turn_id: None,
pending_score: None,
}
}
@@ -649,6 +650,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications
reasoning_effort: ReasoningEffort::None,
parent_turn_id: "turn-1".to_owned(),
root_turn_id: Some("turn-1".to_owned()),
pending_score: None,
})
.await?;
tokio::time::timeout(Duration::from_secs(2), async {
@@ -682,6 +684,7 @@ async fn preconnected_sampler_reuses_authenticated_websocket_for_classifications
reasoning_effort: ReasoningEffort::Medium,
parent_turn_id: "turn-2".to_owned(),
root_turn_id: Some("root-2".to_owned()),
pending_score: None,
})
.await?;
@@ -877,6 +880,7 @@ async fn sampler_returns_classification_token_before_terminal_response_events()
reasoning_effort: ReasoningEffort::None,
parent_turn_id: "turn-1".to_owned(),
root_turn_id: Some("turn-1".to_owned()),
pending_score: None,
}),
)
.await??;
codex-rs/ext/guardian-v2/src/async_scorer/score.rs
+41 -4
View File
@@ -1,6 +1,7 @@
//! 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.
//! Pending classifications expire only when no active action can reuse their score.
use std::collections::BTreeMap;
use std::collections::BTreeSet;
@@ -19,9 +20,42 @@ use super::wrapper_lag::WrapperLag;
pub(super) struct GuardianV2ScoreProgress {
state: Mutex<ScoreState>,
pub(super) updates: tokio::sync::watch::Sender<()>,
observations: tokio::sync::watch::Sender<()>,
pub(super) metrics: Option<Arc<dyn ExtensionMetrics>>,
}
#[derive(Clone)]
pub(super) struct PendingScore {
pub(super) progress: Arc<GuardianV2ScoreProgress>,
pub(super) index: usize,
pub(super) max_lag: usize,
}
impl PendingScore {
pub(super) fn is_expired(&self) -> bool {
let state = self
.progress
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
state
.latest_tool_call
.saturating_sub(self.index)
.saturating_sub(state.wrapper_lag.pending_discount(self.index))
> self.max_lag
}
pub(super) async fn expired(&self) {
// Subscribe before checking so admission cannot race past the expiry check.
let mut observations = self.progress.observations.subscribe();
while !self.is_expired() {
if observations.changed().await.is_err() {
return;
}
}
}
}
#[derive(Default)]
struct ScoreState {
score: Option<SecurityRiskScore>,
@@ -63,6 +97,7 @@ impl GuardianV2ScoreProgress {
state.latest_tool_call = state.latest_tool_call.saturating_add(/*rhs*/ 1);
let index = state.latest_tool_call;
state.wrapper_lag.record(input, index);
self.observations.send_replace(());
index
}
@@ -94,11 +129,13 @@ impl GuardianV2ScoreProgress {
pub(super) fn finish(&self, call_id: &str) {
// Overflow belongs to each active call even after a newer score succeeds.
self.state
let mut state = self
.state
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.oversized_tool_calls
.remove(call_id);
.unwrap_or_else(std::sync::PoisonError::into_inner);
state.oversized_tool_calls.remove(call_id);
state.wrapper_lag.finish(call_id);
self.observations.send_replace(());
}
pub(super) fn publish(
codex-rs/ext/guardian-v2/src/async_scorer/wrapper_lag.rs
+23
View File
@@ -1,6 +1,7 @@
//! Tracks action order and discounts an approval's own unscored code-mode wrapper.
//! Missing provenance keeps the full lag; bounded history never changes scoring or failure order.
//! The score owner's state lock protects this history.
//! Active nested actions retain their wrapper allowance when expiring pending scores.
use std::collections::VecDeque;
@@ -14,6 +15,7 @@ pub(super) struct ToolStart {
pub(super) index: usize,
wrapper_item_id: Option<ResponseItemId>,
parent_wrapper_index: Option<usize>,
finished: bool,
}
#[derive(Default)]
@@ -51,9 +53,30 @@ impl WrapperLag {
index,
wrapper_item_id,
parent_wrapper_index,
finished: false,
});
}
pub(super) fn finish(&mut self, call_id: &str) {
if let Some(start) = self
.starts
.iter_mut()
.rev()
.find(|start| start.call_id == call_id)
{
start.finished = true;
}
}
pub(super) fn pending_discount(&self, index: usize) -> usize {
usize::from(self.starts.iter().any(|start| {
!start.finished
&& start
.parent_wrapper_index
.is_some_and(|parent| parent > index)
}))
}
pub(super) fn get(&self, call_id: Option<&str>) -> Option<&ToolStart> {
let call_id = call_id?;
self.starts