mirror of
https://github.com/openai/codex.git
synced 2026-10-21 17:48:44 +00:00
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:
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user