From 9dc3a1ce7ea53bb818b8dd69416d6dde70d0b680 Mon Sep 17 00:00:00 2001 From: ioanbudea-OAI Date: Mon, 19 Oct 2026 09:40:38 +0000 Subject: [PATCH] Record Guardian sampling retries, attempts, and delivery outcomes (#55234) ## What changed - Add sampling retry counters that distinguish transport retries from authentication recovery, with failure reasons. - Record attempt counts and final outcomes for successful, failed, cancelled, and superseded sampling requests. - Carry retained-score telemetry until the consumer accepts or discards the score, so queued scores that expire are recorded as superseded and accepted early scores remain successful even if response draining fails. ## Testing Add metric assertions for transient recovery, retry exhaustion, authentication recovery, cancellation, supersession, and retained-score acceptance or discard, including expiry while a score is queued. GitOrigin-RevId: 36babd6f431f22cf15b73355e6e026a256a81004 --- .../src/async_scorer/classification.rs | 2 +- .../src/async_scorer/conversation.rs | 6 +- .../src/async_scorer/conversation_tests.rs | 78 +++++++++- .../src/async_scorer/extension_tests.rs | 3 +- .../src/async_scorer/pending_score_tests.rs | 17 ++- .../async_scorer/retained_sampling_tests.rs | 39 ++++- .../guardian-v2/src/async_scorer/sampler.rs | 35 ++++- .../connection_pool_cancellation_tests.rs | 36 ++++- .../src/async_scorer/sampler/execution.rs | 19 ++- .../src/async_scorer/sampler/telemetry.rs | 122 ++++++++++++++++ .../src/async_scorer/sampler_tests.rs | 133 ++++++++++++++++-- 11 files changed, 448 insertions(+), 42 deletions(-) create mode 100644 codex-rs/ext/guardian-v2/src/async_scorer/sampler/telemetry.rs diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs index 9be3f246e3..cf46504867 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/classification.rs @@ -415,7 +415,7 @@ impl Classification { tokio::select! { biased; () = pending_score.expired() => Err(LunaSamplerError::Superseded), - result = score => result.unwrap_or(Err(LunaSamplerError::Superseded)), + result = score => result.unwrap_or(Err(LunaSamplerError::Superseded)).map(super::sampler::SampledScore::accept), } } }; 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 5b2090edc6..7a5d368e61 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/conversation.rs @@ -26,6 +26,7 @@ use super::authorization::ScoreAuthorization; use super::sampler::LunaSampler; use super::sampler::LunaSamplerError; use super::sampler::LunaSamplingRequest; +use super::sampler::SampledScore; use super::score::PendingScore; use super::transcript::CollectedTranscript; @@ -58,7 +59,7 @@ pub(super) struct ConversationRequest { pub(super) reset_token_limit: usize, pub(super) authorization: ScoreAuthorization, pub(super) thread: Arc, - pub(super) ready: oneshot::Sender>, + pub(super) ready: oneshot::Sender>, pub(super) metrics: Option>, pub(super) pending_score: PendingScore, } @@ -108,7 +109,8 @@ impl ConversationBackend { } impl Reservation { - pub(super) async fn run(self, request: ConversationRequest) { + pub(super) async fn run(self, mut request: ConversationRequest) { + request.sampling.pending_score = Some(request.pending_score.clone()); if self.generation.strong_count() == 0 || !request.authorization.is_current(&request.thread).await || request.pending_score.is_expired() 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 916edfc209..5159e7964f 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 @@ -5,6 +5,7 @@ use crate::async_scorer::conversation::ConversationRequest; use crate::async_scorer::sampler::LunaSampler; use crate::async_scorer::sampler::LunaSamplerError; use crate::async_scorer::sampler::LunaSamplingRequest; +use crate::async_scorer::sampler::SampledScore; use crate::async_scorer::sampler::tests::sample_request; use crate::async_scorer::sampler::tests::sampler_config; use crate::async_scorer::score::PendingScore; @@ -23,7 +24,9 @@ async fn pending_request( label: &str, ) -> Result<( ConversationRequest, - tokio::sync::oneshot::Receiver>, + impl std::future::Future< + Output = Result, tokio::sync::oneshot::error::RecvError>, + >, )> { let thread = Arc::clone(&fixture.test.codex); let config = thread @@ -69,7 +72,7 @@ async fn pending_request( ..sample_request(label) }, }, - score, + async move { score.await.map(|result| result.map(SampledScore::accept)) }, )) } @@ -395,10 +398,10 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul }, ]; let (server, _) = start_streaming_sse_server(streams).await; - let sampler = Arc::new(LunaSampler::new(sampler_config(format!( - "{}/v1", - server.uri() - )))); + let metrics = Arc::new(RecordingMetrics::default()); + let mut config = sampler_config(format!("{}/v1", server.uri())); + config.metrics = Some(metrics.clone()); + let sampler = Arc::new(LunaSampler::new(config)); let backend = ConversationBackend::new(sampler); let (request, score) = pending_request(&fixture, "seed").await?; backend.reserve(|| 0).1.unwrap().run(request).await; @@ -455,6 +458,23 @@ async fn concurrent_requests_retain_only_the_newest_completed_history() -> Resul .unwrap_err(); assert!(error.is_cancelled()); } + assert_eq!( + metrics + .0 + .lock() + .unwrap() + .iter() + .filter_map(|metric| match metric { + RecordedMetric::Counter(name, _, tags) + if name == "codex.guardian_v2.sampling.retry_outcome" => + tags.iter() + .find(|(key, _)| key == "outcome") + .map(|(_, value)| value.as_str()), + _ => None, + }) + .collect::>(), + vec!["success"; 4], + ); drop(finish_older); let (request, score) = pending_request(&fixture, "probe").await?; backend.reserve(|| 4).1.unwrap().run(request).await; @@ -626,3 +646,49 @@ async fn expiry_preserves_the_completed_checkpoint_and_transcript_cursor() -> Re server.shutdown().await; Ok(()) } + +#[tokio::test] +async fn expiry_discards_a_queued_score_without_recording_success() -> Result<()> { + skip_if_no_network!(Ok(())); + let fixture = GuardianFailureFixture::new().await?; + let store = fixture.test.codex.thread_extension_data(); + let server = responses::start_mock_server().await; + responses::mount_sse_once( + &server, + responses::sse(vec![ + responses::ev_output_text_delta("low"), + ev_assistant_message("score", "low"), + ev_completed("score"), + ]), + ) + .await; + let metrics = Arc::new(crate::async_scorer::sampler::tests::RecordingMetrics::default()); + let mut config = sampler_config(server.uri()); + config.metrics = Some(metrics.clone()); + let backend = ConversationBackend::new(Arc::new(LunaSampler::new(config))); + let (mut request, score) = pending_request(&fixture, "queued").await?; + request.pending_score.max_lag = 0; + let (index, reservation) = + backend.reserve(|| observe_unscored_call(&request.pending_score.progress, store)); + request.pending_score.index = index; + let pending_score = request.pending_score.clone(); + // Finish the worker while the score stays queued and the consumer remains unpolled. + tokio::time::timeout(ASYNC_TEST_TIMEOUT, reservation.unwrap().run(request)).await?; + assert!( + !metrics + .0 + .lock() + .unwrap() + .iter() + .any(|metric| metric.0 == "codex.guardian_v2.sampling.retry_outcome") + ); + observe_unscored_call(&pending_score.progress, store); + tokio::select! { + biased; + () = pending_score.expired() => {}, + _ = score => panic!("expired score was accepted"), + } + metrics.assert_outcome("superseded", /*attempts*/ 1, "superseded"); + fixture.test.codex.shutdown_and_wait().await?; + Ok(()) +} diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs index 38319bef7a..05e6aedd85 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/extension_tests.rs @@ -432,7 +432,8 @@ struct RecordingMetrics(Mutex>); impl RecordingMetrics { fn classification_samples(&self) -> Vec { self.0.lock().unwrap().iter().filter(|sample| { - !matches!(sample, RecordedMetric::Histogram(name, _, _) if name == codex_guardian_context::SECTION_COST_METRIC || name == codex_guardian_context::REQUEST_TOKENS_METRIC || name.starts_with("codex.guardian_v2.connection.")) + !matches!(sample, RecordedMetric::Histogram(name, _, _) if name == codex_guardian_context::SECTION_COST_METRIC || name == codex_guardian_context::REQUEST_TOKENS_METRIC || name.starts_with("codex.guardian_v2.connection.") || name.starts_with("codex.guardian_v2.sampling.")) + && !matches!(sample, RecordedMetric::Counter(name, _, _) if name.starts_with("codex.guardian_v2.sampling.")) }).cloned().collect() } } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/pending_score_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/pending_score_tests.rs index 4fac093e6a..fba473e9a5 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/pending_score_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/pending_score_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::async_scorer::conversation::ConversationBackend; +use crate::async_scorer::sampler::tests::RecordingMetrics; use crate::async_scorer::sampler::tests::sampler_config; use pretty_assertions::assert_eq; @@ -37,14 +38,17 @@ async fn stale_classifications_close_the_transport_without_invalidating_scores() }]) .await; let store = fixture.test.codex.thread_extension_data(); - store.insert(LunaSampler::new(sampler_config( + let metrics = Arc::new(RecordingMetrics::default()); + let mut config = sampler_config( proxy_websocket_servers_with_http( &[&server], ProxyPrewarmLimit::AllConnections, /*http_url*/ None, ) .await?, - ))); + ); + config.metrics = Some(metrics.clone()); + store.insert(LunaSampler::new(config)); let sampler = store.get::().unwrap(); sampler.prewarm().await; if mode == "conversation" { @@ -95,6 +99,15 @@ async fn stale_classifications_close_the_transport_without_invalidating_scores() assert_eq!(progress.inspect(/*call_id*/ None), expected); assert_eq!(cached_score(store), score); + let (outcome, failure_reason) = if early_score { + ("success", "none") + } else { + ("superseded", "superseded") + }; + // Socket closure can precede the sampling task recording its outcome. + tokio::time::timeout(ASYNC_TEST_TIMEOUT, metrics.1.notified()).await?; + metrics.assert_outcome(outcome, /*attempts*/ 1, failure_reason); + fixture.test.codex.shutdown_and_wait().await?; server.shutdown().await; } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs index a4e2baab40..8b03b1218b 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/retained_sampling_tests.rs @@ -50,7 +50,7 @@ async fn retained_websocket_requests_preserve_model_output_and_ids() -> Result<( section_costs: Vec::new(), }; let (completed, score) = tokio::join!(sampler.sample_retained(prepared, ready), score); - assert_eq!(score??, "low"); + assert_eq!(score??.accept(), "low"); let completed = completed .expect("completed model history") .into_iter() @@ -80,7 +80,7 @@ async fn retained_websocket_requests_preserve_model_output_and_ids() -> Result<( #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<()> { skip_if_no_network!(Ok(())); - for (limit, output, completed) in [ + for (limit, output, completed, receive_score, early_delta) in [ ( 512, vec![ev_assistant_message( @@ -110,8 +110,16 @@ async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<( true, ), (512, vec![], false), - ] { - let mut events = vec![ev_output_text_delta("low")]; + ].into_iter().flat_map(|(limit, output, completed)| { + [(true, true), (false, true), (false, false)].map(|(receive_score, early_delta)| { + (limit, output.clone(), completed, receive_score, early_delta) + }) + }) { + let mut events = if early_delta { + vec![ev_output_text_delta("low")] + } else { + Vec::new() + }; events.extend(output); if completed { events.push(ev_completed("score")); @@ -120,6 +128,8 @@ async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<( responses::mount_sse_once(&server, responses::sse(events)).await; let mut config = sampler_config(server.uri()); config.max_input_tokens = limit; + let metrics = Arc::new(RecordingMetrics::default()); + config.metrics = Some(metrics.clone()); let sampler = LunaSampler::new(config); let (ready, score) = tokio::sync::oneshot::channel(); let request = sample_request("parent"); @@ -145,9 +155,26 @@ async fn retention_overflow_or_incomplete_stream_keeps_early_score() -> Result<( truncations: Vec::new(), section_costs: Vec::new(), }; - let (history, score) = tokio::join!(sampler.sample_retained(prepared, ready), score); - assert_eq!(score??, "low"); + let history = if receive_score { + let (history, score) = tokio::join!(sampler.sample_retained(prepared, ready), score); + assert_eq!(score??.accept(), "low"); + history + } else { + let score = early_delta.then_some(score); + let history = sampler.sample_retained(prepared, ready).await; + // A queued score discarded without expiry remains a genuine cancellation. + drop(score); + history + }; + let (outcome, failure_reason) = if receive_score { + ("success", "none") + } else { + ("cancelled", "cancelled") + }; assert_eq!(history, None); + // Before the first score, an incomplete stream retries against the exhausted mock. + let attempts = if completed || early_delta { 1 } else { 2 }; + metrics.assert_outcome(outcome, attempts, failure_reason); } Ok(()) } 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 3b0d0541b0..59586b053f 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler.rs @@ -4,6 +4,9 @@ mod connection_pool; mod execution; mod response; +mod telemetry; + +pub(super) use telemetry::SampledScore; use super::request::PreparedRequest; use super::score::PendingScore; @@ -94,7 +97,7 @@ 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, - /// Snapshot expiry remains active until its response finishes draining. + /// Sampling expiry remains active until its response finishes draining. pub(super) pending_score: Option, } @@ -337,6 +340,8 @@ impl LunaSampler { }); } // The same cancellation future moves into the detached drain after an early score. + let mut telemetry = + telemetry::SamplingTelemetry::new(self.config.metrics.clone(), pending_score.clone()); let superseded = Box::pin(async move { tokio::select! { _ = superseded => {}, @@ -348,9 +353,15 @@ impl LunaSampler { } => {}, } }); - self.execution(&request, input, auth_owner_generation) - .run(execution::SamplingMode::Snapshot { superseded, scored }) - .await + let result = self + .execution(&request, input, auth_owner_generation) + .run( + execution::SamplingMode::Snapshot { superseded, scored }, + &mut telemetry, + ) + .await; + telemetry.record_result(&result); + result } pub(super) fn max_input_tokens(&self) -> usize { @@ -361,7 +372,7 @@ impl LunaSampler { pub(super) async fn sample_retained( &self, prepared: PreparedRequest<'_>, - ready: oneshot::Sender>, + ready: oneshot::Sender>, ) -> Option> { let auth_owner_generation = self.auth_owner_generation(); let mut completion = execution::RetainedCompletion { @@ -382,10 +393,20 @@ impl LunaSampler { .unzip(); let mut execution = self.execution(prepared.sampling, input, auth_owner_generation); execution.request.include = vec!["reasoning.encrypted_content".to_owned()]; + let mut telemetry = telemetry::SamplingTelemetry::new( + self.config.metrics.clone(), + prepared.sampling.pending_score.clone(), + ); let result = execution - .run(execution::SamplingMode::Retained(&mut completion)) + .run( + execution::SamplingMode::Retained(&mut completion), + &mut telemetry, + ) .await; - execution.request.input.extend(completion.finish(result)?); + execution + .request + .input + .extend(completion.finish(result, &mut telemetry)?); Some( execution .request diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs index 6c4930acbe..2228dd8045 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/connection_pool_cancellation_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::async_scorer::sampler::CLASSIFICATION_TOKEN_USAGE_METRIC; +use crate::async_scorer::sampler::LunaSamplerError; use crate::async_scorer::sampler::response::PendingResponse; use crate::async_scorer::sampler::tests::RecordingMetrics; use crate::async_scorer::sampler::tests::connect_sampler; @@ -33,7 +34,7 @@ async fn cancelled_classifiers_interrupt_and_reuse_only_completed_streams() -> R // None means reuse; explicit failures must close before the one-second drain deadline. let close_now = Some(Duration::from_millis(500)); let close_after_timeout = Some(Duration::from_secs(2)); - for (interrupted, close_timeout, expected_tokens) in [ + for (interrupted, close_timeout, expected_tokens, outcome) in [ (vec![accepted.clone(), completed.clone()], None, Some(42)), (vec![completed.clone(), accepted], None, Some(42)), (vec![rejected.clone()], close_now, None), @@ -50,7 +51,11 @@ async fn cancelled_classifiers_interrupt_and_reuse_only_completed_streams() -> R (vec![completed], close_after_timeout, Some(42)), (vec![error], close_after_timeout, None), (Vec::new(), close_after_timeout, None), - ] { + ] + .into_iter() + .flat_map(|(events, timeout, tokens)| { + ["cancelled", "superseded"].map(|outcome| (events.clone(), timeout, tokens, outcome)) + }) { let server = responses::start_websocket_server(vec![vec![ vec![responses::ev_response_created("stale")], interrupted, @@ -63,15 +68,34 @@ async fn cancelled_classifiers_interrupt_and_reuse_only_completed_streams() -> R config.metrics = Some(metrics.clone()); let sampler = Arc::new(connect_sampler(config).await?); let pending_sampler = Arc::clone(&sampler); - let pending = - tokio::spawn(async move { pending_sampler.sample(sample_request("stale")).await }); + let mut request = sample_request("stale"); + request.pending_score = Some(crate::async_scorer::score::PendingScore { + progress: Arc::default(), + index: 0, + max_lag: 0, + }); + let pending = tokio::spawn(async move { pending_sampler.sample(request).await }); tokio::time::timeout( Duration::from_secs(2), server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 0), ) .await?; - pending.abort(); - assert!(pending.await.unwrap_err().is_cancelled()); + if outcome == "cancelled" { + pending.abort(); + assert!(pending.await.unwrap_err().is_cancelled()); + } else { + sampler + .active_requests + .lock() + .unwrap() + .pop_front() + .unwrap() + .supersede + .send(()) + .unwrap(); + assert!(matches!(pending.await?, Err(LunaSamplerError::Superseded))); + } + metrics.assert_outcome(outcome, /*attempts*/ 1, outcome); let interrupt = tokio::time::timeout( Duration::from_secs(2), server.wait_for_request(/*connection_index*/ 0, /*request_index*/ 1), diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs index c5c8aee51f..0fc65adb41 100644 --- a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/execution.rs @@ -1,6 +1,7 @@ //! Executes one prepared classifier request with the existing retry and streaming rules. //! Snapshot output returns immediately with a detached drain that retains cancellation. //! Retained requests publish the same early result but keep ownership through completion. +//! Sampling outcomes describe score delivery, independently of retained response draining. //! 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. @@ -9,8 +10,11 @@ use super::ConnectionPool; use super::LunaSamplerConfig; use super::LunaSamplerError; use super::MAX_OUTPUT_BYTES; +use super::SampledScore; use super::connection_pool::RequestMode; use super::response::PendingResponse; +use super::telemetry::SamplingTelemetry; +use super::telemetry::record_retry; use codex_api::ApiError; use codex_api::ResponseEvent; use codex_api::ResponsesApiRequest; @@ -50,7 +54,7 @@ pub(super) struct SamplingExecution { } pub(super) struct RetainedCompletion { - pub(super) ready: Option>>, + pub(super) ready: Option>>, pub(super) output: Option>, pub(super) remaining_tokens: usize, pub(super) early_score: Option, @@ -112,13 +116,14 @@ impl RetainedCompletion { pub(super) fn finish( mut self, result: Result, + telemetry: &mut SamplingTelemetry, ) -> Option> { let reusable = result.as_ref().is_ok_and(|score| { matches!(score.as_str(), "high" | "low") && self.early_score.as_ref().is_none_or(|early| early == score) }); if let Some(ready) = self.ready.take() { - let _ = ready.send(result); + telemetry.deliver(ready, result); } self.output.filter(|output| reusable && !output.is_empty()) } @@ -157,6 +162,7 @@ impl SamplingExecution { return false; } self.connections.clear(); + record_retry(self.config.metrics.as_deref(), "http_401", "auth_recovery"); return true; } else { status.is_server_error() || *status == StatusCode::TOO_MANY_REQUESTS @@ -189,6 +195,11 @@ impl SamplingExecution { }; if retryable && *retries < MAX_SAMPLING_RETRIES { *retries += 1; + record_retry( + self.config.metrics.as_deref(), + super::super::metrics::sampler_failure_reason(error), + "transport", + ); return true; } false @@ -197,6 +208,7 @@ impl SamplingExecution { pub(super) async fn run( &mut self, mut mode: SamplingMode<'_>, + telemetry: &mut SamplingTelemetry, ) -> Result { let auth_changes = self .config @@ -244,6 +256,7 @@ impl SamplingExecution { }; 'retry: loop { ensure_account_owner()?; + telemetry.attempts = telemetry.attempts.saturating_add(1); if let SamplingMode::Retained(completion) = &mut mode { completion.output = Some(Vec::new()); completion.remaining_tokens = retention_budget; @@ -390,7 +403,7 @@ impl SamplingExecution { SamplingMode::Retained(ref mut retained) => { retained.early_score = Some(delta.clone()); if let Some(ready) = retained.ready.take() { - let _ = ready.send(Ok(delta)); + telemetry.deliver(ready, Ok(delta)); } continue; } diff --git a/codex-rs/ext/guardian-v2/src/async_scorer/sampler/telemetry.rs b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/telemetry.rs new file mode 100644 index 0000000000..8d909da837 --- /dev/null +++ b/codex-rs/ext/guardian-v2/src/async_scorer/sampler/telemetry.rs @@ -0,0 +1,122 @@ +//! Records sampling retries and attempts through score delivery or cancellation. +//! Retained scores carry telemetry until accepted or discarded by their consumer. + +use super::LunaSamplerError; +use super::PendingScore; +use codex_extension_api::ExtensionMetrics; +use std::sync::Arc; + +pub(in crate::async_scorer) struct SampledScore { + output: String, + telemetry: SamplingTelemetry, +} + +impl SampledScore { + pub(in crate::async_scorer) fn accept(mut self) -> String { + self.telemetry.outcome = Some(("success", "none")); + self.output + } +} + +pub(super) struct SamplingTelemetry { + metrics: Option>, + pending_score: Option, + pub(super) attempts: usize, + outcome: Option<(&'static str, &'static str)>, +} +impl Drop for SamplingTelemetry { + fn drop(&mut self) { + if let Some(metrics) = self.metrics.as_deref() { + let (outcome, failure_reason) = + self.outcome.unwrap_or_else(|| self.undelivered_outcome()); + let attempt_count = self.attempts.to_string(); + let tags = [ + ("outcome", outcome), + ("attempt_count", attempt_count.as_str()), + ("retried", if self.attempts > 1 { "true" } else { "false" }), + ("failure_reason", failure_reason), + ]; + // One initial attempt, two ordinary retries and at most two auth recoveries. + metrics.histogram_with_boundaries( + "codex.guardian_v2.sampling.attempt_count", + i64::try_from(self.attempts).unwrap_or(i64::MAX), + &[0.0, 1.0, 2.0, 3.0, 4.0, 5.0], + &tags, + ); + metrics.counter( + "codex.guardian_v2.sampling.retry_outcome", + /*inc*/ 1, + &tags, + ); + } + } +} + +impl SamplingTelemetry { + pub(super) fn new( + metrics: Option>, + pending_score: Option, + ) -> Self { + Self { + metrics, + pending_score, + attempts: 0, + outcome: None, + } + } + + fn undelivered_outcome(&self) -> (&'static str, &'static str) { + // An expiry-first select can drop sampling before it returns Superseded. + if self + .pending_score + .as_ref() + .is_some_and(PendingScore::is_expired) + { + ("superseded", "superseded") + } else { + ("cancelled", "cancelled") + } + } + + pub(super) fn record_result(&mut self, result: &Result) { + self.outcome.get_or_insert_with(|| match result { + Ok(_) => ("success", "none"), + Err(LunaSamplerError::Superseded) => ("superseded", "superseded"), + Err(error) => ( + "failure", + super::super::metrics::sampler_failure_reason(error), + ), + }); + } + + pub(super) fn deliver( + &mut self, + ready: tokio::sync::oneshot::Sender>, + result: Result, + ) { + if result.is_err() { + self.record_result(&result); + } + let result = result.map(|output| SampledScore { + output, + telemetry: std::mem::replace(self, Self::new(/*metrics*/ None, /*pending_score*/ None)), + }); + if ready.send(result).is_err() { + self.outcome = Some(self.undelivered_outcome()); + } + } +} + +pub(super) fn record_retry( + metrics: Option<&dyn ExtensionMetrics>, + failure_reason: &'static str, + kind: &'static str, +) { + if let Some(metrics) = metrics { + metrics.counter( + "codex.guardian_v2.sampling.retry", + /*inc*/ 1, + &[("failure_reason", failure_reason), ("kind", kind)], + ); + } +} 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 cce5cb311f..80a7225723 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 @@ -345,7 +345,53 @@ async fn classifier_preserves_api_key_cyber_program_policy() -> Result<()> { type RecordedMetric = (String, i64, Vec<(String, String)>); #[derive(Default)] -pub(super) struct RecordingMetrics(pub(super) Mutex>); +pub(in crate::async_scorer) struct RecordingMetrics( + pub(in crate::async_scorer) Mutex>, + pub(in crate::async_scorer) tokio::sync::Notify, +); + +impl RecordingMetrics { + pub(in crate::async_scorer) fn assert_outcome( + &self, + outcome: &str, + attempts: i64, + failure_reason: &str, + ) { + let tags = json!([ + ["outcome", outcome], + ["attempt_count", attempts.to_string()], + ["retried", if attempts > 1 { "true" } else { "false" }], + ["failure_reason", failure_reason] + ]); + let samples = self.0.lock().unwrap(); + for (name, value) in [ + ("codex.guardian_v2.sampling.attempt_count", attempts), + ("codex.guardian_v2.sampling.retry_outcome", 1), + ] { + assert_eq!( + json!( + samples + .iter() + .filter(|metric| metric.0 == name) + .map(|metric| (&metric.1, &metric.2)) + .collect::>() + ), + json!([[value, tags]]), + ); + } + } + + fn retries(&self) -> serde_json::Value { + let samples = self.0.lock().unwrap(); + let mut retries = samples + .iter() + .filter(|metric| metric.0 == "codex.guardian_v2.sampling.retry") + .map(|metric| (&metric.1, &metric.2)) + .collect::>(); + retries.sort(); + json!(retries) + } +} impl ExtensionMetrics for RecordingMetrics { fn histogram_with_boundaries( @@ -358,7 +404,12 @@ impl ExtensionMetrics for RecordingMetrics { self.histogram(name, value, tags); } - fn counter(&self, _name: &str, _inc: i64, _tags: &[(&str, &str)]) {} + fn counter(&self, name: &str, inc: i64, tags: &[(&str, &str)]) { + self.histogram(name, inc, tags); + if name == "codex.guardian_v2.sampling.retry_outcome" { + self.1.notify_one(); + } + } fn histogram(&self, name: &str, value: i64, tags: &[(&str, &str)]) { if name == "codex.guardian_v2.connection.duration_ms" { @@ -1290,17 +1341,33 @@ async fn sampler_reconnects_after_transient_service_failures() -> Result<()> { ]), ) .await; - let sampler = connect_sampler(sampler_config( + let metrics = Arc::new(RecordingMetrics::default()); + let mut config = sampler_config( proxy_websocket_servers_with_http( &[&first, &second], ProxyPrewarmLimit::AllConnections, Some(&http.uri()), ) .await?, - )) - .await?; + ); + config.metrics = Some(metrics.clone()); + let sampler = connect_sampler(config).await?; assert_eq!(sampler.sample(sample_request("turn-1")).await?, "low"); + metrics.assert_outcome("success", /*attempts*/ 3, "none"); + assert_eq!( + metrics.retries(), + json!([ + [ + 1, + [ + ["failure_reason", "flex_unavailable"], + ["kind", "transport"] + ] + ], + [1, [["failure_reason", "http_5xx"], ["kind", "transport"]]], + ]) + ); assert_eq!(first.single_connection().len(), 1); assert_eq!(second.single_connection().len(), 1); let request = recovered.single_request(); @@ -1346,15 +1413,17 @@ async fn sampler_limits_transient_recovery_attempts() -> Result<()> { })]), ) .await; - let sampler = connect_sampler(sampler_config( + let metrics = Arc::new(RecordingMetrics::default()); + let mut config = sampler_config( proxy_websocket_servers_with_http( &[&first, &second], ProxyPrewarmLimit::AllConnections, Some(&http.uri()), ) .await?, - )) - .await?; + ); + config.metrics = Some(metrics.clone()); + let sampler = connect_sampler(config).await?; let error = sampler .sample(sample_request("turn-1")) @@ -1365,6 +1434,7 @@ async fn sampler_limits_transient_recovery_attempts() -> Result<()> { error, LunaSamplerError::Api(ApiError::FlexUnavailable) )); + metrics.assert_outcome("failure", /*attempts*/ 3, "flex_unavailable"); assert_eq!(first.single_connection().len(), 1); assert_eq!(second.single_connection().len(), 1); assert_eq!(third.requests().len(), 1); @@ -1398,6 +1468,8 @@ async fn parent_response_id_survives_classifier_transport_retry() -> Result<()> )), ); } + let metrics = Arc::new(RecordingMetrics::default()); + config.metrics = Some(metrics.clone()); let sampler = connect_sampler(config).await?; let expired_id = expired.single_handshake().header("thread-id").unwrap(); sampler @@ -1410,6 +1482,7 @@ async fn parent_response_id_survives_classifier_transport_retry() -> Result<()> let mut request = sample_request("turn-1"); request.parent_response_id = Some(parent_response_id.to_owned()); assert_eq!(sampler.sample(request).await?, "low"); + metrics.assert_outcome("success", /*attempts*/ 2, "none"); for server in [&expired, &healthy] { let requests = server.single_connection(); assert_eq!(requests.len(), 1); @@ -1427,3 +1500,47 @@ async fn parent_response_id_survives_classifier_transport_retry() -> Result<()> Ok(()) } + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sampler_records_401_recovery_separately_from_transport_retries() -> Result<()> { + skip_if_no_network!(Ok(())); + let rejected = vec![vec![json!({"type": "error", "status": 401, + "error": {"type": "authentication_error", "message": "expired"}})]]; + let ws = responses::start_websocket_server(vec![rejected.clone(), rejected]).await; + let http = responses::start_mock_server().await; + let recovered = responses::mount_sse_once( + &http, + responses::sse(vec![ + ev_assistant_message("recovered", "low"), + ev_completed("recovered"), + ]), + ) + .await; + let metrics = Arc::new(RecordingMetrics::default()); + let base_url = proxy_websocket_servers_with_http( + &[&ws, &ws], + ProxyPrewarmLimit::AllConnections, + Some(&http.uri()), + ) + .await?; + let mut config = sampler_config(base_url); + config.metrics = Some(metrics.clone()); + let auth = config.provider.auth_manager().unwrap(); + auth.set_external_auth(Arc::new(RefreshableAuth(Mutex::new("refreshed")))) + .await?; + let sampler = connect_sampler(config).await?; + assert_eq!(sampler.sample(sample_request("turn-1")).await?, "low"); + assert_eq!( + recovered.single_request().header("authorization"), + Some("Bearer refreshed".to_owned()) + ); + metrics.assert_outcome("success", /*attempts*/ 2, "none"); + assert_eq!( + metrics.retries(), + json!([[ + 1, + [["failure_reason", "http_401"], ["kind", "auth_recovery"]] + ],]) + ); + Ok(()) +}