Refactor Guardian v2 Decisions snapshot comparison (#55623)

## What changed

- Rename the snapshot comparison entry point to `start_snapshot_comparison` and extract request spawning into `spawn_request`.
- Accept a `ContentItem` rubric and an iterator of evidence in the request adapter.
- Read the Decisions comparison feature flag once per classification and simplify the budget rejection test setup.

GitOrigin-RevId: b4a294eeecb6069db0f80ff6e4f13f558cbd7815
This commit is contained in:
swastik-oai
2026-10-20 16:31:41 +00:00
committed by copyberry
parent 5ac9888436
commit aa081afb1b
4 changed files with 39 additions and 63 deletions
@@ -313,6 +313,9 @@ impl Classification {
truncations.extend(std::mem::take(&mut context.truncations));
super::metrics::record_section_costs(metrics.as_deref(), context.section_costs());
}
let decisions_comparison_enabled = config
.features
.enabled(codex_features::Feature::GuardianV2DecisionsComparison);
let mut decisions_task = None;
let mut responses_duration = None;
let mut failure_reason = "invalid_output";
@@ -363,12 +366,11 @@ impl Classification {
sampling.pending_score = Some(pending_score.clone());
sampling.input = context.into_messages();
if let Some(decisions_sampler) = decisions_sampler.as_ref() {
decisions_task =
Some(decisions_sampler.spawn(&sampling, sampler.max_input_tokens()));
} else if config
.features
.enabled(codex_features::Feature::GuardianV2DecisionsComparison)
{
decisions_task = Some(decisions_sampler.start_snapshot_comparison(
&sampling,
sampler.max_input_tokens(),
));
} else if decisions_comparison_enabled {
super::metrics::record_decisions_comparison_outcome(
metrics.as_deref(),
"skipped",
@@ -390,10 +392,7 @@ impl Classification {
reservation,
} => {
// The retained prefix is owned by the conversation backend, not this snapshot.
if config
.features
.enabled(codex_features::Feature::GuardianV2DecisionsComparison)
{
if decisions_comparison_enabled {
super::metrics::record_decisions_comparison_outcome(
metrics.as_deref(),
"skipped",
@@ -554,9 +553,7 @@ impl Classification {
baseline_risk.zip(responses_duration),
metrics.as_deref(),
);
} else if config
.features
.enabled(codex_features::Feature::GuardianV2DecisionsComparison)
} else if decisions_comparison_enabled
&& let Some(risk) = baseline_risk
&& let Some(metrics) = metrics.as_deref()
{
@@ -6,7 +6,6 @@
use super::sampler::LunaSampler;
use super::sampler::LunaSamplingRequest;
use codex_context_fragments::RenderedFragment;
use codex_history::ResponseItemEnvelope;
use codex_http_client::HttpClient;
use codex_protocol::models::ContentItem;
@@ -117,15 +116,14 @@ impl DecisionsSampler {
})
}
/// Retain the newest requests; the caller joins after baseline publication.
pub(super) fn spawn(
pub(super) fn start_snapshot_comparison(
self: &Arc<Self>,
request: &LunaSamplingRequest,
max_input_tokens: usize,
) -> DecisionsTask {
let started = Instant::now();
let body = request_body(
&request.instructions,
request.instructions.annotated_content().content(),
&request.input,
request.parent_compaction.as_ref(),
)
@@ -147,6 +145,14 @@ impl DecisionsSampler {
}
Ok(body)
});
self.spawn_request(body, started)
}
fn spawn_request(
self: &Arc<Self>,
body: Result<Value, DecisionsError>,
started: Instant,
) -> DecisionsTask {
let body = match body {
Ok(body) => body,
Err(error) => {
@@ -212,15 +218,15 @@ impl DecisionsSampler {
}
}
fn request_body(
instructions: &RenderedFragment,
evidence: &[ResponseItem],
fn request_body<'a>(
instructions: &ContentItem,
evidence: impl IntoIterator<Item = &'a ResponseItem>,
parent_compaction: Option<&ResponseItem>,
) -> Result<Value, DecisionsError> {
if parent_compaction.is_some() {
return Err(DecisionsError::UnsupportedEvidence);
}
let ContentItem::InputText { text: rubric } = instructions.annotated_content().content() else {
let ContentItem::InputText { text: rubric } = instructions else {
return Err(DecisionsError::UnsupportedEvidence);
};
let mut image_bytes = 0usize;
@@ -200,25 +200,8 @@ async fn guardian_budget_rejects_both_backends_before_sending() -> Result<()> {
.await?;
let server = responses::start_mock_server().await;
let store = fixture.test.codex.thread_extension_data();
let mut config = fixture.test.config.clone();
config.features.enable(Feature::GuardianApproval)?;
config.features.enable(Feature::GuardianV2)?;
config
.features
.enable(Feature::GuardianV2DecisionsComparison)?;
let metrics = Arc::new(RecordingMetrics::default());
fixture.registry.thread_lifecycle_contributors()[0]
.on_thread_start(ThreadStartInput {
config: &config,
session_source: &SessionSource::Exec,
persistent_thread_state_available: false,
environments: &[],
mcp_resource_client: None,
extension_metrics: Some(metrics.clone()),
session_store: &fixture.session_store,
thread_store: store,
})
.await;
store.insert(GuardianV2ScoreProgress::new(Some(metrics.clone())));
let mut sampler_config = crate::async_scorer::sampler::tests::sampler_config(server.uri());
sampler_config.max_input_tokens = 256;
store.insert(LunaSampler::new(sampler_config));
@@ -227,7 +210,6 @@ async fn guardian_budget_rejects_both_backends_before_sending() -> Result<()> {
"synthetic-key".into(),
server.uri(),
)?);
metrics.0.lock().unwrap().clear();
fixture.score_tool(ToolName::plain("read_file")).await;
fixture.assert_fails_closed("elevated_risk").await?;
tokio::time::timeout(ASYNC_TEST_TIMEOUT, async {
@@ -2,6 +2,7 @@ use super::super::startup::decisions_api_key;
use super::*;
use codex_context_fragments::AnnotatedContent;
use codex_context_fragments::ContextualUserFragment;
use codex_context_fragments::RenderedFragment;
use codex_guardian_context::TrustedTool;
use codex_guardian_context::TrustedToolSource;
use codex_http_client::HttpClientBuilder;
@@ -21,14 +22,9 @@ use wiremock::matchers::path;
#[test]
fn adapter_rejects_unsupported_evidence() {
let rubric = RenderedFragment::new(
"developer",
AnnotatedContent::text(
"Classify risk.",
ContentItemKind("guardian.test".into()),
ContentItemMetadata::harness(),
),
);
let rubric = ContentItem::InputText {
text: "Classify risk.".into(),
};
let user: ResponseItem = serde_json::from_value(json!({"type":"message", "role":"user", "content":[
{"type":"input_text", "text":"Authorized task"}, {"type":"input_text", "text":"Planned action"}
]})).unwrap();
@@ -66,14 +62,9 @@ fn adapter_rejects_unsupported_evidence() {
#[test]
fn adapter_preserves_trusted_tool_authority_and_scope() {
let rubric = RenderedFragment::new(
"developer",
AnnotatedContent::text(
"Classify risk.",
ContentItemKind("guardian.test".into()),
ContentItemMetadata::harness(),
),
);
let rubric = ContentItem::InputText {
text: "Classify risk.".into(),
};
let tool = TrustedTool {
server: "example".into(),
connector_id: None,
@@ -256,7 +247,7 @@ async fn http_contract_and_untrusted_response_validation() {
)
.unwrap();
let mut sampler = Arc::new(sampler);
let task = sampler.spawn(&request, /*max_input_tokens*/ 128_000);
let task = sampler.start_snapshot_comparison(&request, /*max_input_tokens*/ 128_000);
assert_eq!(task.finish().await.unwrap().0, Ok("low"));
let mut invalid = response;
invalid["answers"][0]["choice"] = json!("untrusted server text");
@@ -293,11 +284,11 @@ async fn decisions_admits_newest_request_by_cancelling_oldest() {
.acquire_many(MAX_CONCURRENT_REQUESTS as u32)
.await
.unwrap();
let oldest = sampler.spawn(&request, /*max_input_tokens*/ 128_000);
let oldest = sampler.start_snapshot_comparison(&request, /*max_input_tokens*/ 128_000);
let mut remaining = (1..MAX_CONCURRENT_REQUESTS)
.map(|_| sampler.spawn(&request, /*max_input_tokens*/ 128_000))
.map(|_| sampler.start_snapshot_comparison(&request, /*max_input_tokens*/ 128_000))
.collect::<Vec<_>>();
remaining.push(sampler.spawn(&request, /*max_input_tokens*/ 128_000));
remaining.push(sampler.start_snapshot_comparison(&request, /*max_input_tokens*/ 128_000));
assert!(oldest.finish().await.unwrap_err().is_cancelled());
assert!(server.received_requests().await.unwrap().is_empty());
drop(permits);
@@ -320,7 +311,7 @@ fn adapter_preserves_inline_images_and_bounds_image_bytes() {
.unwrap();
request.input = vec![image];
let body = request_body(
&request.instructions,
request.instructions.annotated_content().content(),
&request.input,
/*parent_compaction*/ None,
)
@@ -340,7 +331,7 @@ fn adapter_preserves_inline_images_and_bounds_image_bytes() {
];
assert_eq!(
request_body(
&request.instructions,
request.instructions.annotated_content().content(),
&request.input,
/*parent_compaction*/ None
),
@@ -373,7 +364,7 @@ async fn cancelling_a_decisions_request_releases_thread_capacity() {
.unwrap(),
);
let request = super::super::sampler::tests::sample_request("turn");
let task = sampler.spawn(&request, /*max_input_tokens*/ 128_000);
let task = sampler.start_snapshot_comparison(&request, /*max_input_tokens*/ 128_000);
tokio::time::timeout(Duration::from_secs(5), received.notified())
.await
.unwrap();