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

Bucket admitted Guardian request metrics by input size (#55226)

Browse Source 
## Why

Guardian request counts grouped by input size should reflect admitted sampling inputs without counting synchronous transport retries more than once.

## What changed

- Add a shared `token_bucket` label to admitted total-input observations in `codex.guardian.context.request_tokens` for synchronous and asynchronous reviews. Use inclusive upper bounds in 20k-token steps through 240k, followed by 256k and overflow buckets.
- Record synchronous request tokens after budget admission, once per sampling step, including the complete input. A fresh step counts again even when its input is identical.
- Preserve asynchronous component histograms for oversized classifications, but omit `token_bucket` so rejected inputs do not contribute to counts grouped by that label.

## Testing

Add coverage for bucket boundaries, synchronous retry deduplication and budget rejection, asynchronous retained-history reuse and reset, and admitted versus oversized classification metrics.

GitOrigin-RevId: fac2f102a4c48c5fd21d56bfa701b95a1ff852d0
This commit is contained in:
ioanbudea-OAI
2026-10-19 09:03:12 +00:00
committed by copyberry
parent 89461ea683
commit 99f617bf74
12 changed files with 337 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/core/src/client.rs
-9
View File
@@ -1861,9 +1861,6 @@ impl ResponsesHttpExecutor<'_> {
)?;
self.client
.prepare_response_items_for_request(&mut request.input);
if crate::guardian::is_basic_session_source(&self.client.state.session_source) {
crate::guardian::observe_guardian_request(session_telemetry, &request);
}
let request_session_telemetry =
session_telemetry_for_request(session_telemetry, &request);
options.extra_headers.extend(responses_headers);
@@ -2074,12 +2071,6 @@ impl ModelClientSession {
Err(err) => return Err(provider.map_api_error(err)),
}
// Measure the complete logical request, not only the websocket delta.
if !warmup
&& crate::guardian::is_basic_session_source(&self.client.state.session_source)
{
crate::guardian::observe_guardian_request(session_telemetry, &request);
}
let mut client_metadata = self.client.build_ws_client_metadata(
responses_metadata,
include_internal,
codex-rs/core/src/guardian/mod.rs
-1
View File
@@ -16,7 +16,6 @@ pub(crate) use input_budget::should_compact as should_compact_guardian_input;
pub(crate) use permissions::for_tool as tool_permission_context;
mod request_budget;
pub(crate) use request_budget::ExhaustedReviewBudget;
pub(crate) use request_budget::observe as observe_guardian_request;
pub(crate) use request_budget::prepare_prompt as prepare_guardian_prompt;
mod review;
mod review_session;
codex-rs/core/src/guardian/request_budget.rs
+27 -14
View File
@@ -1,5 +1,6 @@
//! Restores compacted retained evidence and checks complete synchronous requests.
//! Restored records are budgeted and persisted before sampling, including continuations.
//! Admitted sampling inputs are observed once per step, independent of transport retries.
use codex_api::ResponsesApiRequest;
use codex_context_fragments::set_annotated_content;
@@ -8,7 +9,6 @@ use codex_guardian_context::REQUEST_TOKENS_BOUNDARIES;
use codex_guardian_context::REQUEST_TOKENS_METRIC;
use codex_guardian_context::RequestBudget;
use codex_guardian_context::effective_input_token_limit;
use codex_otel::SessionTelemetry;
use codex_protocol::error::CodexErr;
use codex_protocol::error::Result as CodexResult;
use codex_protocol::models::ContentItem;
@@ -33,18 +33,8 @@ pub(crate) enum ExhaustedReviewBudget {
Compacting,
}
pub(crate) fn observe(telemetry: &SessionTelemetry, request: &ResponsesApiRequest) -> usize {
let total = estimate_request_tokens(request);
// The assembled input already includes inherited history and the current
// review. Do not report a guessed old/new split after context injection.
telemetry.histogram_with_boundaries(
REQUEST_TOKENS_METRIC,
i64::try_from(total).unwrap_or(i64::MAX),
REQUEST_TOKENS_BOUNDARIES,
&[("target", "sync"), ("component", "total")],
);
total
}
/// Guardian-owned marker scoped to one sampling step, including all transport retries.
struct ObservedRequest;
pub(super) fn estimate_request_tokens(request: &ResponsesApiRequest) -> usize {
let input = request
@@ -157,7 +147,8 @@ pub(crate) async fn prepare_prompt(
metadata,
/*include_internal*/ true,
)?;
if estimate_request_tokens(&request) > maximum {
let total = estimate_request_tokens(&request);
if total > maximum {
session
.services
.thread_extension_data
@@ -182,5 +173,27 @@ pub(crate) async fn prepare_prompt(
.services
.thread_extension_data
.remove::<ExhaustedReviewBudget>();
if step
.extension_data
.insert_if(ObservedRequest, |existing| existing.is_none())
{
step.session_telemetry.histogram_with_boundaries(
REQUEST_TOKENS_METRIC,
i64::try_from(total).unwrap_or(i64::MAX),
REQUEST_TOKENS_BOUNDARIES,
&[
("target", "sync"),
("component", "total"),
(
"token_bucket",
codex_guardian_context::request_token_bucket(total),
),
],
);
}
Ok(())
}
#[cfg(test)]
#[path = "request_budget_tests.rs"]
mod tests;
codex-rs/core/src/guardian/request_budget_tests.rs
+126
View File
@@ -0,0 +1,126 @@
//! Admitted complete inputs count once per sampling step, including retry preparation.
use super::*;
use codex_otel::MetricsClient;
use codex_otel::MetricsConfig;
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use pretty_assertions::assert_eq;
use std::collections::BTreeMap;
use std::sync::Arc;
#[tokio::test]
#[test_case::test_case(128_000; "admitted")]
#[test_case::test_case(1; "rejected")]
async fn request_histogram_counts_admitted_steps_without_counting_retry_preparation(window: i64) {
let metrics = MetricsClient::new(
MetricsConfig::in_memory(
"test",
"codex-core",
env!("CARGO_PKG_VERSION"),
InMemoryMetricExporter::default(),
)
.with_runtime_reader(),
)
.unwrap();
let (session, mut turn) = crate::session::tests::make_session_and_context().await;
turn.session_telemetry = turn.session_telemetry.clone().with_metrics(metrics.clone());
Arc::make_mut(&mut turn.config).model_context_window = Some(window);
let session = Arc::new(session);
let turn = Arc::new(turn);
let input = ["retained history", "new review"]
.map(|role| {
serde_json::from_value(serde_json::json!({
"type": "message",
"role": "user",
"content": [{
"type": "input_text",
"text": format!("{role}: {}", "x".repeat(48_000)),
}],
}))
.expect("valid user message")
})
.to_vec();
let mut sum = 0;
// Identical input in a fresh step must count again without an explicit reset.
for _ in 0..2 {
let step = session
.capture_step_context(
Arc::clone(&turn),
&tokio_util::sync::CancellationToken::new(),
)
.await
.unwrap();
let mut prompt = crate::session::turn::build_prompt(
input.clone(),
&step,
session.get_prompt_base_instructions().await,
/*incremental_tools*/ false,
);
let metadata = session
.responses_metadata(
&step,
crate::responses_metadata::CodexResponsesRequestKind::Turn,
)
.await;
let request = session
.services
.model_client
.build_responses_request(
&prompt,
&step.settings.model_info,
/*effort*/ None,
codex_protocol::config_types::ReasoningSummary::None,
/*service_tier*/ None,
&metadata,
/*include_internal*/ true,
)
.unwrap();
sum += estimate_request_tokens(&request);
// The retry loop invokes Guardian preparation again on the same step.
for _ in 0..2 {
let result = prepare_prompt(&session, &mut prompt, &step, &metadata).await;
if window > 1 {
result.unwrap();
} else {
assert!(matches!(
result.unwrap_err().details(),
codex_protocol::error::CodexErrorDetails::ContextWindowExceeded
));
}
}
}
let snapshot = metrics.snapshot().unwrap();
let mut histograms = BTreeMap::new();
for metric in snapshot
.scope_metrics()
.flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics)
.filter(|metric| metric.name() == REQUEST_TOKENS_METRIC)
{
let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else {
panic!("expected request histogram")
};
for point in histogram.data_points() {
let tags = point
.attributes()
.filter(|attr| matches!(attr.key.as_str(), "target" | "component" | "token_bucket"))
.map(|attr| (attr.key.as_str().to_owned(), attr.value.to_string()))
.collect::<BTreeMap<_, _>>();
histograms.insert(tags, (point.count(), point.sum()));
}
}
let expected = if window > 1 {
BTreeMap::from([(
BTreeMap::from([
("target".to_owned(), "sync".to_owned()),
("component".to_owned(), "total".to_owned()),
("token_bucket".to_owned(), "02_20k-40k".to_owned()),
]),
(2, sum as f64),
)])
} else {
BTreeMap::new()
};
assert_eq!(histograms, expected);
}
codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs
+1
View File
@@ -167,6 +167,7 @@ impl Reservation {
request.metrics.as_deref(),
prepared.existing_context_tokens,
prepared.input_tokens,
super::metrics::RequestAdmission::Admitted,
);
drop(request.evidence);
let cursor = prepared.cursor;
codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs
+39 -1
View File
@@ -89,13 +89,51 @@ async fn configured_reset_limit_rebuilds_history_without_rejecting_fresh_evidenc
ev_completed("score"),
]);
let mock = responses::mount_sse_sequence(&server, vec![events; 2]).await;
let metrics = Arc::new(RecordingMetrics::default());
for index in 0..2 {
let (request, score) = pending_request(&fixture, "parent").await?;
let (mut request, score) = pending_request(&fixture, "parent").await?;
request.metrics = Some(metrics.clone());
backend.reserve(|| index).1.unwrap().run(request).await;
assert_eq!(score.await??, "low");
}
let requests = mock.requests();
assert_eq!(requests.len(), 2);
let expected_buckets = requests
.iter()
.map(|request| {
let input: Vec<ResponseItem> =
serde_json::from_value(request.body_json()["input"].clone()).unwrap();
let tokens = input
.iter()
.map(codex_guardian_context::estimate_input_tokens)
.sum();
RecordedMetric::Histogram(
codex_guardian_context::REQUEST_TOKENS_METRIC.to_owned(),
i64::try_from(tokens).unwrap(),
vec![
("target".to_owned(), "async".to_owned()),
("component".to_owned(), "total".to_owned()),
(
"token_bucket".to_owned(),
codex_guardian_context::request_token_bucket(tokens).to_owned(),
),
],
)
})
.collect::<Vec<_>>();
let buckets = metrics
.0
.lock()
.unwrap()
.iter()
.filter(|sample| {
matches!(sample, RecordedMetric::Histogram(name, _, tags)
if name == codex_guardian_context::REQUEST_TOKENS_METRIC
&& tags.iter().any(|(key, _)| key == "token_bucket"))
})
.cloned()
.collect::<Vec<_>>();
assert_eq!(buckets, expected_buckets);
let input = requests[1].input();
assert_eq!(input.iter().any(|item| item["id"] == "score"), reused);
assert_eq!(
codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs
+16 -1
View File
@@ -145,10 +145,17 @@ pub(super) fn record_section_costs(
}
}
#[derive(Clone, Copy)]
pub(super) enum RequestAdmission {
Admitted,
Rejected,
}
pub(super) fn record_request_tokens(
metrics: Option<&dyn ExtensionMetrics>,
existing: usize,
total: usize,
admission: RequestAdmission,
) {
let Some(metrics) = metrics else {
return;
@@ -158,11 +165,19 @@ pub(super) fn record_request_tokens(
("new_input", total.saturating_sub(existing)),
("total", total),
] {
let mut tags = vec![("target", "async"), ("component", component)];
// Histogram Count grouped by token_bucket counts only admitted complete inputs.
if component == "total" && matches!(admission, RequestAdmission::Admitted) {
tags.push((
"token_bucket",
codex_guardian_context::request_token_bucket(total),
));
}
metrics.histogram_with_boundaries(
codex_guardian_context::REQUEST_TOKENS_METRIC,
i64::try_from(tokens).unwrap_or(i64::MAX),
codex_guardian_context::REQUEST_TOKENS_BOUNDARIES,
&[("target", "async"), ("component", component)],
&tags,
);
}
}
codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs
+8 -1
View File
@@ -298,13 +298,20 @@ impl LunaSampler {
.into_iter()
.map(ResponseItemEnvelope::into_item)
.collect();
let admission =
if total_tokens > self.config.max_input_tokens.saturating_sub(/*rhs*/ 256) {
super::metrics::RequestAdmission::Rejected
} else {
super::metrics::RequestAdmission::Admitted
};
super::metrics::record_request_tokens(
self.config.metrics.as_deref(),
/*existing*/ 0,
total_tokens,
admission,
);
// Oversized classifications defer to sync with the existing failure score.
if total_tokens > self.config.max_input_tokens.saturating_sub(/*rhs*/ 256) {
if matches!(admission, super::metrics::RequestAdmission::Rejected) {
return Err(LunaSamplerError::InputTooLarge);
}
let (supersede, superseded) = oneshot::channel();
codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs
+59 -8
View File
@@ -446,18 +446,69 @@ async fn sampler_records_token_usage_after_returning_an_early_classification() -
.iter()
.map(codex_guardian_context::estimate_input_tokens)
.sum::<usize>();
assert!(metrics.0.lock().unwrap().contains(&(
codex_guardian_context::REQUEST_TOKENS_METRIC.to_owned(),
i64::try_from(estimated)?,
vec![
("target".to_owned(), "async".to_owned()),
("component".to_owned(), "total".to_owned()),
],
)));
assert_eq!(
metrics
.0
.lock()
.unwrap()
.iter()
.filter(
|metric| metric.0 == codex_guardian_context::REQUEST_TOKENS_METRIC
&& metric.2.iter().any(|(key, _)| key == "token_bucket")
)
.cloned()
.collect::<Vec<_>>(),
vec![(
codex_guardian_context::REQUEST_TOKENS_METRIC.to_owned(),
i64::try_from(estimated)?,
vec![
("target".to_owned(), "async".to_owned()),
("component".to_owned(), "total".to_owned()),
(
"token_bucket".to_owned(),
codex_guardian_context::request_token_bucket(estimated).to_owned()
),
],
)],
);
Ok(())
}
#[tokio::test]
async fn oversized_classification_keeps_untagged_request_histograms() {
let metrics = Arc::new(RecordingMetrics::default());
let mut config = sampler_config("http://localhost".to_owned());
config.max_input_tokens = 1;
config.metrics = Some(metrics.clone());
let sampler = LunaSampler::new(config);
assert!(matches!(
sampler.sample(sample_request("parent")).await,
Err(LunaSamplerError::InputTooLarge)
));
assert_eq!(
metrics
.0
.lock()
.unwrap()
.iter()
.filter(|metric| metric.2.iter().any(|(key, _)| key == "token_bucket"))
.cloned()
.collect::<Vec<_>>(),
Vec::<RecordedMetric>::new(),
);
assert_eq!(
metrics
.0
.lock()
.unwrap()
.iter()
.filter(|metric| metric.0 == codex_guardian_context::REQUEST_TOKENS_METRIC)
.count(),
3,
);
}
struct RefreshableAuth(std::sync::Mutex<&'static str>);
impl ExternalAuth for RefreshableAuth {
fn resolve(&self) -> ExternalAuthFuture<'_, CodexAuth> {
codex-rs/guardian-context/src/budget.rs
+27
View File
@@ -30,6 +30,33 @@ pub const REQUEST_TOKENS_BOUNDARIES: &[f64] = &[
1_000_000.0,
2_000_000.0,
];
/// Returns a bounded, sortable label with 20k steps and a final 256k cutoff.
/// Upper bounds are inclusive; zero belongs to the first bucket and values above 256k to the last.
pub fn request_token_bucket(tokens: usize) -> &'static str {
const BOUNDARIES: [usize; 13] = [
20_000, 40_000, 60_000, 80_000, 100_000, 120_000, 140_000, 160_000, 180_000, 200_000,
220_000, 240_000, 256_000,
];
const LABELS: [&str; 14] = [
"01_0-20k",
"02_20k-40k",
"03_40k-60k",
"04_60k-80k",
"05_80k-100k",
"06_100k-120k",
"07_120k-140k",
"08_140k-160k",
"09_160k-180k",
"10_180k-200k",
"11_200k-220k",
"12_220k-240k",
"13_240k-256k",
"14_256k-max",
];
LABELS[BOUNDARIES.partition_point(|bound| tokens > *bound)]
}
/// One shared scale for the section metric's count, token and byte measurements.
/// Buckets belong to the metric name, not its measurement tag.
pub const SECTION_COST_BOUNDARIES: &[f64] = &[
codex-rs/guardian-context/src/budget_tests.rs
+33
View File
@@ -7,6 +7,39 @@ use crate::composition::user_message as message;
use codex_protocol::models::ImageReference;
use pretty_assertions::assert_eq;
#[test]
fn request_token_ranges_keep_boundary_values_in_the_lower_bucket() {
let tokens = [
0,
19_999,
20_000,
20_001,
40_000,
40_001,
240_000,
240_001,
256_000,
256_001,
usize::MAX,
];
assert_eq!(
tokens.map(request_token_bucket),
[
"01_0-20k",
"01_0-20k",
"01_0-20k",
"02_20k-40k",
"02_20k-40k",
"03_40k-60k",
"12_220k-240k",
"13_240k-256k",
"13_240k-256k",
"14_256k-max",
"14_256k-max"
]
);
}
#[test]
fn section_costs_keep_multimodal_payloads_separate() {
let context = ComposedContext {
codex-rs/guardian-context/src/lib.rs
+1
View File
@@ -73,6 +73,7 @@ pub use budget::SectionCost;
pub use budget::effective_input_token_limit;
pub use budget::estimate_input_tokens;
pub use budget::estimate_input_tokens_from_serialized_len;
pub use budget::request_token_bucket;
pub use cursor::TranscriptCursor;
pub use cursor::TranscriptMode;
pub use cursor::TranscriptSelection;