From e16d098c00e47aa45719e8afc8f6fe301424bcca Mon Sep 17 00:00:00 2001 From: jif Date: Tue, 25 Aug 2026 14:18:53 +0000 Subject: [PATCH] Instrument exec-server shell snapshot capture (#40620) ## What changed - Emit `codex.shell_snapshot` and `codex.shell_snapshot.duration_ms` for each exec-server capture attempt, including success and failure tags. - Label exec-server metrics as `v2` and existing core snapshot metrics as `v1` so the two implementations can be distinguished. - Use the exec server's metrics client when available and fall back to the configured global client for local execution. ## Testing - Extend the bounded-retry and single-flight snapshot test to verify metric counts and tags for first-attempt success, retries, recovery, and exhaustion. GitOrigin-RevId: 94f73e58a4a1a90c518fa102f138ca7ee29b5627 --- codex-rs/core/src/shell_snapshot.rs | 4 +- codex-rs/exec-server/src/local_process.rs | 2 +- codex-rs/exec-server/src/shell_snapshot.rs | 18 ++-- .../exec-server/src/shell_snapshot_tests.rs | 84 ++++++++++++++++++- codex-rs/exec-server/src/telemetry.rs | 25 ++++++ 5 files changed, 121 insertions(+), 12 deletions(-) diff --git a/codex-rs/core/src/shell_snapshot.rs b/codex-rs/core/src/shell_snapshot.rs index b0782e5c5a..2e0d734a98 100644 --- a/codex-rs/core/src/shell_snapshot.rs +++ b/codex-rs/core/src/shell_snapshot.rs @@ -94,7 +94,7 @@ impl ShellSnapshot { async { let timer = config .session_telemetry - .start_timer("codex.shell_snapshot.duration_ms", &[]); + .start_timer("codex.shell_snapshot.duration_ms", &[("version", "v1")]); let snapshot = ShellSnapshot::try_create( &config.codex_home, config.session_id, @@ -105,7 +105,7 @@ impl ShellSnapshot { .await; let success_tag = if snapshot.is_ok() { "true" } else { "false" }; let _ = timer.map(|timer| timer.record(&[("success", success_tag)])); - let mut counter_tags = vec![("success", success_tag)]; + let mut counter_tags = vec![("version", "v1"), ("success", success_tag)]; if let Some(failure_reason) = snapshot.as_ref().err() { counter_tags.push(("failure_reason", *failure_reason)); } diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 927276f36e..27746255ce 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -355,7 +355,7 @@ impl LocalProcess { #[cfg(unix)] self.inner .shell_snapshots - .prepare(¶ms, &mut prepared) + .prepare(¶ms, &mut prepared, &self.inner.telemetry) .await?; if prepared.command.is_empty() { return Err(invalid_params("argv must not be empty".to_string())); diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index e352d63641..324580a66f 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -26,6 +26,7 @@ use crate::protocol::ExecParams; use crate::protocol::ShellSnapshotRequest; use crate::rpc::internal_error; use crate::rpc::invalid_params; +use crate::telemetry::ExecServerTelemetry; const MAX_CACHED_SNAPSHOTS: usize = 16; const MAX_SNAPSHOT_BYTES: usize = 512 * 1024; @@ -60,6 +61,7 @@ impl ShellSnapshotCache { &self, params: &ExecParams, prepared: &mut PreparedExecRequest, + telemetry: &ExecServerTelemetry, ) -> Result<(), JSONRPCErrorError> { let Some(request) = params.shell_snapshot.as_ref() else { return Ok(()); @@ -134,12 +136,16 @@ impl ShellSnapshotCache { }; let Ok(snapshot) = snapshot .get_or_init(|| async { - capture_snapshot(params, prepared, shell_type) - .await - .map_err(|err| { - tracing::warn!("failed to capture shell snapshot: {err:?}"); - Instant::now() + SNAPSHOT_RETRY_BACKOFF - }) + let started_at = std::time::Instant::now(); + let result = capture_snapshot(params, prepared, shell_type).await; + telemetry.shell_snapshot_captured( + started_at.elapsed(), + result.as_ref().map(|_| ()).map_err(|_| "capture_failed"), + ); + result.map_err(|err| { + tracing::warn!("failed to capture shell snapshot: {err:?}"); + Instant::now() + SNAPSHOT_RETRY_BACKOFF + }) }) .await else { diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index 1f24ef5de8..f8a380a3f1 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -1,6 +1,12 @@ +use std::collections::BTreeMap; use std::collections::HashMap; +use codex_otel::MetricsClient; +use codex_otel::MetricsConfig; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; +use opentelemetry_sdk::metrics::InMemoryMetricExporter; +use opentelemetry_sdk::metrics::data::AggregatedMetrics; +use opentelemetry_sdk::metrics::data::MetricData; use pretty_assertions::assert_eq; use test_case::test_case; @@ -14,7 +20,9 @@ use crate::protocol::ExecParams; use crate::protocol::ProcessId; use crate::protocol::ShellInfo; use crate::protocol::ShellSnapshotRequest; +use crate::telemetry::ExecServerTelemetry; +#[test_case(1; "succeeds_on_first_attempt")] #[test_case(2; "recovers_on_second_attempt")] #[test_case(3; "recovers_on_last_attempt")] #[test_case(4; "stops_after_three_failures")] @@ -57,6 +65,16 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( network_proxy: None, }; let cache = ShellSnapshotCache::default(); + let metrics = MetricsClient::new( + MetricsConfig::in_memory( + "test", + "codex-exec-server", + env!("CARGO_PKG_VERSION"), + InMemoryMetricExporter::default(), + ) + .with_runtime_reader(), + )?; + let telemetry = ExecServerTelemetry::new(metrics.clone()); for attempt in 1..=5 { if attempt == recovery_attempt { @@ -84,8 +102,8 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( .await .expect("prepare concurrent capture"); let (first, second) = tokio::join!( - cache.prepare(¶ms, &mut prepared), - cache.prepare(¶ms, &mut concurrent), + cache.prepare(¶ms, &mut prepared, &telemetry), + cache.prepare(¶ms, &mut concurrent, &telemetry), ); first.expect("capture failure must preserve command fallback"); second.expect("concurrent request must share the capture attempt"); @@ -97,7 +115,7 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( tokio::time::pause(); if attempt < recovery_attempt || recovery_attempt > MAX_SNAPSHOT_ATTEMPTS { cache - .prepare(¶ms, &mut prepared) + .prepare(¶ms, &mut prepared, &telemetry) .await .expect("capture must stay cached during backoff"); assert_eq!( @@ -114,6 +132,66 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( tokio::time::advance(SNAPSHOT_RETRY_BACKOFF).await; tokio::time::resume(); } + + let snapshot = metrics.snapshot()?; + let mut counters = BTreeMap::new(); + let mut durations = BTreeMap::new(); + for metric in snapshot + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + { + match metric.name() { + "codex.shell_snapshot" => { + let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else { + panic!("expected shell snapshot counter"); + }; + for point in sum.data_points() { + let tags = point + .attributes() + .map(|attribute| (attribute.key.to_string(), attribute.value.to_string())) + .collect::>(); + counters.insert(tags, point.value()); + } + } + "codex.shell_snapshot.duration_ms" => { + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("expected shell snapshot duration histogram"); + }; + for point in histogram.data_points() { + let tags = point + .attributes() + .map(|attribute| (attribute.key.to_string(), attribute.value.to_string())) + .collect::>(); + durations.insert(tags, point.count()); + } + } + _ => {} + } + } + let mut expected_counters = BTreeMap::new(); + let mut expected_durations = BTreeMap::new(); + let failures = (recovery_attempt - 1).min(MAX_SNAPSHOT_ATTEMPTS) as u64; + for (success, count) in [ + ("false", failures), + ("true", u64::from(recovery_attempt <= MAX_SNAPSHOT_ATTEMPTS)), + ] { + if count == 0 { + continue; + } + let mut tags = BTreeMap::from([ + ("version".to_string(), "v2".to_string()), + ("success".to_string(), success.to_string()), + ]); + expected_durations.insert(tags.clone(), count); + if success == "false" { + tags.insert("failure_reason".to_string(), "capture_failed".to_string()); + } + expected_counters.insert(tags, count); + } + assert_eq!( + (counters, durations), + (expected_counters, expected_durations) + ); Ok(()) } diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs index e78d2ec455..13dbbb8a4f 100644 --- a/codex-rs/exec-server/src/telemetry.rs +++ b/codex-rs/exec-server/src/telemetry.rs @@ -159,6 +159,31 @@ impl ExecServerTelemetry { }); } + #[cfg(unix)] + pub(crate) fn shell_snapshot_captured( + &self, + duration: Duration, + result: Result<(), &'static str>, + ) { + // Local execution has no exec-server telemetry owner. Use the host's + // configured metrics client while preserving an explicit server client. + let Some(metrics) = self + .inner + .as_ref() + .map(|inner| inner.metrics.clone()) + .or_else(codex_otel::global) + else { + return; + }; + let success = if result.is_ok() { "true" } else { "false" }; + let mut tags = vec![("version", "v2"), ("success", success)]; + let _ = metrics.record_duration("codex.shell_snapshot.duration_ms", duration, &tags); + if let Err(failure_reason) = result { + tags.push(("failure_reason", failure_reason)); + } + let _ = metrics.counter("codex.shell_snapshot", /*inc*/ 1, &tags); + } + pub(crate) fn remote_registration_completed(&self, result: &'static str, duration: Duration) { self.record_operation(REMOTE_REGISTRATION_METRICS, result, duration); }