diff --git a/codex-rs/cli/src/exec_server_telemetry.rs b/codex-rs/cli/src/exec_server_telemetry.rs index 4d688a015..1645bd574 100644 --- a/codex-rs/cli/src/exec_server_telemetry.rs +++ b/codex-rs/cli/src/exec_server_telemetry.rs @@ -7,7 +7,7 @@ const OTEL_SERVICE_NAME: &str = "codex-exec-server"; pub(crate) fn init( config: Option<&codex_core::config::Config>, -) -> Result> { +) -> (impl Send + Sync, codex_exec_server::ExecServerTelemetry) { let fmt_layer = tracing_subscriber::fmt::layer() .with_writer(std::io::stderr) .with_filter(stderr_env_filter()); @@ -17,21 +17,30 @@ pub(crate) fn init( env!("CARGO_PKG_VERSION"), Some(OTEL_SERVICE_NAME), DEFAULT_ANALYTICS_ENABLED, - ), - None => Ok(None), + ) + .unwrap_or_else(|error| { + eprintln!("Could not create otel exporter: {error}"); + None + }), + None => None, }; - let provider = otel.as_ref().ok().and_then(Option::as_ref); + let provider = otel.as_ref(); codex_core::otel_init::record_process_start(provider, OTEL_SERVICE_NAME); let otel_logger_layer = provider.and_then(|otel| otel.logger_layer()); let otel_tracing_layer = provider.and_then(|otel| otel.tracing_layer()); + let telemetry = provider + .and_then(|otel| otel.metrics()) + .cloned() + .map(codex_exec_server::ExecServerTelemetry::new) + .unwrap_or_default(); let _ = tracing_subscriber::registry() .with(fmt_layer) .with(otel_tracing_layer) .with(otel_logger_layer) .try_init(); tracing::callsite::rebuild_interest_cache(); - otel + (otel, telemetry) } fn stderr_env_filter() -> EnvFilter { diff --git a/codex-rs/cli/src/main.rs b/codex-rs/cli/src/main.rs index 959ec3a3b..79af9831d 100644 --- a/codex-rs/cli/src/main.rs +++ b/codex-rs/cli/src/main.rs @@ -1693,9 +1693,7 @@ async fn run_exec_server_command( .environment_id .ok_or_else(|| anyhow::anyhow!("--environment-id is required when --remote is set"))?; let config = load_exec_server_config(root_config_overrides, strict_config).await?; - let _otel = exec_server_telemetry::init(Some(&config)) - .inspect_err(|err| eprintln!("Could not create otel exporter: {err}")) - .ok(); + let (_otel, telemetry) = exec_server_telemetry::init(Some(&config)); let auth_provider = load_exec_server_remote_auth_provider(&config, &base_url, cmd.use_agent_identity_auth) .await?; @@ -1707,6 +1705,7 @@ async fn run_exec_server_command( if let Some(name) = cmd.name { remote_config.name = name; } + let remote_config = remote_config.with_telemetry(telemetry); codex_exec_server::run_remote_environment(remote_config, runtime_paths).await?; Ok(()) } else { @@ -1716,14 +1715,12 @@ async fn run_exec_server_command( } else { config_result.ok() }; - let _otel = exec_server_telemetry::init(config.as_ref()) - .inspect_err(|err| eprintln!("Could not create otel exporter: {err}")) - .ok(); + let (_otel, telemetry) = exec_server_telemetry::init(config.as_ref()); let listen_url = cmd .listen .as_deref() .unwrap_or(codex_exec_server::DEFAULT_LISTEN_URL); - codex_exec_server::run_main(listen_url, runtime_paths) + codex_exec_server::run_main_with_telemetry(listen_url, runtime_paths, telemetry) .await .map_err(anyhow::Error::from_boxed) } diff --git a/codex-rs/cli/tests/exec_server.rs b/codex-rs/cli/tests/exec_server.rs index 5e412a9de..c54e9c40b 100644 --- a/codex-rs/cli/tests/exec_server.rs +++ b/codex-rs/cli/tests/exec_server.rs @@ -1,9 +1,20 @@ use std::path::Path; +use std::process::Stdio; +use std::time::Duration; use anyhow::Result; use predicates::prelude::PredicateBooleanExt; use predicates::str::contains; use tempfile::TempDir; +use tokio::io::AsyncBufReadExt; +use tokio::io::AsyncReadExt; +use tokio::io::AsyncWriteExt; +use tokio::io::BufReader; +use wiremock::Mock; +use wiremock::MockServer; +use wiremock::ResponseTemplate; +use wiremock::matchers::method; +use wiremock::matchers::path; fn codex_command(codex_home: &Path) -> Result { let mut cmd = assert_cmd::Command::new(codex_utils_cargo_bin::cargo_bin("codex")?); @@ -48,3 +59,225 @@ fn local_exec_server_ignores_invalid_config_without_strict_config() -> Result<() Ok(()) } + +#[tokio::test] +async fn local_exec_server_flushes_telemetry_on_stdio_disconnect() -> Result<()> { + let collector = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/metrics")) + .respond_with(ResponseTemplate::new(202)) + .mount(&collector) + .await; + let codex_home = TempDir::new()?; + let base_url = collector.uri(); + std::fs::write( + codex_home.path().join("config.toml"), + format!( + r#" +[analytics] +enabled = true + +[otel] +environment = "test" +metrics_exporter = {{ otlp-http = {{ endpoint = "{base_url}/v1/metrics", protocol = "json" }} }} +"# + ), + )?; + + let cwd = url::Url::from_directory_path(std::env::current_dir()?) + .map_err(|()| anyhow::anyhow!("could not convert cwd to file URL"))?; + #[cfg(windows)] + let argv = vec!["ping.exe", "-n", "61", "127.0.0.1"]; + #[cfg(not(windows))] + let argv = vec!["/bin/sleep", "60"]; + let codex_bin = codex_utils_cargo_bin::cargo_bin("codex")?; + let codex_home = codex_home.path().to_path_buf(); + let subprocess = async move { + let mut command = tokio::process::Command::new(codex_bin); + command + .env("CODEX_HOME", codex_home) + .env("NO_PROXY", "127.0.0.1,localhost") + .env("no_proxy", "127.0.0.1,localhost") + .args(["exec-server", "--listen", "stdio"]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .kill_on_drop(true); + let mut child = command.spawn()?; + let mut stdin = child + .stdin + .take() + .ok_or_else(|| anyhow::anyhow!("exec-server stdin was not piped"))?; + let stdout = child + .stdout + .take() + .ok_or_else(|| anyhow::anyhow!("exec-server stdout was not piped"))?; + let mut stdout = BufReader::new(stdout); + send_json_line( + &mut stdin, + &serde_json::json!({ + "id": 1, + "method": "initialize", + "params": {"clientName": "otel-test", "resumeSessionId": null} + }), + ) + .await?; + wait_for_response(&mut stdout, /*expected_id*/ 1).await?; + send_json_line( + &mut stdin, + &serde_json::json!({"method": "initialized", "params": {}}), + ) + .await?; + send_json_line( + &mut stdin, + &serde_json::json!({ + "id": 2, + "method": "process/start", + "params": { + "processId": "otel-process", + "argv": argv, + "cwd": cwd, + "env": {}, + "tty": false, + "pipeStdin": false, + "arg0": null + } + }), + ) + .await?; + wait_for_response(&mut stdout, /*expected_id*/ 2).await?; + drop(stdin); + let mut remaining_stdout = String::new(); + stdout.read_to_string(&mut remaining_stdout).await?; + let status = child.wait().await?; + anyhow::ensure!( + status.success(), + "exec-server exited with {status}; remaining stdout: {remaining_stdout}" + ); + Ok::<(), anyhow::Error>(()) + }; + let subprocess_result = tokio::time::timeout(Duration::from_secs(30), subprocess) + .await + .map_err(|_| anyhow::anyhow!("exec-server subprocess timed out"))?; + subprocess_result?; + + let requests = collector + .received_requests() + .await + .ok_or_else(|| anyhow::anyhow!("failed to read OTLP collector requests"))?; + let metrics = requests + .iter() + .filter(|request| request.url.path() == "/v1/metrics") + .map(|request| serde_json::from_slice::(&request.body)) + .collect::>>()?; + assert_metric_point( + &metrics, + "exec_server_connections_active", + &[("transport", "stdio")], + Some(0), + ); + assert_metric_point( + &metrics, + "exec_server_connections_total", + &[("transport", "stdio")], + Some(1), + ); + assert_metric_point( + &metrics, + "exec_server_requests_total", + &[("method", "process/start"), ("result", "success")], + Some(1), + ); + assert_metric_point(&metrics, "exec_server_processes_active", &[], Some(0)); + assert_metric_point( + &metrics, + "exec_server_processes_finished_total", + &[("result", "terminated")], + Some(1), + ); + assert_metric_point( + &metrics, + "exec_server_request_duration_seconds", + &[("method", "process/start"), ("result", "success")], + /*value*/ None, + ); + assert_metric_point( + &metrics, + "exec_server_process_duration_seconds", + &[("result", "terminated")], + /*value*/ None, + ); + Ok(()) +} + +async fn send_json_line( + stdin: &mut (impl tokio::io::AsyncWrite + Unpin), + message: &serde_json::Value, +) -> Result<()> { + let mut encoded = serde_json::to_vec(message)?; + encoded.push(b'\n'); + stdin.write_all(&encoded).await?; + stdin.flush().await?; + Ok(()) +} + +async fn wait_for_response( + stdout: &mut (impl tokio::io::AsyncBufRead + Unpin), + expected_id: i64, +) -> Result<()> { + loop { + let mut line = String::new(); + if stdout.read_line(&mut line).await? == 0 { + anyhow::bail!("exec-server stdout closed before response {expected_id}"); + } + let message: serde_json::Value = serde_json::from_str(&line)?; + if message["id"].as_i64() == Some(expected_id) { + anyhow::ensure!( + message.get("error").is_none(), + "exec-server request {expected_id} failed: {message}" + ); + return Ok(()); + } + } +} + +fn assert_metric_point( + payloads: &[serde_json::Value], + name: &str, + attributes: &[(&str, &str)], + value: Option, +) { + let found = payloads + .iter() + .flat_map(|payload| payload["resourceMetrics"].as_array().into_iter().flatten()) + .flat_map(|resource| resource["scopeMetrics"].as_array().into_iter().flatten()) + .flat_map(|scope| scope["metrics"].as_array().into_iter().flatten()) + .filter(|metric| metric["name"].as_str() == Some(name)) + .flat_map(|metric| { + ["gauge", "sum", "histogram"] + .into_iter() + .find_map(|kind| metric[kind]["dataPoints"].as_array()) + .into_iter() + .flatten() + }) + .any(|point| { + let actual_attributes = point["attributes"] + .as_array() + .map(Vec::as_slice) + .unwrap_or_default(); + let attributes_match = actual_attributes.len() == attributes.len() + && attributes.iter().all(|(expected_key, expected_value)| { + actual_attributes.iter().any(|actual| { + actual["key"].as_str() == Some(*expected_key) + && actual["value"]["stringValue"].as_str() == Some(*expected_value) + }) + }); + let actual_value = point["asInt"] + .as_i64() + .or_else(|| point["asInt"].as_str()?.parse().ok()); + attributes_match && value.is_none_or(|expected| actual_value == Some(expected)) + }); + assert!( + found, + "metric {name} with attributes {attributes:?} and value {value:?} missing" + ); +} diff --git a/codex-rs/exec-server/src/lib.rs b/codex-rs/exec-server/src/lib.rs index 92c7f85ed..87593374c 100644 --- a/codex-rs/exec-server/src/lib.rs +++ b/codex-rs/exec-server/src/lib.rs @@ -27,6 +27,7 @@ mod rpc; mod runtime_paths; mod sandboxed_file_system; mod server; +mod telemetry; use codex_exec_server_protocol as protocol; @@ -149,3 +150,5 @@ pub use runtime_paths::ExecServerRuntimePaths; pub use server::DEFAULT_LISTEN_URL; pub use server::ExecServerListenUrlParseError; pub use server::run_main; +pub use server::run_main_with_telemetry; +pub use telemetry::ExecServerTelemetry; diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 973019808..7335bd1ee 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -59,6 +59,8 @@ use crate::rpc::RpcServerOutboundMessage; use crate::rpc::internal_error; use crate::rpc::invalid_params; use crate::rpc::invalid_request; +use crate::telemetry::ExecServerTelemetry; +use crate::telemetry::ProcessMetricGuard; const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024; const NOTIFICATION_CHANNEL_CAPACITY: usize = 256; @@ -91,6 +93,8 @@ struct RunningProcess { output_notify: Arc, open_streams: usize, closed: bool, + metrics: Option, + termination_requested: bool, sandbox: SandboxType, sandbox_denied: bool, } @@ -136,6 +140,7 @@ enum ProcessEntry { struct Inner { notifications: std::sync::RwLock>, processes: Mutex>, + telemetry: ExecServerTelemetry, } #[derive(Clone)] @@ -166,24 +171,31 @@ impl LocalProcess { let (outgoing_tx, mut outgoing_rx) = mpsc::channel::(NOTIFICATION_CHANNEL_CAPACITY); tokio::spawn(async move { while outgoing_rx.recv().await.is_some() {} }); - Self::with_runtime_paths(RpcNotificationSender::new(outgoing_tx), runtime_paths) + Self::with_runtime_paths( + RpcNotificationSender::new(outgoing_tx), + ExecServerTelemetry::default(), + runtime_paths, + ) } pub(crate) fn new( notifications: RpcNotificationSender, + telemetry: ExecServerTelemetry, runtime_paths: ExecServerRuntimePaths, ) -> Self { - Self::with_runtime_paths(notifications, Some(runtime_paths)) + Self::with_runtime_paths(notifications, telemetry, Some(runtime_paths)) } fn with_runtime_paths( notifications: RpcNotificationSender, + telemetry: ExecServerTelemetry, runtime_paths: Option, ) -> Self { Self { inner: Arc::new(Inner { notifications: std::sync::RwLock::new(Some(notifications)), processes: Mutex::new(HashMap::new()), + telemetry, }), runtime_paths, } @@ -200,7 +212,10 @@ impl LocalProcess { }) .collect::>() }; - for process in remaining { + for mut process in remaining { + if let Some(metrics) = process.metrics.take() { + metrics.finish("terminated"); + } process.session.terminate(); } } @@ -319,12 +334,13 @@ impl LocalProcess { output_notify: Arc::clone(&output_notify), open_streams: 2, closed: false, + metrics: Some(self.inner.telemetry.process_started()), + termination_requested: false, sandbox: prepared.sandbox, sandbox_denied: false, })), ); } - tokio::spawn(stream_output( process_id.clone(), if params.tty { @@ -536,11 +552,12 @@ impl LocalProcess { ) -> Result { let running = { let mut process_map = self.inner.processes.lock().await; - match process_map.get(¶ms.process_id) { + match process_map.get_mut(¶ms.process_id) { Some(ProcessEntry::Running(process)) => { if process.exit_code.is_some() { return Ok(TerminateResponse { running: false }); } + process.termination_requested = true; process.session.terminate(); true } @@ -807,11 +824,23 @@ async fn watch_exit( ) { let exit_code = exit_rx.await.unwrap_or(-1); let sandboxed = { - let processes = inner.processes.lock().await; - matches!( - processes.get(&process_id), - Some(ProcessEntry::Running(process)) if process.sandbox != SandboxType::None - ) + let mut processes = inner.processes.lock().await; + match processes.get_mut(&process_id) { + Some(ProcessEntry::Running(process)) => { + let sandboxed = process.sandbox != SandboxType::None; + if let Some(metrics) = process.metrics.take() { + metrics.finish(if process.termination_requested { + "terminated" + } else if exit_code == 0 { + "success" + } else { + "error" + }); + } + sandboxed + } + Some(ProcessEntry::Starting(_)) | None => false, + } }; if sandboxed { let _ = tokio::time::timeout(Duration::from_millis(20), output_notify.notified()).await; @@ -945,9 +974,13 @@ fn notification_sender(inner: &Inner) -> Option { #[cfg(test)] mod tests { use super::*; + use codex_otel::MetricsConfig; use codex_protocol::config_types::ShellEnvironmentPolicyInherit; use codex_utils_path_uri::PathUri; use codex_utils_pty::ProcessDriver; + use opentelemetry_sdk::metrics::InMemoryMetricExporter; + use opentelemetry_sdk::metrics::data::AggregatedMetrics; + use opentelemetry_sdk::metrics::data::MetricData; use pretty_assertions::assert_eq; use tokio::sync::oneshot; use tokio::time::timeout; @@ -969,6 +1002,62 @@ mod tests { } } + fn telemetry_backend() -> ( + LocalProcess, + codex_otel::MetricsClient, + InMemoryMetricExporter, + ) { + let exporter = InMemoryMetricExporter::default(); + let metrics = codex_otel::MetricsClient::new(MetricsConfig::in_memory( + "test", + "codex-exec-server", + env!("CARGO_PKG_VERSION"), + exporter.clone(), + )) + .expect("metrics"); + let telemetry = ExecServerTelemetry::new(metrics.clone()); + let (outgoing_tx, mut outgoing_rx) = mpsc::channel(NOTIFICATION_CHANNEL_CAPACITY); + tokio::spawn(async move { while outgoing_rx.recv().await.is_some() {} }); + ( + LocalProcess::with_runtime_paths( + RpcNotificationSender::new(outgoing_tx), + telemetry, + /*runtime_paths*/ None, + ), + metrics, + exporter, + ) + } + + fn assert_finished_process_result( + metrics: codex_otel::MetricsClient, + exporter: &InMemoryMetricExporter, + expected: &str, + ) { + metrics.shutdown().expect("shutdown metrics"); + let resource_metrics = exporter + .get_finished_metrics() + .expect("finished metrics") + .into_iter() + .last() + .expect("metrics export"); + let finished_processes = resource_metrics + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .find(|metric| metric.name() == "exec_server_processes_finished_total") + .expect("finished process metric"); + let AggregatedMetrics::U64(MetricData::Sum(sum)) = finished_processes.data() else { + panic!("finished process metric should be a u64 sum"); + }; + let results = sum + .data_points() + .flat_map(opentelemetry_sdk::metrics::data::SumDataPoint::attributes) + .filter(|attribute| attribute.key.as_str() == "result") + .map(|attribute| attribute.value.as_str().into_owned()) + .collect::>(); + assert_eq!(results, vec![expected.to_string()]); + } + #[tokio::test] async fn start_process_rejects_non_native_cwd_before_launch() { #[cfg(unix)] @@ -1028,6 +1117,50 @@ mod tests { assert_eq!(child_env(¶ms), expected); } + #[tokio::test] + async fn exit_before_shutdown_records_success() { + let (backend, metrics, exporter) = telemetry_backend(); + let mut process = spawn_test_process(&backend, "exit-before-shutdown").await; + + process.exit(/*exit_code*/ 0); + let _ = read_process_until_change(&backend, &process.process_id, /*after_seq*/ None).await; + backend.shutdown().await; + + assert_finished_process_result(metrics, &exporter, "success"); + } + + #[tokio::test] + async fn termination_request_before_exit_records_terminated() { + let (backend, metrics, exporter) = telemetry_backend(); + let mut process = spawn_test_process(&backend, "terminate-before-exit").await; + + assert_eq!( + backend + .terminate_process(TerminateParams { + process_id: process.process_id.clone(), + }) + .await + .expect("terminate process"), + TerminateResponse { running: true }, + ); + process.exit(/*exit_code*/ 0); + let _ = read_process_until_change(&backend, &process.process_id, /*after_seq*/ None).await; + backend.shutdown().await; + + assert_finished_process_result(metrics, &exporter, "terminated"); + } + + #[tokio::test] + async fn shutdown_before_exit_records_terminated() { + let (backend, metrics, exporter) = telemetry_backend(); + let mut process = spawn_test_process(&backend, "shutdown-before-exit").await; + + backend.shutdown().await; + process.exit(/*exit_code*/ 0); + + assert_finished_process_result(metrics, &exporter, "terminated"); + } + #[tokio::test] async fn exited_process_retains_late_output_past_retention() { let backend = LocalProcess::default(); @@ -1170,6 +1303,8 @@ mod tests { output_notify: Arc::clone(&output_notify), open_streams: 2, closed: false, + metrics: Some(backend.inner.telemetry.process_started()), + termination_requested: false, sandbox: SandboxType::None, sandbox_denied: false, })), diff --git a/codex-rs/exec-server/src/noise_relay/executor_stream.rs b/codex-rs/exec-server/src/noise_relay/executor_stream.rs index 57a05579d..89f33d938 100644 --- a/codex-rs/exec-server/src/noise_relay/executor_stream.rs +++ b/codex-rs/exec-server/src/noise_relay/executor_stream.rs @@ -27,6 +27,7 @@ use crate::relay::encode_relay_message_frame; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; use crate::server::ConnectionProcessor; +use crate::telemetry::ConnectionTransport; /// Identifies one completed virtual-stream instance. /// @@ -169,7 +170,9 @@ pub(crate) fn spawn_noise_virtual_stream( transport: JsonRpcTransport::Plain, }; tokio::spawn(async move { - processor.run_connection(connection).await; + processor + .run_connection(connection, ConnectionTransport::Relay) + .await; let _ = processor_closed_stream_tx .send(ClosedNoiseVirtualStream { stream_id: processor_stream_id, diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index 199e78754..6609b6cd9 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -25,6 +25,7 @@ use crate::EnvironmentRegistryRegistrationRequest; use crate::EnvironmentRegistryRegistrationResponse; use crate::ExecServerError; use crate::ExecServerRuntimePaths; +use crate::ExecServerTelemetry; use crate::NoiseChannelIdentity; use crate::NoiseChannelPublicKey; use crate::NoiseRendezvousConnectBundle; @@ -373,6 +374,7 @@ pub struct RemoteEnvironmentConfig { pub environment_id: String, pub name: String, auth_provider: SharedAuthProvider, + telemetry: ExecServerTelemetry, } impl std::fmt::Debug for RemoteEnvironmentConfig { @@ -398,8 +400,14 @@ impl RemoteEnvironmentConfig { environment_id, name: "codex-exec-server".to_string(), auth_provider, + telemetry: ExecServerTelemetry::default(), }) } + + pub fn with_telemetry(mut self, telemetry: ExecServerTelemetry) -> Self { + self.telemetry = telemetry; + self + } } /// Register an exec-server for remote use and serve requests over Noise. @@ -420,7 +428,8 @@ pub async fn run_remote_environment( ensure_rustls_crypto_provider(); let client = EnvironmentRegistryClient::new(config.base_url.clone(), config.auth_provider.clone())?; - let processor = ConnectionProcessor::new(runtime_paths); + let processor = + ConnectionProcessor::new_with_telemetry(runtime_paths, config.telemetry.clone()); let identity = NoiseChannelIdentity::generate().map_err(|error| { ExecServerError::Protocol(format!("failed to generate Noise relay identity: {error}")) })?; diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index 4613c41f6..a941dce41 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -212,8 +212,10 @@ where ); } - pub(crate) fn request_route(&self, method: &str) -> Option<&RequestRoute> { - self.request_routes.get(method) + pub(crate) fn request_route(&self, method: &str) -> Option<(&'static str, &RequestRoute)> { + self.request_routes + .get_key_value(method) + .map(|(&method, route)| (method, route)) } pub(crate) fn notification_route(&self, method: &str) -> Option<&NotificationRoute> { diff --git a/codex-rs/exec-server/src/server.rs b/codex-rs/exec-server/src/server.rs index d443b0ed1..7988e0870 100644 --- a/codex-rs/exec-server/src/server.rs +++ b/codex-rs/exec-server/src/server.rs @@ -12,15 +12,71 @@ pub use transport::DEFAULT_LISTEN_URL; pub use transport::ExecServerListenUrlParseError; use crate::ExecServerRuntimePaths; +use crate::ExecServerTelemetry; + +pub async fn run_main( + listen_url: &str, + runtime_paths: ExecServerRuntimePaths, +) -> Result<(), Box> { + run_main_with_telemetry(listen_url, runtime_paths, ExecServerTelemetry::default()).await +} #[tracing::instrument( name = "codex.exec_server", skip_all, fields(otel.kind = "internal") )] -pub async fn run_main( +pub async fn run_main_with_telemetry( listen_url: &str, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { - transport::run_transport(listen_url, runtime_paths).await + transport::run_transport(listen_url, runtime_paths, telemetry).await +} + +#[cfg(test)] +mod tests { + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::InMemorySpanExporter; + use opentelemetry_sdk::trace::SdkTracerProvider; + use tracing::instrument::WithSubscriber; + use tracing_subscriber::prelude::*; + + use super::run_main_with_telemetry; + use crate::ExecServerRuntimePaths; + use crate::ExecServerTelemetry; + + #[tokio::test] + async fn telemetry_entrypoint_emits_root_span() { + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let subscriber = tracing_subscriber::registry() + .with(tracing_opentelemetry::layer().with_tracer(provider.tracer("exec-server-test"))); + + async { + tracing::callsite::rebuild_interest_cache(); + run_main_with_telemetry( + "invalid", + ExecServerRuntimePaths::new( + std::env::current_exe().expect("current executable"), + /*codex_linux_sandbox_exe*/ None, + ) + .expect("runtime paths"), + ExecServerTelemetry::default(), + ) + .await + .expect_err("invalid listen URL should fail"); + } + .with_subscriber(subscriber) + .await; + + provider.force_flush().expect("flush traces"); + let spans = exporter.get_finished_spans().expect("span export"); + assert!( + spans.iter().any(|span| span.name == "codex.exec_server"), + "root exec-server span missing: {spans:?}" + ); + } } diff --git a/codex-rs/exec-server/src/server/handler/tests.rs b/codex-rs/exec-server/src/server/handler/tests.rs index 2cea99b21..47416039c 100644 --- a/codex-rs/exec-server/src/server/handler/tests.rs +++ b/codex-rs/exec-server/src/server/handler/tests.rs @@ -82,7 +82,7 @@ fn test_runtime_paths() -> ExecServerRuntimePaths { async fn initialized_handler() -> Arc { let (outgoing_tx, _outgoing_rx) = mpsc::channel(16); - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let handler = Arc::new(ExecServerHandler::new( registry, RpcNotificationSender::new(outgoing_tx), @@ -160,7 +160,7 @@ async fn terminate_reports_false_after_process_exit() { #[tokio::test] async fn long_poll_read_fails_after_session_resume() { let (first_tx, _first_rx) = mpsc::channel(16); - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let first_handler = Arc::new(ExecServerHandler::new( Arc::clone(®istry), RpcNotificationSender::new(first_tx), @@ -233,7 +233,7 @@ async fn long_poll_read_fails_after_session_resume() { #[tokio::test] async fn active_session_resume_is_rejected() { let (first_tx, _first_rx) = mpsc::channel(16); - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let first_handler = Arc::new(ExecServerHandler::new( Arc::clone(®istry), RpcNotificationSender::new(first_tx), @@ -277,7 +277,7 @@ async fn active_session_resume_is_rejected() { async fn output_and_exit_are_retained_after_notification_receiver_closes() { let (outgoing_tx, outgoing_rx) = mpsc::channel(16); let handler = Arc::new(ExecServerHandler::new( - SessionRegistry::new(), + SessionRegistry::new(crate::ExecServerTelemetry::default()), RpcNotificationSender::new(outgoing_tx), test_runtime_paths(), )); diff --git a/codex-rs/exec-server/src/server/process_handler.rs b/codex-rs/exec-server/src/server/process_handler.rs index 9ff669041..4d933afbe 100644 --- a/codex-rs/exec-server/src/server/process_handler.rs +++ b/codex-rs/exec-server/src/server/process_handler.rs @@ -13,6 +13,7 @@ use crate::protocol::TerminateResponse; use crate::protocol::WriteParams; use crate::protocol::WriteResponse; use crate::rpc::RpcNotificationSender; +use crate::telemetry::ExecServerTelemetry; #[derive(Clone)] pub(crate) struct ProcessHandler { @@ -22,10 +23,11 @@ pub(crate) struct ProcessHandler { impl ProcessHandler { pub(crate) fn new( notifications: RpcNotificationSender, + telemetry: ExecServerTelemetry, runtime_paths: ExecServerRuntimePaths, ) -> Self { Self { - process: LocalProcess::new(notifications, runtime_paths), + process: LocalProcess::new(notifications, telemetry, runtime_paths), } } diff --git a/codex-rs/exec-server/src/server/processor.rs b/codex-rs/exec-server/src/server/processor.rs index 3eedfb2f2..63f6bd45d 100644 --- a/codex-rs/exec-server/src/server/processor.rs +++ b/codex-rs/exec-server/src/server/processor.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::time::Instant; use tokio::sync::mpsc; use tracing::Instrument; @@ -17,36 +18,61 @@ use crate::rpc::method_not_found; use crate::server::ExecServerHandler; use crate::server::registry::build_router; use crate::server::session_registry::SessionRegistry; +use crate::telemetry::ConnectionTransport; +use crate::telemetry::ExecServerTelemetry; #[derive(Clone)] pub(crate) struct ConnectionProcessor { session_registry: Arc, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, } impl ConnectionProcessor { + #[cfg(test)] pub(crate) fn new(runtime_paths: ExecServerRuntimePaths) -> Self { + Self::new_with_telemetry(runtime_paths, ExecServerTelemetry::default()) + } + + pub(crate) fn new_with_telemetry( + runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, + ) -> Self { Self { - session_registry: SessionRegistry::new(), + session_registry: SessionRegistry::new(telemetry.clone()), runtime_paths, + telemetry, } } - pub(crate) async fn run_connection(&self, connection: JsonRpcConnection) { + pub(crate) async fn run_connection( + &self, + connection: JsonRpcConnection, + transport: ConnectionTransport, + ) { run_connection( connection, Arc::clone(&self.session_registry), self.runtime_paths.clone(), + self.telemetry.clone(), + transport, ) .await; } + + pub(crate) async fn shutdown(&self) { + self.session_registry.shutdown().await; + } } async fn run_connection( connection: JsonRpcConnection, session_registry: Arc, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, + transport: ConnectionTransport, ) { + let _connection_metrics = telemetry.connection_started(transport); let router = Arc::new(build_router()); let JsonRpcConnection { outgoing_tx: json_outgoing_tx, @@ -101,12 +127,18 @@ async fn run_connection( } JsonRpcConnectionEvent::Message(message) => match message { codex_exec_server_protocol::JSONRPCMessage::Request(request) => { - if let Some(route) = router.request_route(request.method.as_str()) { - let request_span = request_span(request.method.as_str(), &request); + let request_started_at = Instant::now(); + if let Some((method, route)) = router.request_route(request.method.as_str()) { + let request_span = request_span(method, &request); let message = tokio::select! { message = route(Arc::clone(&handler), request).instrument(request_span.clone()) => message, _ = disconnected_rx.changed() => { request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); debug!("exec-server transport disconnected while handling request"); break; } @@ -116,11 +148,19 @@ async fn run_connection( && outgoing_tx.send(message).await.is_err() { request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); break; } request_span.record("result", result); + telemetry.request_completed(method, result, request_started_at.elapsed()); + drop(request_span); } else { - let request_span = request_span("unknown", &request); + let method = "unknown"; + let request_span = request_span(method, &request); if outgoing_tx .send(RpcServerOutboundMessage::Error { request_id: request.id, @@ -133,9 +173,15 @@ async fn run_connection( .is_err() { request_span.record("result", "disconnected"); + telemetry.request_completed( + method, + "disconnected", + request_started_at.elapsed(), + ); break; } request_span.record("result", "error"); + telemetry.request_completed(method, "error", request_started_at.elapsed()); } } codex_exec_server_protocol::JSONRPCMessage::Notification(notification) => { @@ -330,7 +376,7 @@ mod tests { #[tokio::test] async fn connection_accepts_pipelined_scalar_requests() { - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let (mut writer, mut lines, task) = spawn_test_connection(registry, "pipelined-scalar"); send_request( @@ -362,7 +408,7 @@ mod tests { #[tokio::test] async fn transport_disconnect_detaches_session_during_in_flight_read() { - let registry = SessionRegistry::new(); + let registry = SessionRegistry::new(crate::ExecServerTelemetry::default()); let (mut first_writer, mut first_lines, first_task) = spawn_test_connection(Arc::clone(®istry), "first"); @@ -458,7 +504,13 @@ mod tests { let (server_writer, client_reader) = duplex(1 << 20); let connection = JsonRpcConnection::from_stdio(server_reader, server_writer, label.to_string()); - let task = tokio::spawn(run_connection(connection, registry, test_runtime_paths())); + let task = tokio::spawn(run_connection( + connection, + registry, + test_runtime_paths(), + crate::ExecServerTelemetry::default(), + crate::telemetry::ConnectionTransport::Stdio, + )); (client_writer, BufReader::new(client_reader).lines(), task) } diff --git a/codex-rs/exec-server/src/server/session_registry.rs b/codex-rs/exec-server/src/server/session_registry.rs index 4d2ef7f89..89a861ab7 100644 --- a/codex-rs/exec-server/src/server/session_registry.rs +++ b/codex-rs/exec-server/src/server/session_registry.rs @@ -12,6 +12,7 @@ use crate::rpc::RpcNotificationSender; use crate::rpc::invalid_request; use crate::rpc::session_already_attached; use crate::server::process_handler::ProcessHandler; +use crate::telemetry::ExecServerTelemetry; #[cfg(test)] const DETACHED_SESSION_TTL: Duration = Duration::from_millis(200); @@ -20,6 +21,7 @@ const DETACHED_SESSION_TTL: Duration = Duration::from_secs(30); pub(crate) struct SessionRegistry { sessions: Mutex>>, + telemetry: ExecServerTelemetry, } struct SessionEntry { @@ -51,9 +53,10 @@ pub(crate) struct SessionHandle { } impl SessionRegistry { - pub(crate) fn new() -> Arc { + pub(crate) fn new(telemetry: ExecServerTelemetry) -> Arc { Arc::new(Self { sessions: Mutex::new(HashMap::new()), + telemetry, }) } @@ -97,7 +100,7 @@ impl SessionRegistry { let session_id = Uuid::new_v4().to_string(); let entry = Arc::new(SessionEntry::new( session_id.clone(), - ProcessHandler::new(notifications, runtime_paths), + ProcessHandler::new(notifications, self.telemetry.clone(), runtime_paths), connection_id, )); sessions.insert(session_id, Arc::clone(&entry)); @@ -119,6 +122,13 @@ impl SessionRegistry { }) } + pub(crate) async fn shutdown(&self) { + let sessions = std::mem::take(&mut *self.sessions.lock().await); + for entry in sessions.into_values() { + entry.process.shutdown().await; + } + } + async fn expire_if_detached(&self, session_id: String, connection_id: ConnectionId) { tokio::time::sleep(DETACHED_SESSION_TTL).await; @@ -143,6 +153,7 @@ impl Default for SessionRegistry { fn default() -> Self { Self { sessions: Mutex::new(HashMap::new()), + telemetry: ExecServerTelemetry::default(), } } } diff --git a/codex-rs/exec-server/src/server/transport.rs b/codex-rs/exec-server/src/server/transport.rs index 57d7048e2..2affab007 100644 --- a/codex-rs/exec-server/src/server/transport.rs +++ b/codex-rs/exec-server/src/server/transport.rs @@ -22,8 +22,10 @@ use tracing::info; use tracing::warn; use crate::ExecServerRuntimePaths; +use crate::ExecServerTelemetry; use crate::connection::JsonRpcConnection; use crate::server::processor::ConnectionProcessor; +use crate::telemetry::ConnectionTransport; pub const DEFAULT_LISTEN_URL: &str = "ws://127.0.0.1:0"; @@ -80,49 +82,54 @@ pub(crate) fn parse_listen_url( pub(crate) async fn run_transport( listen_url: &str, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { match parse_listen_url(listen_url)? { ExecServerListenTransport::WebSocket(bind_address) => { - run_websocket_listener(bind_address, runtime_paths).await + run_websocket_listener(bind_address, runtime_paths, telemetry).await } - ExecServerListenTransport::Stdio => run_stdio_connection(runtime_paths).await, + ExecServerListenTransport::Stdio => run_stdio_connection(runtime_paths, telemetry).await, } } async fn run_stdio_connection( runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { - run_stdio_connection_with_io(io::stdin(), io::stdout(), runtime_paths).await + run_stdio_connection_with_io(io::stdin(), io::stdout(), runtime_paths, telemetry).await } async fn run_stdio_connection_with_io( reader: R, writer: W, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> where R: AsyncRead + Unpin + Send + 'static, W: AsyncWrite + Unpin + Send + 'static, { - let processor = ConnectionProcessor::new(runtime_paths); + let processor = ConnectionProcessor::new_with_telemetry(runtime_paths, telemetry); tracing::info!("codex-exec-server listening on stdio"); processor - .run_connection(JsonRpcConnection::from_stdio( - reader, - writer, - "exec-server stdio".to_string(), - )) + .run_connection( + JsonRpcConnection::from_stdio(reader, writer, "exec-server stdio".to_string()), + ConnectionTransport::Stdio, + ) .await; + // Stdio serves exactly one connection, so detached sessions cannot be resumed. + processor.shutdown().await; Ok(()) } async fn run_websocket_listener( bind_address: SocketAddr, runtime_paths: ExecServerRuntimePaths, + telemetry: ExecServerTelemetry, ) -> Result<(), Box> { let listener = TcpListener::bind(bind_address).await?; let local_addr = listener.local_addr()?; - let processor = ConnectionProcessor::new(runtime_paths); + let processor = ConnectionProcessor::new_with_telemetry(runtime_paths, telemetry); info!("codex-exec-server listening on ws://{local_addr}"); println!("ws://{local_addr}"); std::io::stdout().flush()?; @@ -174,10 +181,13 @@ async fn websocket_upgrade_handler( websocket.on_upgrade(move |stream| async move { state .processor - .run_connection(JsonRpcConnection::from_axum_websocket( - stream, - format!("exec-server websocket {peer_addr}"), - )) + .run_connection( + JsonRpcConnection::from_axum_websocket( + stream, + format!("exec-server websocket {peer_addr}"), + ), + ConnectionTransport::WebSocket, + ) .await; }) } diff --git a/codex-rs/exec-server/src/server/transport_tests.rs b/codex-rs/exec-server/src/server/transport_tests.rs index a8db4250b..3016effa6 100644 --- a/codex-rs/exec-server/src/server/transport_tests.rs +++ b/codex-rs/exec-server/src/server/transport_tests.rs @@ -61,6 +61,7 @@ async fn stdio_listen_transport_serves_initialize() { server_reader, server_writer, test_runtime_paths(), + crate::ExecServerTelemetry::default(), )); let mut client_lines = BufReader::new(client_reader).lines(); diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs new file mode 100644 index 000000000..4624a2ff2 --- /dev/null +++ b/codex-rs/exec-server/src/telemetry.rs @@ -0,0 +1,282 @@ +use std::sync::Arc; +use std::sync::Mutex; +use std::time::Duration; +use std::time::Instant; + +use codex_otel::MetricsClient; +use tracing::warn; + +const CONNECTIONS_ACTIVE_METRIC: &str = "exec_server_connections_active"; +const CONNECTIONS_ACTIVE_DESCRIPTION: &str = "Number of active exec-server connections."; +const CONNECTIONS_TOTAL_METRIC: &str = "exec_server_connections_total"; +const CONNECTIONS_TOTAL_DESCRIPTION: &str = "Total number of accepted exec-server connections."; +const REQUESTS_TOTAL_METRIC: &str = "exec_server_requests_total"; +const REQUESTS_TOTAL_DESCRIPTION: &str = "Total number of exec-server requests."; +const REQUEST_DURATION_METRIC: &str = "exec_server_request_duration_seconds"; +const REQUEST_DURATION_DESCRIPTION: &str = "Duration of exec-server requests in seconds."; +const PROCESSES_ACTIVE_METRIC: &str = "exec_server_processes_active"; +const PROCESSES_ACTIVE_DESCRIPTION: &str = "Number of active exec-server processes."; +const PROCESSES_FINISHED_TOTAL_METRIC: &str = "exec_server_processes_finished_total"; +const PROCESSES_FINISHED_TOTAL_DESCRIPTION: &str = + "Total number of finished exec-server processes."; +const PROCESS_DURATION_METRIC: &str = "exec_server_process_duration_seconds"; +const PROCESS_DURATION_DESCRIPTION: &str = "Duration of exec-server processes in seconds."; + +#[derive(Clone, Copy)] +pub(crate) enum ConnectionTransport { + Relay, + Stdio, + WebSocket, +} + +impl ConnectionTransport { + fn metric_tag(self) -> &'static str { + match self { + Self::Relay => "relay", + Self::Stdio => "stdio", + Self::WebSocket => "websocket", + } + } +} + +#[derive(Clone, Default)] +pub struct ExecServerTelemetry { + inner: Option>, +} + +struct ExecServerTelemetryInner { + metrics: MetricsClient, + active: Arc>, +} + +#[derive(Default)] +struct ActiveCounts { + relay_connections: i64, + stdio_connections: i64, + websocket_connections: i64, + processes: i64, +} + +impl ActiveCounts { + fn connections(&self, transport: ConnectionTransport) -> i64 { + match transport { + ConnectionTransport::Relay => self.relay_connections, + ConnectionTransport::Stdio => self.stdio_connections, + ConnectionTransport::WebSocket => self.websocket_connections, + } + } +} + +pub(crate) struct ConnectionMetricGuard { + telemetry: ExecServerTelemetry, + transport: ConnectionTransport, +} + +pub(crate) struct ProcessMetricGuard { + telemetry: ExecServerTelemetry, + started_at: Instant, + result: &'static str, +} + +impl ExecServerTelemetry { + pub fn new(metrics: MetricsClient) -> Self { + let active = Arc::new(Mutex::new(ActiveCounts::default())); + register_active_gauges(&metrics, &active); + Self { + inner: Some(Arc::new(ExecServerTelemetryInner { metrics, active })), + } + } + + pub(crate) fn connection_started( + &self, + transport: ConnectionTransport, + ) -> ConnectionMetricGuard { + self.with_inner(|inner| { + inner.adjust_connection_count(transport, /*delta*/ 1); + inner.counter( + CONNECTIONS_TOTAL_METRIC, + CONNECTIONS_TOTAL_DESCRIPTION, + &[("transport", transport.metric_tag())], + ); + }); + ConnectionMetricGuard { + telemetry: self.clone(), + transport, + } + } + + pub(crate) fn request_completed( + &self, + method: &'static str, + result: &'static str, + duration: Duration, + ) { + self.with_inner(|inner| { + let tags = [("method", method), ("result", result)]; + inner.counter(REQUESTS_TOTAL_METRIC, REQUESTS_TOTAL_DESCRIPTION, &tags); + inner.duration( + REQUEST_DURATION_METRIC, + REQUEST_DURATION_DESCRIPTION, + duration, + &tags, + ); + }); + } + + pub(crate) fn process_started(&self) -> ProcessMetricGuard { + self.with_inner(|inner| { + inner.adjust_process_count(/*delta*/ 1); + }); + ProcessMetricGuard { + telemetry: self.clone(), + started_at: Instant::now(), + result: "unknown", + } + } + + fn process_finished(&self, result: &'static str, duration: Duration) { + self.with_inner(|inner| { + inner.adjust_process_count(/*delta*/ -1); + inner.counter( + PROCESSES_FINISHED_TOTAL_METRIC, + PROCESSES_FINISHED_TOTAL_DESCRIPTION, + &[("result", result)], + ); + inner.duration( + PROCESS_DURATION_METRIC, + PROCESS_DURATION_DESCRIPTION, + duration, + &[("result", result)], + ); + }); + } + + fn connection_finished(&self, transport: ConnectionTransport) { + self.with_inner(|inner| { + inner.adjust_connection_count(transport, /*delta*/ -1); + }); + } + + fn with_inner(&self, emit: impl FnOnce(&ExecServerTelemetryInner)) { + if let Some(inner) = &self.inner { + emit(inner); + } + } +} + +impl Drop for ConnectionMetricGuard { + fn drop(&mut self) { + self.telemetry.connection_finished(self.transport); + } +} + +impl ProcessMetricGuard { + pub(crate) fn finish(mut self, result: &'static str) { + self.result = result; + } +} + +impl Drop for ProcessMetricGuard { + fn drop(&mut self) { + self.telemetry + .process_finished(self.result, self.started_at.elapsed()); + } +} + +impl ExecServerTelemetryInner { + fn active_counts(&self) -> std::sync::MutexGuard<'_, ActiveCounts> { + // These are independent integer counts, so a panic cannot leave a cross-field invariant + // half-updated. Recovering a poisoned lock preserves the last completed count update. + self.active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + fn adjust_connection_count(&self, transport: ConnectionTransport, delta: i64) { + let mut active = self.active_counts(); + let count = match transport { + ConnectionTransport::Relay => &mut active.relay_connections, + ConnectionTransport::Stdio => &mut active.stdio_connections, + ConnectionTransport::WebSocket => &mut active.websocket_connections, + }; + *count += delta; + } + + fn adjust_process_count(&self, delta: i64) { + let mut active = self.active_counts(); + active.processes += delta; + } + + fn counter(&self, name: &str, description: &str, tags: &[(&str, &str)]) { + if self + .metrics + .counter_with_description(name, description, /*inc*/ 1, tags) + .is_err() + { + warn!(metric = name, "failed to emit exec-server counter"); + } + } + + fn duration(&self, name: &str, description: &str, duration: Duration, tags: &[(&str, &str)]) { + if self + .metrics + .record_duration_seconds_with_description(name, description, duration, tags) + .is_err() + { + warn!(metric = name, "failed to emit exec-server duration"); + } + } +} + +fn register_active_gauges(metrics: &MetricsClient, active: &Arc>) { + for transport in [ + ConnectionTransport::Relay, + ConnectionTransport::Stdio, + ConnectionTransport::WebSocket, + ] { + register_active_gauge( + metrics, + active, + CONNECTIONS_ACTIVE_METRIC, + CONNECTIONS_ACTIVE_DESCRIPTION, + &[("transport", transport.metric_tag())], + move |active| active.connections(transport), + ); + } + + register_active_gauge( + metrics, + active, + PROCESSES_ACTIVE_METRIC, + PROCESSES_ACTIVE_DESCRIPTION, + &[], + |active| active.processes, + ); +} + +fn register_active_gauge( + metrics: &MetricsClient, + active: &Arc>, + name: &str, + description: &str, + tags: &[(&str, &str)], + read: impl Fn(&ActiveCounts) -> i64 + Send + Sync + 'static, +) { + let active = Arc::clone(active); + if metrics + .register_observable_gauge_with_description( + name, + description, + move || { + let active = active + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + read(&active) + }, + tags, + ) + .is_err() + { + warn!(metric = name, "failed to register exec-server gauge"); + } +} diff --git a/codex-rs/otel/src/metrics/client.rs b/codex-rs/otel/src/metrics/client.rs index 73fc3c661..251fc74d7 100644 --- a/codex-rs/otel/src/metrics/client.rs +++ b/codex-rs/otel/src/metrics/client.rs @@ -188,6 +188,24 @@ impl MetricsClientInner { Ok(()) } + fn register_observable_gauge( + &self, + name: &str, + description: &str, + observe: impl Fn() -> i64 + Send + Sync + 'static, + tags: &[(&str, &str)], + ) -> Result<()> { + validate_metric_name(name)?; + let attributes = self.attributes(tags)?; + let _gauge = self + .meter + .i64_observable_gauge(name.to_string()) + .with_description(description.to_string()) + .with_callback(move |observer| observer.observe(observe(), &attributes)) + .build(); + Ok(()) + } + fn duration_histogram( &self, name: &str, @@ -354,6 +372,18 @@ impl MetricsClient { self.0.gauge(name, Some(description), value, tags) } + /// Register a gauge callback that reports the current value on every collection. + pub fn register_observable_gauge_with_description( + &self, + name: &str, + description: &str, + observe: impl Fn() -> i64 + Send + Sync + 'static, + tags: &[(&str, &str)], + ) -> Result<()> { + self.0 + .register_observable_gauge(name, description, observe, tags) + } + /// Record a duration in milliseconds using a histogram. pub fn record_duration( &self, diff --git a/codex-rs/otel/tests/suite/snapshot.rs b/codex-rs/otel/tests/suite/snapshot.rs index e63110000..ffc2458a0 100644 --- a/codex-rs/otel/tests/suite/snapshot.rs +++ b/codex-rs/otel/tests/suite/snapshot.rs @@ -62,6 +62,38 @@ fn snapshot_collects_metrics_without_shutdown() -> Result<()> { Ok(()) } +#[test] +fn observable_gauge_is_collected_on_every_delta_snapshot() -> Result<()> { + let exporter = InMemoryMetricExporter::default(); + let config = MetricsConfig::in_memory("test", "codex-cli", env!("CARGO_PKG_VERSION"), exporter) + .with_runtime_reader(); + let metrics = MetricsClient::new(config)?; + metrics.register_observable_gauge_with_description( + "codex.active", + "Number of active operations.", + || 1, + &[("component", "test")], + )?; + + for snapshot in [metrics.snapshot()?, metrics.snapshot()?] { + let gauge = find_metric(&snapshot, "codex.active").expect("gauge metric missing"); + let point = match gauge.data() { + AggregatedMetrics::I64(MetricData::Gauge(gauge)) => { + gauge.data_points().next().expect("gauge point") + } + _ => panic!("unexpected gauge metric data type"), + }; + assert_eq!(point.value(), 1); + assert_eq!( + attributes_to_map(point.attributes()), + BTreeMap::from([("component".to_string(), "test".to_string())]) + ); + } + + metrics.shutdown()?; + Ok(()) +} + #[test] fn manager_snapshot_metrics_collects_without_shutdown() -> Result<()> { let exporter = InMemoryMetricExporter::default();