diff --git a/codex-rs/core/src/tools/runtimes/unified_exec.rs b/codex-rs/core/src/tools/runtimes/unified_exec.rs index 82a76ded4c..0a7b9baf75 100644 --- a/codex-rs/core/src/tools/runtimes/unified_exec.rs +++ b/codex-rs/core/src/tools/runtimes/unified_exec.rs @@ -64,6 +64,7 @@ use std::time::Duration; use tokio_util::sync::CancellationToken; mod launch; +mod snapshot_metrics; use launch::with_launch_failure_events; @@ -313,6 +314,11 @@ impl<'a> ToolRuntime for UnifiedExecRunt } else { environment_shell }; + let snapshot_metrics = snapshot_metrics::SnapshotMetrics::start( + req, + &ctx.step_context.turn.config, + &self.shell_mode, + ); let shell_snapshot = if environment_is_remote || credential_broker_available && launch_sandbox_permissions.requires_escalated_permissions() @@ -339,6 +345,9 @@ impl<'a> ToolRuntime for UnifiedExecRunt ) .await }; + let snapshot_wait = snapshot_metrics + .as_ref() + .map(|metrics| metrics.started_at.elapsed()); let shell_snapshot_location = shell_snapshot.as_ref().map(|snapshot| snapshot.path()); let mut env = exec_env_for_sandbox_permissions(&req.env, launch_sandbox_permissions); let snapshot_credential_context = if let Some(snapshot) = shell_snapshot.as_ref() @@ -527,8 +536,9 @@ impl<'a> ToolRuntime for UnifiedExecRunt { command = command_with_path_prepends; } + let mut snapshot_used = false; if !environment_is_remote { - command = maybe_wrap_shell_lc_with_snapshot( + let wrapped = maybe_wrap_shell_lc_with_snapshot( &command, shell, shell_snapshot_location.as_ref(), @@ -536,6 +546,8 @@ impl<'a> ToolRuntime for UnifiedExecRunt &env, &runtime_path_prepends, ); + snapshot_used = wrapped != command; + command = wrapped; } let brokered_shell_snapshot_missing = !environment_is_remote && managed_network.is_some() @@ -561,6 +573,12 @@ impl<'a> ToolRuntime for UnifiedExecRunt } network.restore_and_disable_brokered_credentials(&mut env, &mut command); } + if let Some(metrics) = snapshot_metrics + && let Some(wait) = snapshot_wait + { + let outcome = if snapshot_used { "used" } else { "fallback" }; + metrics.record(&ctx.step_context.session_telemetry, wait, outcome); + } if req.shell_snapshot.is_some() { let exports = runtime_path_prepends.shell_exports_after_snapshot(&explicit_env_overrides); diff --git a/codex-rs/core/src/tools/runtimes/unified_exec/snapshot_metrics.rs b/codex-rs/core/src/tools/runtimes/unified_exec/snapshot_metrics.rs new file mode 100644 index 0000000000..2f5ff3ecd1 --- /dev/null +++ b/codex-rs/core/src/tools/runtimes/unified_exec/snapshot_metrics.rs @@ -0,0 +1,77 @@ +//! Command-level v1 snapshot observations. Count replay selection, not successful +//! restoration or command completion; background captures are not commands. + +use std::time::Duration; +use std::time::Instant; + +use codex_features::Feature; +use codex_otel::SessionTelemetry; +use codex_tools::UnifiedExecShellMode; + +use super::UnifiedExecRequest; +use crate::config::Config; + +pub(super) struct SnapshotMetrics { + pub(super) started_at: Instant, + state: &'static str, +} + +impl SnapshotMetrics { + pub(super) fn start( + req: &UnifiedExecRequest, + config: &Config, + shell_mode: &UnifiedExecShellMode, + ) -> Option { + let brokered = config + .permissions + .network + .as_ref() + .is_some_and(|network| network.enabled() && network.credential_broker_enabled()); + // Executor snapshots emit at their own selection boundary. Disabled, + // remote legacy, and non-login invocations are outside this denominator. + if !cfg!(unix) + || req.shell_snapshot.is_some() + || req.turn_environment.environment.is_remote() + || !config.features.enabled(Feature::ShellSnapshot) + || !matches!(shell_mode, UnifiedExecShellMode::Direct) + || !req.shell.is_posix_login() + || (req.cwd != req.turn_environment.selection.cwd && !brokered) + { + return None; + } + let state = if brokered { + // Protected snapshots use a separate lazy cache, not the prewarm task. + "protected" + } else { + match req.turn_environment.shell_snapshot.peek() { + None => "prewarm_pending", + Some(Some(snapshot)) + if &snapshot.shell_environment_policy + == req.turn_environment.shell_environment_policy() => + { + "prewarm_ready" + } + Some(_) => "unavailable", + } + }; + Some(Self { + started_at: Instant::now(), + state, + }) + } + + pub(super) fn record( + self, + telemetry: &SessionTelemetry, + wait: Duration, + outcome: &'static str, + ) { + let tags = [ + ("version", "v1"), + ("state", self.state), + ("outcome", outcome), + ]; + telemetry.counter("codex.shell_snapshot.command", /*inc*/ 1, &tags); + telemetry.record_duration("codex.shell_snapshot.wait_ms", wait, &tags); + } +} diff --git a/codex-rs/core/tests/suite/approvals.rs b/codex-rs/core/tests/suite/approvals.rs index 5c0afeaaf5..22f28f9212 100644 --- a/codex-rs/core/tests/suite/approvals.rs +++ b/codex-rs/core/tests/suite/approvals.rs @@ -2754,6 +2754,14 @@ async fn shell_startup_credentials_are_brokered( allow_login_shell: bool, login: bool, ) -> Result<()> { + use codex_otel::MetricsClient; + use codex_otel::MetricsConfig; + use opentelemetry_sdk::metrics::data::AggregatedMetrics; + use opentelemetry_sdk::metrics::data::HistogramDataPoint; + use opentelemetry_sdk::metrics::data::MetricData; + use opentelemetry_sdk::metrics::data::SumDataPoint; + use std::collections::BTreeMap; + skip_if_no_network!(Ok(())); const GH_HOST: &str = "github.example.com"; @@ -2915,9 +2923,24 @@ ZDOTDIR = "{}" } }); let server = start_mock_server().await; - let test = builder.build(&server).await?; + let mut test = builder.build(&server).await?; + let metrics = if login { + // Keep the collector local to this thread so parallel tests cannot affect counts. + let metrics = MetricsClient::new( + MetricsConfig::in_memory("test", "test", "1", Default::default()).with_runtime_reader(), + )?; + let mut options = test.start_thread_options().await; + options.thread_extension_init.insert(metrics.clone()); + let thread = test.thread_manager.start_thread(options).await?; + test.codex.shutdown_and_wait().await?; + test.codex = thread.thread; + test.session_configured = thread.session_configured; + Some(metrics) + } else { + None + }; let snapshot_dir = test.home.path().join("shell_snapshots"); - let workdir = if allow_login_shell || shell_mode == "zsh_fork" { + let workdir = if (allow_login_shell && !login) || shell_mode == "zsh_fork" { test.cwd.path().to_path_buf() } else { let workdir = test.cwd.path().join("brokered-subdirectory"); @@ -3025,6 +3048,44 @@ ZDOTDIR = "{}" } let result = parse_result(&output); assert_eq!(result.exit_code, Some(0), "command failed: {result:?}"); + if let Some(metrics) = metrics { + let snapshot = metrics.snapshot()?; + let mut counts = BTreeMap::new(); + for metric in snapshot + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + { + let count = match (metric.name(), metric.data()) { + ("codex.shell_snapshot.command", AggregatedMetrics::U64(MetricData::Sum(sum))) => { + for point in sum.data_points() { + let tags = point + .attributes() + .map(|tag| (tag.key.as_str(), tag.value.to_string())) + .collect::>(); + assert_eq!( + ["version", "state", "outcome"] + .map(|key| tags.get(key).map(String::as_str)), + [Some("v1"), Some("protected"), Some("used")], + ); + } + sum.data_points().map(SumDataPoint::value).sum::() + } + ( + "codex.shell_snapshot.wait_ms", + AggregatedMetrics::F64(MetricData::Histogram(histogram)), + ) => histogram.data_points().map(HistogramDataPoint::count).sum(), + _ => continue, + }; + counts.insert(metric.name(), count); + } + assert_eq!( + counts, + BTreeMap::from([ + ("codex.shell_snapshot.command", 1), + ("codex.shell_snapshot.wait_ms", 1), + ]), + ); + } assert!(!result.stdout.contains(REAL_GITHUB_TOKEN)); assert!(!result.stdout.contains(REAL_CUSTOM_TOKEN)); let values = result.stdout.lines().collect::>(); diff --git a/codex-rs/exec-server/src/shell_snapshot.rs b/codex-rs/exec-server/src/shell_snapshot.rs index 5d399e263e..511ad9520f 100644 --- a/codex-rs/exec-server/src/shell_snapshot.rs +++ b/codex-rs/exec-server/src/shell_snapshot.rs @@ -1,5 +1,6 @@ //! Cache captured shell state in executor memory. Replay through a per-launch //! unnamed reader when the capture sandbox permits it; otherwise retain env replay. +//! Command metrics exclude prewarm and measure lookup/capture wait before replay. use std::collections::HashMap; use std::collections::VecDeque; @@ -63,6 +64,7 @@ struct CachedShellSnapshot { } struct ShellSnapshot { + prewarmed: bool, state: String, file_source: bool, environment: HashMap, @@ -121,6 +123,7 @@ impl ShellSnapshotCache { } }; + let wait_started_at = std::time::Instant::now(); let (snapshot, attempt) = { let mut entries = self.entries.lock().await; let position = entries.iter().position(|entry| { @@ -166,14 +169,27 @@ impl ShellSnapshotCache { (snapshot, 1) } }; + let mut availability = match snapshot.get() { + Some(Ok(snapshot)) if snapshot.prewarmed => "prewarm_ready", + Some(Ok(_)) => "cache_hit", + Some(Err(_)) => "unavailable", + None => "capture_pending", + }; let capture = async { + let prewarmed = purpose == CapturePurpose::Prewarm; let attempt = attempt.to_string(); let purpose = match purpose { CapturePurpose::Execution => "execution", CapturePurpose::Prewarm => "prewarm", }; let started_at = std::time::Instant::now(); - let result = capture_snapshot(params, prepared, shell_type).await; + let result = + capture_snapshot(params, prepared, shell_type) + .await + .map(|mut snapshot| { + snapshot.prewarmed = prewarmed; + snapshot + }); telemetry.shell_snapshot_captured( started_at.elapsed(), result.as_ref().map(|_| ()).map_err(|(reason, _)| *reason), @@ -190,6 +206,7 @@ impl ShellSnapshotCache { CapturePurpose::Execution => { snapshot .get_or_init(|| async { + availability = "on_demand"; capture.await.map_err(|err| { tracing::warn!("failed to capture shell snapshot: {err:?}"); Instant::now() + SNAPSHOT_RETRY_BACKOFF @@ -205,12 +222,17 @@ impl ShellSnapshotCache { .await? } }; - let Ok(snapshot) = snapshot else { - return Ok(None); - }; + let wait = wait_started_at.elapsed(); if purpose == CapturePurpose::Prewarm { return Ok(None); } + let Ok(snapshot) = snapshot else { + telemetry.shell_snapshot_command(wait, availability, "fallback"); + return Ok(None); + }; + if availability == "capture_pending" && snapshot.prewarmed { + availability = "prewarm_pending"; + } // POSIX sh cannot portably close arbitrary descriptors above 9. Keep its // existing replay until the launcher supports remapping child descriptors. @@ -227,6 +249,7 @@ impl ShellSnapshotCache { ?error, "cannot prepare shell snapshot transport; using normal startup" ); + telemetry.shell_snapshot_command(wait, availability, "fallback"); return Ok(None); } } @@ -293,6 +316,7 @@ impl ShellSnapshotCache { params.argv[2] ); + telemetry.shell_snapshot_command(wait, availability, "used"); Ok(reader) } } @@ -500,6 +524,7 @@ fn parse_snapshot( environment.retain(|name, _| !shell_environment::is_non_inheritable_env_var(name)); Ok(ShellSnapshot { + prewarmed: false, state, file_source: false, environment, diff --git a/codex-rs/exec-server/src/shell_snapshot_tests.rs b/codex-rs/exec-server/src/shell_snapshot_tests.rs index c79eb72589..b3852d548a 100644 --- a/codex-rs/exec-server/src/shell_snapshot_tests.rs +++ b/codex-rs/exec-server/src/shell_snapshot_tests.rs @@ -85,6 +85,7 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( )?; let telemetry = ExecServerTelemetry::new(metrics.clone()); + let mut expected_commands = BTreeMap::new(); for attempt in 1..=5 { if attempt == recovery_attempt { std::fs::write( @@ -142,6 +143,13 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( (&concurrent.command, &concurrent.env) ); + let outcome = if concurrent.command != params.argv { + "used" + } else { + "fallback" + }; + *expected_commands.entry(outcome.to_string()).or_insert(0) += + if prewarming { 1 } else { 2 }; tokio::time::pause(); if attempt < recovery_attempt || recovery_attempt > MAX_SNAPSHOT_ATTEMPTS { cache @@ -153,6 +161,7 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( ) .await .expect("capture must stay cached during backoff"); + *expected_commands.entry("fallback".to_string()).or_insert(0) += 1; assert_eq!( (&prepared.command, &prepared.env), (¶ms.argv, ¶ms.env) @@ -174,6 +183,8 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( let snapshot = metrics.snapshot()?; let mut counters = BTreeMap::new(); let mut durations = BTreeMap::new(); + let mut commands = BTreeMap::new(); + let mut waits = BTreeMap::new(); for metric in snapshot .scope_metrics() .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) @@ -203,9 +214,41 @@ async fn snapshot_failure_retries_are_bounded_and_single_flight( durations.insert(tags, point.count()); } } + "codex.shell_snapshot.command" => { + let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else { + panic!("expected command counter"); + }; + for point in sum.data_points() { + let outcome = point + .attributes() + .find(|tag| tag.key.as_str() == "outcome") + .unwrap() + .value + .to_string(); + *commands.entry(outcome).or_insert(0) += point.value(); + } + } + "codex.shell_snapshot.wait_ms" => { + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("expected wait histogram"); + }; + for point in histogram.data_points() { + let outcome = point + .attributes() + .find(|tag| tag.key.as_str() == "outcome") + .unwrap() + .value + .to_string(); + *waits.entry(outcome).or_insert(0) += point.count(); + } + } _ => {} } } + assert_eq!( + (commands, waits), + (expected_commands.clone(), expected_commands) + ); let mut expected_counters = BTreeMap::new(); let mut expected_durations = BTreeMap::new(); let captures = prewarm_fails_first diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs index 16fe4d138a..ad84835fe7 100644 --- a/codex-rs/exec-server/src/telemetry.rs +++ b/codex-rs/exec-server/src/telemetry.rs @@ -217,6 +217,28 @@ impl ExecServerTelemetry { let _ = metrics.counter("codex.shell_snapshot", /*inc*/ 1, &tags); } + /// One observation per eligible execution preparation, excluding prewarm. + /// `used` means selected for replay, not that restoration or execution succeeded. + #[cfg(unix)] + pub(crate) fn shell_snapshot_command( + &self, + wait: Duration, + state: &'static str, + outcome: &'static str, + ) { + let Some(metrics) = self + .inner + .as_ref() + .map(|inner| inner.metrics.clone()) + .or_else(codex_otel::global) + else { + return; + }; + let tags = [("version", "v2"), ("state", state), ("outcome", outcome)]; + let _ = metrics.counter("codex.shell_snapshot.command", /*inc*/ 1, &tags); + let _ = metrics.record_duration("codex.shell_snapshot.wait_ms", wait, &tags); + } + pub(crate) fn remote_registration_completed(&self, result: &'static str, duration: Duration) { self.record_operation(REMOTE_REGISTRATION_METRICS, result, duration); }