diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 981e015434..2b7e5cebe2 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -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, diff --git a/codex-rs/core/src/guardian/mod.rs b/codex-rs/core/src/guardian/mod.rs index a0f4dedddf..7602d6721e 100644 --- a/codex-rs/core/src/guardian/mod.rs +++ b/codex-rs/core/src/guardian/mod.rs @@ -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; diff --git a/codex-rs/core/src/guardian/request_budget.rs b/codex-rs/core/src/guardian/request_budget.rs index ff6ad0b5d9..a6850c9387 100644 --- a/codex-rs/core/src/guardian/request_budget.rs +++ b/codex-rs/core/src/guardian/request_budget.rs @@ -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::(); + 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; diff --git a/codex-rs/core/src/guardian/request_budget_tests.rs b/codex-rs/core/src/guardian/request_budget_tests.rs new file mode 100644 index 0000000000..0a0a049584 --- /dev/null +++ b/codex-rs/core/src/guardian/request_budget_tests.rs @@ -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::>(); + 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); +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs b/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs index bbca31cafc..5b2090edc6 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs @@ -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; diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs index 57a587bb09..916edfc209 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/conversation_tests.rs @@ -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 = + 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::>(); + 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::>(); + assert_eq!(buckets, expected_buckets); let input = requests[1].input(); assert_eq!(input.iter().any(|item| item["id"] == "score"), reused); assert_eq!( diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs index 7037b6fc5e..eefcdd12f7 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/metrics.rs @@ -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, ); } } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs index da8fc86f94..3b0d0541b0 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -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(); diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs index cf3e8343ea..cce5cb311f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler_tests.rs @@ -446,18 +446,69 @@ async fn sampler_records_token_usage_after_returning_an_early_classification() - .iter() .map(codex_guardian_context::estimate_input_tokens) .sum::(); - 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![( + 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::::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> { diff --git a/codex-rs/guardian-context/src/budget.rs b/codex-rs/guardian-context/src/budget.rs index 53ef168330..999fe5aba0 100644 --- a/codex-rs/guardian-context/src/budget.rs +++ b/codex-rs/guardian-context/src/budget.rs @@ -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] = &[ diff --git a/codex-rs/guardian-context/src/budget_tests.rs b/codex-rs/guardian-context/src/budget_tests.rs index 6f5e02fa3c..e1434db709 100644 --- a/codex-rs/guardian-context/src/budget_tests.rs +++ b/codex-rs/guardian-context/src/budget_tests.rs @@ -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 { diff --git a/codex-rs/guardian-context/src/lib.rs b/codex-rs/guardian-context/src/lib.rs index c86dc2a388..af09a72acc 100644 --- a/codex-rs/guardian-context/src/lib.rs +++ b/codex-rs/guardian-context/src/lib.rs @@ -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;