diff --git a/codex-rs/core/src/unified_exec/process.rs b/codex-rs/core/src/unified_exec/process.rs index 5f785adfa75a..e3d9c9c92d82 100644 --- a/codex-rs/core/src/unified_exec/process.rs +++ b/codex-rs/core/src/unified_exec/process.rs @@ -14,6 +14,7 @@ use tokio_util::sync::CancellationToken; use crate::exec::is_likely_sandbox_denied; use codex_exec_server::ExecProcess; +use codex_exec_server::ExecProcessEvent; use codex_exec_server::ProcessSignal as ExecServerProcessSignal; use codex_exec_server::ReadResponse as ExecReadResponse; use codex_exec_server::StartedExecProcess; @@ -425,15 +426,70 @@ impl UnifiedExecProcess { cancellation_token, } = output_handles; let process = started.process; - let mut wake_rx = process.subscribe_wake(); + let mut events = process.subscribe_events(); tokio::spawn(async move { - let mut after_seq = None; + let mut last_seq = 0; loop { - match process - .read(after_seq, /*max_bytes*/ None, /*wait_ms*/ Some(0)) - .await - { - Ok(response) => { + match events.recv().await { + Ok(ExecProcessEvent::Output(chunk)) => { + if chunk.seq > last_seq { + last_seq = chunk.seq; + let bytes = chunk.chunk.into_inner(); + let mut guard = output_buffer.lock().await; + guard.push_chunk(bytes.clone()); + drop(guard); + let _ = output_tx.send(bytes); + output_notify.notify_waiters(); + } + } + Ok(ExecProcessEvent::Exited { seq, exit_code }) => { + if seq > last_seq { + last_seq = seq; + let state = state_tx.borrow().clone(); + let _ = state_tx.send_replace(state.exited(Some(exit_code))); + } + } + Ok(ExecProcessEvent::Closed { + seq: _, + sandbox_denied, + }) => { + if sandbox_denied { + let mut state = state_tx.borrow().clone(); + state.sandbox_denied = true; + let _ = state_tx.send_replace(state); + } + output_closed.store(true, Ordering::Release); + output_closed_notify.notify_waiters(); + cancellation_token.cancel(); + break; + } + Ok(ExecProcessEvent::Failed(message)) => { + let state = state_tx.borrow().clone(); + let _ = state_tx.send_replace(state.failed(message)); + output_closed.store(true, Ordering::Release); + output_closed_notify.notify_waiters(); + cancellation_token.cancel(); + break; + } + Err(broadcast::error::RecvError::Lagged(_)) => { + let response = match process + .read( + Some(last_seq), + /*max_bytes*/ None, + /*wait_ms*/ Some(0), + ) + .await + { + Ok(response) => response, + Err(err) => { + let state = state_tx.borrow().clone(); + let _ = state_tx.send_replace(state.failed(err.to_string())); + output_closed.store(true, Ordering::Release); + output_closed_notify.notify_waiters(); + cancellation_token.cancel(); + break; + } + }; let ExecReadResponse { chunks, next_seq, @@ -443,8 +499,7 @@ impl UnifiedExecProcess { failure, sandbox_denied, } = response; - - for chunk in chunks { + for chunk in chunks.into_iter().filter(|chunk| chunk.seq > last_seq) { let bytes = chunk.chunk.into_inner(); let mut guard = output_buffer.lock().await; guard.push_chunk(bytes.clone()); @@ -452,7 +507,7 @@ impl UnifiedExecProcess { let _ = output_tx.send(bytes); output_notify.notify_waiters(); } - + last_seq = last_seq.max(next_seq.saturating_sub(1)); if let Some(message) = failure { let state = state_tx.borrow().clone(); let _ = state_tx.send_replace(state.failed(message)); @@ -461,7 +516,6 @@ impl UnifiedExecProcess { cancellation_token.cancel(); break; } - if sandbox_denied { let mut state = state_tx.borrow().clone(); state.sandbox_denied = true; @@ -472,37 +526,24 @@ impl UnifiedExecProcess { let state = state_tx.borrow().clone(); let _ = state_tx.send_replace(state.exited(exit_code)); } - if closed { output_closed.store(true, Ordering::Release); output_closed_notify.notify_waiters(); cancellation_token.cancel(); - } - - after_seq = next_seq.checked_sub(1); - if output_closed.load(Ordering::Acquire) { break; } } - Err(err) => { + Err(broadcast::error::RecvError::Closed) => { let state = state_tx.borrow().clone(); - let _ = state_tx.send_replace(state.failed(err.to_string())); + let _ = state_tx.send_replace( + state.failed("exec-server process event stream closed".to_string()), + ); output_closed.store(true, Ordering::Release); output_closed_notify.notify_waiters(); cancellation_token.cancel(); break; } } - - if wake_rx.changed().await.is_err() { - let state = state_tx.borrow().clone(); - let _ = state_tx - .send_replace(state.failed("exec-server wake channel closed".to_string())); - output_closed.store(true, Ordering::Release); - output_closed_notify.notify_waiters(); - cancellation_token.cancel(); - break; - } } }) } diff --git a/codex-rs/core/src/unified_exec/process_tests.rs b/codex-rs/core/src/unified_exec/process_tests.rs index db64d461ec66..e041ddcdfaa0 100644 --- a/codex-rs/core/src/unified_exec/process_tests.rs +++ b/codex-rs/core/src/unified_exec/process_tests.rs @@ -1,6 +1,7 @@ use super::process::UnifiedExecProcess; use crate::unified_exec::UnifiedExecError; use codex_exec_server::ExecProcess; +use codex_exec_server::ExecProcessEvent; use codex_exec_server::ExecProcessEventReceiver; use codex_exec_server::ExecProcessFuture; use codex_exec_server::ExecServerError; @@ -15,7 +16,6 @@ use std::collections::VecDeque; use std::sync::Arc; use tokio::sync::Mutex; use tokio::sync::watch; -use tokio::time::Duration; struct MockExecProcess { process_id: ProcessId, @@ -23,6 +23,7 @@ struct MockExecProcess { read_responses: Mutex>, terminate_error: Option, wake_tx: watch::Sender, + events: Vec, } impl MockExecProcess { @@ -61,7 +62,7 @@ impl ExecProcess for MockExecProcess { } fn subscribe_events(&self) -> ExecProcessEventReceiver { - ExecProcessEventReceiver::empty() + ExecProcessEventReceiver::from_events(self.events.clone()) } fn read( @@ -100,6 +101,7 @@ async fn remote_process( read_responses: Mutex::new(VecDeque::new()), terminate_error, wake_tx, + events: Vec::new(), }), }; @@ -183,25 +185,22 @@ async fn remote_process_waits_for_early_exit_event() { write_response: WriteResponse { status: WriteStatus::Accepted, }, - read_responses: Mutex::new(VecDeque::from([ReadResponse { - chunks: Vec::new(), - next_seq: 2, - exited: true, - exit_code: Some(17), - closed: true, - failure: None, - sandbox_denied: false, - }])), + read_responses: Mutex::new(VecDeque::new()), terminate_error: None, - wake_tx: wake_tx.clone(), + wake_tx, + events: vec![ + ExecProcessEvent::Exited { + seq: 1, + exit_code: 17, + }, + ExecProcessEvent::Closed { + seq: 2, + sandbox_denied: false, + }, + ], }), }; - tokio::spawn(async move { - tokio::time::sleep(Duration::from_millis(10)).await; - let _ = wake_tx.send(1); - }); - let process = UnifiedExecProcess::from_exec_server_started(started) .await .expect("remote process should observe early exit"); @@ -209,3 +208,35 @@ async fn remote_process_waits_for_early_exit_event() { assert!(process.has_exited()); assert_eq!(process.exit_code(), Some(17)); } + +#[tokio::test] +async fn remote_process_preserves_streamed_sandbox_denial() { + let (wake_tx, _wake_rx) = watch::channel(0); + let started = StartedExecProcess { + process: Arc::new(MockExecProcess { + process_id: "sandbox-denied".to_string().into(), + write_response: WriteResponse { + status: WriteStatus::Accepted, + }, + read_responses: Mutex::new(VecDeque::new()), + terminate_error: None, + wake_tx, + events: vec![ + ExecProcessEvent::Exited { + seq: 1, + exit_code: 1, + }, + ExecProcessEvent::Closed { + seq: 2, + sandbox_denied: true, + }, + ], + }), + }; + + let error = UnifiedExecProcess::from_exec_server_started(started) + .await + .expect_err("sandbox denial should be preserved"); + + assert!(matches!(error, UnifiedExecError::SandboxDenied { .. })); +} diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index 491fc8cf97ba..4a1994a3b515 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -534,6 +534,8 @@ pub struct ExecExitedNotification { pub struct ExecClosedNotification { pub process_id: ProcessId, pub seq: u64, + #[serde(default)] + pub sandbox_denied: bool, } mod base64_bytes { @@ -564,6 +566,7 @@ mod base64_bytes { #[cfg(test)] mod tests { use super::EnvironmentInfo; + use super::ExecClosedNotification; use super::ExecParams; use super::FsReadFileParams; use super::HttpRequestParams; @@ -708,4 +711,15 @@ mod tests { ("req-explicit-timeout", Some(1234)) ); } + + #[test] + fn closed_notification_accepts_legacy_payload_without_sandbox_denied() { + let notification: ExecClosedNotification = serde_json::from_value(serde_json::json!({ + "processId": "proc-1", + "seq": 3, + })) + .expect("legacy closed notification should deserialize"); + + assert!(!notification.sandbox_denied); + } } diff --git a/codex-rs/exec-server/README.md b/codex-rs/exec-server/README.md index b1f927918660..116b6441e922 100644 --- a/codex-rs/exec-server/README.md +++ b/codex-rs/exec-server/README.md @@ -334,10 +334,35 @@ Params: ```json { - "processId": "proc-1" + "processId": "proc-1", + "seq": 3, + "sandboxDenied": false } ``` +`sandboxDenied` lets streaming clients preserve executor-side sandbox denial +detection without issuing a final `process/read` request. Older servers omit +the field, which clients interpret as `false`. + +## Remote latency benchmark + +Run the repeatable benchmark against an already registered remote exec-server: + +```bash +CODEX_EXEC_SERVER_NOISE_REGISTRY_URL="$REGISTRY_URL" \ +CODEX_EXEC_SERVER_NOISE_ENVIRONMENT_ID="$ENVIRONMENT_ID" \ +CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN="$AUTH_TOKEN" \ +cargo run -p codex-exec-server --example remote_latency +``` + +It performs 5 warmups and 30 measured filesystem and process operations by +default, then prints connection timing, p50/p95 samples, per-RPC phase timing, +and bounded Rendezvous route dimensions. Set +`EXEC_SERVER_LATENCY_WARMUP_ITERATIONS` or `EXEC_SERVER_LATENCY_ITERATIONS` to +change sample counts. Process completion uses pushed events; set +`EXEC_SERVER_LATENCY_PROCESS_COMPLETION=read` to measure the legacy extra +`process/read` round trip as a control. + ## Filesystem RPCs Filesystem methods require valid `file:` URI strings and return JSON-RPC errors @@ -432,5 +457,5 @@ Terminate it: {"id":4,"method":"process/terminate","params":{"processId":"proc-1"}} {"id":4,"result":{"running":true}} {"method":"process/exited","params":{"processId":"proc-1","seq":3,"exitCode":0}} -{"method":"process/closed","params":{"processId":"proc-1"}} +{"method":"process/closed","params":{"processId":"proc-1","seq":4,"sandboxDenied":false}} ``` diff --git a/codex-rs/exec-server/examples/remote_latency.rs b/codex-rs/exec-server/examples/remote_latency.rs new file mode 100644 index 000000000000..cf2bf3129283 --- /dev/null +++ b/codex-rs/exec-server/examples/remote_latency.rs @@ -0,0 +1,328 @@ +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::time::Instant; + +use anyhow::Context; +use anyhow::Result; +use codex_exec_server::EnvironmentManager; +use codex_exec_server::ExecParams; +use codex_exec_server::ExecProcessEvent; +use codex_exec_server::ProcessId; +use opentelemetry::Value as OtelValue; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::InMemorySpanExporter; +use opentelemetry_sdk::trace::SdkTracerProvider; +use opentelemetry_sdk::trace::SpanData; +use serde_json::Value; +use serde_json::json; +use tracing::Instrument; +use tracing_subscriber::EnvFilter; +use tracing_subscriber::prelude::*; + +const ITERATIONS_ENV: &str = "EXEC_SERVER_LATENCY_ITERATIONS"; +const WARMUP_ITERATIONS_ENV: &str = "EXEC_SERVER_LATENCY_WARMUP_ITERATIONS"; +const PROCESS_COMPLETION_MODE_ENV: &str = "EXEC_SERVER_LATENCY_PROCESS_COMPLETION"; +const RPC_CLIENT_SPAN_NAME: &str = "codex.exec_server.rpc.client_call"; +const RPC_DURATION_FIELDS: [&str; 6] = [ + "duration_ms", + "pending_registration_ms", + "serialize_ms", + "enqueue_ms", + "response_wait_ms", + "deserialize_ms", +]; + +#[derive(Clone, Copy)] +enum ProcessCompletionMode { + Events, + Read, +} + +impl ProcessCompletionMode { + fn from_env() -> Result { + match std::env::var(PROCESS_COMPLETION_MODE_ENV) + .unwrap_or_else(|_| "events".to_string()) + .as_str() + { + "events" => Ok(Self::Events), + "read" => Ok(Self::Read), + value => anyhow::bail!( + "{PROCESS_COMPLETION_MODE_ENV} must be `events` or `read`, got {value:?}" + ), + } + } + + fn as_str(self) -> &'static str { + match self { + Self::Events => "events", + Self::Read => "read", + } + } +} + +#[tokio::main] +async fn main() -> Result<()> { + let span_exporter = InMemorySpanExporter::default(); + let tracer_provider = SdkTracerProvider::builder() + .with_simple_exporter(span_exporter.clone()) + .build(); + let tracer = tracer_provider.tracer("exec-server-remote-latency"); + let fmt_layer = tracing_subscriber::fmt::layer() + .with_writer(std::io::stderr) + .with_filter( + EnvFilter::try_from_default_env() + .unwrap_or_else(|_| EnvFilter::new("info,hyper=warn,reqwest=warn")), + ); + tracing_subscriber::registry() + .with(fmt_layer) + .with(tracing_opentelemetry::layer().with_tracer(tracer)) + .try_init() + .context("initialize benchmark tracing")?; + tracing::callsite::rebuild_interest_cache(); + + let iterations = positive_env_usize(ITERATIONS_ENV, 30)?; + let warmup_iterations = positive_env_usize(WARMUP_ITERATIONS_ENV, 5)?; + let process_completion_mode = ProcessCompletionMode::from_env()?; + let benchmark_span = tracing::info_span!( + "codex.exec_server.remote.benchmark", + otel.kind = "client", + otel.name = "codex.exec_server.remote.benchmark", + iterations, + warmup_iterations, + process_completion_mode = process_completion_mode.as_str(), + ); + + let mut report = async move { + let trace = codex_otel::current_span_w3c_trace_context(); + let connection_started_at = Instant::now(); + let manager = EnvironmentManager::from_env(/*local_runtime_paths*/ None).await?; + let environment = manager + .default_environment() + .context("Noise remote environment is not configured")?; + environment.wait_until_ready().await?; + let connection_ms = elapsed_ms(connection_started_at); + + let environment_info = environment.info().await?; + let cwd = environment_info + .cwd + .context("remote environment did not report a working directory")?; + let filesystem = environment.get_filesystem(); + let exec_backend = environment.get_exec_backend(); + + for iteration in 0..warmup_iterations { + filesystem.get_metadata(&cwd, /*sandbox*/ None).await?; + let started = exec_backend + .start(ExecParams { + process_id: ProcessId::from(format!("latency-warmup-{iteration}")), + argv: vec!["/usr/bin/true".to_string()], + cwd: cwd.clone(), + env_policy: None, + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + }) + .await?; + wait_for_process_completion(started.process, process_completion_mode).await?; + } + + let mut metadata_ms = Vec::with_capacity(iterations); + let mut process_start_ms = Vec::with_capacity(iterations); + let mut process_completion_wait_ms = Vec::with_capacity(iterations); + let mut process_completion_ms = Vec::with_capacity(iterations); + for iteration in 0..iterations { + let sample_span = tracing::info_span!( + "codex.exec_server.remote.benchmark.sample", + otel.kind = "client", + otel.name = "codex.exec_server.remote.benchmark.sample", + iteration, + ); + async { + let started_at = Instant::now(); + filesystem.get_metadata(&cwd, /*sandbox*/ None).await?; + metadata_ms.push(elapsed_ms(started_at)); + + let completion_started_at = Instant::now(); + let started_at = Instant::now(); + let started = exec_backend + .start(ExecParams { + process_id: ProcessId::from(format!("latency-measured-{iteration}")), + argv: vec!["/usr/bin/true".to_string()], + cwd: cwd.clone(), + env_policy: None, + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + }) + .await?; + process_start_ms.push(elapsed_ms(started_at)); + + let started_at = Instant::now(); + wait_for_process_completion(started.process, process_completion_mode).await?; + process_completion_wait_ms.push(elapsed_ms(started_at)); + process_completion_ms.push(elapsed_ms(completion_started_at)); + Result::<()>::Ok(()) + } + .instrument(sample_span) + .await?; + } + + Result::::Ok(json!({ + "environment_id": std::env::var("CODEX_EXEC_SERVER_NOISE_ENVIRONMENT_ID").ok(), + "traceparent": trace.and_then(|context| context.traceparent), + "iterations": iterations, + "warmup_iterations": warmup_iterations, + "process_completion_mode": process_completion_mode.as_str(), + "connection_ms": connection_ms, + "fs_get_metadata_ms": summarize(&metadata_ms), + "process_start_ms": summarize(&process_start_ms), + "process_completion_wait_ms": summarize(&process_completion_wait_ms), + "process_completion_ms": summarize(&process_completion_ms), + })) + } + .instrument(benchmark_span) + .await?; + + tracer_provider + .force_flush() + .context("flush benchmark spans")?; + let spans = span_exporter + .get_finished_spans() + .context("read benchmark spans")?; + report["rpc_client_ms"] = rpc_phase_summary(&spans); + report["rendezvous_cluster"] = span_string_attribute( + &spans, + "codex.exec_server.remote.harness.registry_connect_bundle", + "rendezvous.cluster", + ); + report["rendezvous_route_id"] = span_string_attribute( + &spans, + "codex.exec_server.remote.harness.registry_connect_bundle", + "rendezvous.route_id", + ); + println!("{}", serde_json::to_string_pretty(&report)?); + Ok(()) +} + +async fn wait_for_process_completion( + process: std::sync::Arc, + mode: ProcessCompletionMode, +) -> Result<()> { + match mode { + ProcessCompletionMode::Events => { + let mut events = process.subscribe_events(); + loop { + match events.recv().await { + Ok(ExecProcessEvent::Closed { .. }) => return Ok(()), + Ok(ExecProcessEvent::Failed(message)) => anyhow::bail!(message), + Ok(ExecProcessEvent::Output(_) | ExecProcessEvent::Exited { .. }) => {} + Err(error) => anyhow::bail!("process event stream failed: {error}"), + } + } + } + ProcessCompletionMode::Read => { + let mut after_seq = None; + loop { + let response = process + .read(after_seq, /*max_bytes*/ None, Some(1_000)) + .await?; + if let Some(message) = response.failure { + anyhow::bail!(message); + } + if response.closed { + return Ok(()); + } + after_seq = response.next_seq.checked_sub(1); + } + } + } +} + +fn positive_env_usize(name: &str, default: usize) -> Result { + let Some(value) = std::env::var(name).ok() else { + return Ok(default); + }; + let value = value + .parse::() + .with_context(|| format!("{name} must be a positive integer"))?; + anyhow::ensure!(value > 0, "{name} must be a positive integer"); + Ok(value) +} + +fn elapsed_ms(started_at: Instant) -> f64 { + started_at.elapsed().as_secs_f64() * 1_000.0 +} + +fn summarize(samples: &[f64]) -> Value { + let mut sorted = samples.to_vec(); + sorted.sort_by(f64::total_cmp); + let percentile = |percentile: usize| { + let index = (sorted.len() * percentile).div_ceil(100).saturating_sub(1); + sorted[index] + }; + json!({ + "min": sorted[0], + "p50": percentile(50), + "p95": percentile(95), + "max": sorted[sorted.len() - 1], + "mean": sorted.iter().sum::() / sorted.len() as f64, + "samples": samples, + }) +} + +fn rpc_phase_summary(spans: &[SpanData]) -> Value { + let mut samples_by_method = BTreeMap::>>::new(); + for span in spans + .iter() + .filter(|span| span.name.as_ref() == RPC_CLIENT_SPAN_NAME) + { + let Some(OtelValue::String(method)) = span_attribute(span, "rpc.method") else { + continue; + }; + let fields = samples_by_method.entry(method.to_string()).or_default(); + for field in RPC_DURATION_FIELDS { + if let Some(OtelValue::F64(duration_ms)) = span_attribute(span, field) { + fields.entry(field).or_default().push(*duration_ms); + } + } + } + + Value::Object( + samples_by_method + .into_iter() + .map(|(method, fields)| { + let fields = fields + .into_iter() + .map(|(field, samples)| (field.to_string(), summarize(&samples))) + .collect(); + (method, Value::Object(fields)) + }) + .collect(), + ) +} + +fn span_attribute<'a>(span: &'a SpanData, name: &str) -> Option<&'a OtelValue> { + span.attributes + .iter() + .find(|attribute| attribute.key.as_str() == name) + .map(|attribute| &attribute.value) +} + +fn span_string_attribute(spans: &[SpanData], span_name: &str, attribute_name: &str) -> Value { + spans + .iter() + .find(|span| span.name.as_ref() == span_name) + .and_then(|span| span_attribute(span, attribute_name)) + .and_then(|value| match value { + OtelValue::String(value) => Some(Value::String(value.to_string())), + _ => None, + }) + .unwrap_or(Value::Null) +} diff --git a/codex-rs/exec-server/src/client.rs b/codex-rs/exec-server/src/client.rs index e2429d8f1440..2a1df3eb989e 100644 --- a/codex-rs/exec-server/src/client.rs +++ b/codex-rs/exec-server/src/client.rs @@ -7,6 +7,7 @@ use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU64; use std::sync::atomic::Ordering; use std::time::Duration; +use std::time::Instant; use arc_swap::ArcSwap; use codex_exec_server_protocol::JSONRPCNotification; @@ -103,6 +104,9 @@ use crate::protocol::WriteParams; use crate::protocol::WriteResponse; use crate::rpc::RpcCallError; use crate::rpc::RpcClient; +use crate::telemetry::REMOTE_PHASE_HARNESS_INITIALIZE; +use crate::telemetry::RemotePhaseGuard; +use crate::telemetry::record_remote_phase; pub(crate) mod http_client; #[path = "client_recovery.rs"] @@ -505,6 +509,58 @@ impl ExecServerClient { })? } + #[tracing::instrument( + name = "codex.exec_server.remote.harness.initialize", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.initialize", + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + ) + )] + async fn initialize_noise_rpc( + &self, + rpc_client: &RpcClient, + options: ExecServerClientConnectOptions, + ) -> Result { + let started_at = Instant::now(); + let result = self.initialize_rpc(rpc_client, options).await; + let result_name = match &result { + Ok(_) => "success", + Err(ExecServerError::InitializeTimedOut { .. }) => "timeout", + Err(_) => "error", + }; + record_remote_phase( + &tracing::Span::current(), + REMOTE_PHASE_HARNESS_INITIALIZE, + result_name, + started_at.elapsed(), + ); + result + } + + async fn initialize_with_ready_phase( + &self, + rpc_client: &RpcClient, + options: ExecServerClientConnectOptions, + ready_phase: Option, + ) -> Result { + let Some(ready_phase) = ready_phase else { + return self.initialize_rpc(rpc_client, options).await; + }; + let result = self + .initialize_noise_rpc(rpc_client, options) + .instrument(ready_phase.span().clone()) + .await; + ready_phase.finish(match &result { + Ok(_) => "success", + Err(ExecServerError::InitializeTimedOut { .. }) => "timeout", + Err(_) => "error", + }); + result + } + pub async fn exec(&self, params: ExecParams) -> Result { self.call(EXEC_METHOD, ¶ms).await } @@ -725,13 +781,17 @@ impl ExecServerClient { connection: JsonRpcConnection, options: ExecServerClientConnectOptions, ) -> Result { - Self::connect_with_recovery(connection, options, /*reconnect_strategy*/ None).await + Self::connect_with_recovery( + connection, options, /*reconnect_strategy*/ None, /*ready_phase*/ None, + ) + .await } pub(crate) async fn connect_with_recovery( connection: JsonRpcConnection, options: ExecServerClientConnectOptions, reconnect_strategy: Option, + ready_phase: Option, ) -> Result { let (rpc_client, events_rx) = RpcClient::new(connection); let rpc_client = Arc::new(rpc_client); @@ -757,7 +817,9 @@ impl ExecServerClient { // before initialize returns. Drain them immediately so a burst cannot // fill the bounded event channel and block the initialize response. client.spawn_rpc_reader(&rpc_client, events_rx); - client.initialize_rpc(&rpc_client, options).await?; + client + .initialize_with_ready_phase(&rpc_client, options, ready_phase) + .await?; Ok(client) } @@ -1175,8 +1237,10 @@ async fn handle_server_notification( // Closed is terminal, but it can arrive before tail output or // exited. Keep routing this process until the ordered publisher // says Closed has actually been delivered. - let published_closed = - session.publish_ordered_event(ExecProcessEvent::Closed { seq: params.seq }); + let published_closed = session.publish_ordered_event(ExecProcessEvent::Closed { + seq: params.seq, + sandbox_denied: params.sandbox_denied, + }); if published_closed { inner.remove_session_if(¶ms.process_id, &session); } @@ -1385,7 +1449,19 @@ mod tests { assert_eq!(session.process_id(), &process_id); let trace = server.await.expect("server task").expect("trace context"); - assert_eq!(trace, expected_trace); + let expected_traceparent = expected_trace + .traceparent + .as_deref() + .expect("parent traceparent"); + let actual_traceparent = trace.traceparent.as_deref().expect("request traceparent"); + assert_eq!( + actual_traceparent.split('-').nth(1), + expected_traceparent.split('-').nth(1), + ); + assert_ne!( + actual_traceparent.split('-').nth(2), + expected_traceparent.split('-').nth(2), + ); } async fn accept_websocket(listener: &TcpListener) -> WebSocketStream { @@ -1714,6 +1790,7 @@ mod tests { serde_json::to_value(ExecClosedNotification { process_id: process_id.clone(), seq: 4, + sandbox_denied: true, }) .expect("closed notification should serialize"), ), @@ -1787,7 +1864,10 @@ mod tests { seq: 3, exit_code: 0, }, - ExecProcessEvent::Closed { seq: 4 }, + ExecProcessEvent::Closed { + seq: 4, + sandbox_denied: true, + }, ] ); diff --git a/codex-rs/exec-server/src/client_recovery.rs b/codex-rs/exec-server/src/client_recovery.rs index 61cee009cdd5..b13485b2bd25 100644 --- a/codex-rs/exec-server/src/client_recovery.rs +++ b/codex-rs/exec-server/src/client_recovery.rs @@ -58,7 +58,7 @@ impl SessionState { exit_code, closed, failure, - sandbox_denied: _, + sandbox_denied, } = response; if let Some(message) = failure { return Err(ExecServerError::Protocol(format!( @@ -87,17 +87,26 @@ impl SessionState { } } if closed { - match ordered_events.pending.get(&target_seq) { - Some(ExecProcessEvent::Closed { .. }) => {} + match ordered_events.pending.get_mut(&target_seq) { + Some(ExecProcessEvent::Closed { + sandbox_denied: pending_sandbox_denied, + .. + }) => { + *pending_sandbox_denied |= sandbox_denied; + } Some(_) => { return Err(ExecServerError::Protocol(format!( "process close sequence {target_seq} conflicts with recovered output" ))); } None => { - ordered_events - .pending - .insert(target_seq, ExecProcessEvent::Closed { seq: target_seq }); + ordered_events.pending.insert( + target_seq, + ExecProcessEvent::Closed { + seq: target_seq, + sandbox_denied, + }, + ); } } } @@ -382,7 +391,7 @@ impl Inner { .reconnect_strategy .as_ref() .ok_or_else(|| ExecServerError::Protocol("missing reconnect strategy".to_string()))?; - let (connection, options) = reconnect_strategy.resume(session_id).await?; + let (connection, options, ready_phase) = reconnect_strategy.resume(session_id).await?; let (rpc_client, events_rx) = RpcClient::new(connection); let rpc_client = Arc::new(rpc_client); let client = ExecServerClient { @@ -393,7 +402,9 @@ impl Inner { // burst cannot fill the bounded event channel and block the initialize // response behind it. client.spawn_rpc_reader(&rpc_client, events_rx); - client.initialize_rpc(&rpc_client, options).await?; + client + .initialize_with_ready_phase(&rpc_client, options, ready_phase) + .await?; self.recover_processes(&rpc_client).await?; Ok(rpc_client) diff --git a/codex-rs/exec-server/src/client_transport.rs b/codex-rs/exec-server/src/client_transport.rs index 752bf270378c..9a8f5034bd98 100644 --- a/codex-rs/exec-server/src/client_transport.rs +++ b/codex-rs/exec-server/src/client_transport.rs @@ -1,6 +1,7 @@ use std::process::Stdio; use std::sync::Arc; use std::time::Duration; +use std::time::Instant; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; @@ -31,6 +32,11 @@ use crate::noise_relay::NoiseHarnessConnectionArgs; use crate::noise_relay::noise_harness_connection_from_websocket; use crate::noise_relay::noise_relay_websocket_config; use crate::relay::harness_connection_from_websocket; +use crate::telemetry::REMOTE_PHASE_HARNESS_READY; +use crate::telemetry::REMOTE_PHASE_HARNESS_WEBSOCKET; +use crate::telemetry::REMOTE_PHASE_REGISTRY_CONNECT_BUNDLE; +use crate::telemetry::RemotePhaseGuard; +use crate::telemetry::record_remote_phase; use crate::trace_context::current_trace_context_headers; const ENVIRONMENT_CLIENT_NAME: &str = "codex-environment"; @@ -56,13 +62,20 @@ impl ExecServerReconnectStrategy { pub(crate) async fn resume( &self, session_id: &str, - ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> { + ) -> Result< + ( + JsonRpcConnection, + ExecServerClientConnectOptions, + Option, + ), + ExecServerError, + > { match self { Self::WebSocket(args) => { let mut args = args.clone(); args.resume_session_id = Some(session_id.to_string()); let connection = ExecServerClient::open_websocket_connection(&args).await?; - Ok((connection, args.into())) + Ok((connection, args.into(), None)) } Self::NoiseRendezvous { provider, @@ -71,16 +84,20 @@ impl ExecServerReconnectStrategy { connect_timeout, initialize_timeout, } => { - let bundle = provider.connect_bundle(identity.public_key()).await?; - ExecServerClient::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { - bundle, - harness_identity: identity.clone(), - client_name: client_name.clone(), - connect_timeout: *connect_timeout, - initialize_timeout: *initialize_timeout, - resume_session_id: Some(session_id.to_string()), - }) - .await + let bundle = request_noise_connect_bundle(provider, identity.public_key()).await?; + let (connection, options, ready_phase) = + ExecServerClient::open_noise_rendezvous_connection( + NoiseRendezvousConnectArgs { + bundle, + harness_identity: identity.clone(), + client_name: client_name.clone(), + connect_timeout: *connect_timeout, + initialize_timeout: *initialize_timeout, + resume_session_id: Some(session_id.to_string()), + }, + ) + .await?; + Ok((connection, options, Some(ready_phase))) } } } @@ -119,9 +136,15 @@ impl ExecServerClient { connect_timeout: DEFAULT_REMOTE_EXEC_SERVER_CONNECT_TIMEOUT, initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, }; - let (connection, options) = + let (connection, options, ready_phase) = Self::open_initial_noise_rendezvous_connection(&provider, &identity).await?; - Self::connect_with_recovery(connection, options, Some(reconnect_strategy)).await + Self::connect_with_recovery( + connection, + options, + Some(reconnect_strategy), + Some(ready_phase), + ) + .await } crate::client_api::ExecServerTransportParams::StdioCommand { command, @@ -141,7 +164,14 @@ impl ExecServerClient { async fn open_initial_noise_rendezvous_connection( provider: &Arc, identity: &NoiseChannelIdentity, - ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> { + ) -> Result< + ( + JsonRpcConnection, + ExecServerClientConnectOptions, + RemotePhaseGuard, + ), + ExecServerError, + > { let open_connection = |bundle: NoiseRendezvousConnectBundle| { Self::open_noise_rendezvous_connection(NoiseRendezvousConnectArgs { bundle, @@ -152,7 +182,7 @@ impl ExecServerClient { resume_session_id: None, }) }; - let bundle = provider.connect_bundle(identity.public_key()).await?; + let bundle = request_noise_connect_bundle(provider, identity.public_key()).await?; match open_connection(bundle).await { Err(error) if matches!( @@ -165,7 +195,7 @@ impl ExecServerClient { ) ) => { - let bundle = provider.connect_bundle(identity.public_key()).await?; + let bundle = request_noise_connect_bundle(provider, identity.public_key()).await?; open_connection(bundle).await } result => result, @@ -181,6 +211,7 @@ impl ExecServerClient { connection, options, Some(ExecServerReconnectStrategy::WebSocket(args)), + /*ready_phase*/ None, ) .await } @@ -229,13 +260,37 @@ impl ExecServerClient { pub async fn connect_noise_rendezvous( args: NoiseRendezvousConnectArgs, ) -> Result { - let (connection, options) = Self::open_noise_rendezvous_connection(args).await?; - Self::connect(connection, options).await + let (connection, options, ready_phase) = + Self::open_noise_rendezvous_connection(args).await?; + Self::connect_with_recovery( + connection, + options, + /*reconnect_strategy*/ None, + Some(ready_phase), + ) + .await } + #[tracing::instrument( + name = "codex.exec_server.remote.harness.websocket", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.websocket", + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + ) + )] pub(crate) async fn open_noise_rendezvous_connection( args: NoiseRendezvousConnectArgs, - ) -> Result<(JsonRpcConnection, ExecServerClientConnectOptions), ExecServerError> { + ) -> Result< + ( + JsonRpcConnection, + ExecServerClientConnectOptions, + RemotePhaseGuard, + ), + ExecServerError, + > { ensure_rustls_crypto_provider(); // Keep the registry-issued URL, key, and authorization together for this // connection attempt. @@ -259,34 +314,59 @@ impl ExecServerClient { .next() .unwrap_or(websocket_url.as_str()) .to_string(); - let mut request = websocket_url - .as_str() - .into_client_request() + let started_at = Instant::now(); + let websocket = async { + let mut request = websocket_url + .as_str() + .into_client_request() + .map_err(|source| ExecServerError::WebSocketConnect { + url: diagnostic_url.clone(), + source, + })?; + request + .headers_mut() + .extend(current_trace_context_headers()); + timeout( + connect_timeout, + connect_async_with_config( + request, + Some(noise_relay_websocket_config()), + /*disable_nagle*/ true, + ), + ) + .await + .map_err(|_| ExecServerError::WebSocketConnectTimeout { + url: diagnostic_url.clone(), + timeout: connect_timeout, + })? .map_err(|source| ExecServerError::WebSocketConnect { url: diagnostic_url.clone(), source, - })?; - request - .headers_mut() - .extend(current_trace_context_headers()); - let (stream, _) = timeout( - connect_timeout, - connect_async_with_config( - request, - Some(noise_relay_websocket_config()), - /*disable_nagle*/ false, - ), - ) - .await - .map_err(|_| ExecServerError::WebSocketConnectTimeout { - url: diagnostic_url.clone(), - timeout: connect_timeout, - })? - .map_err(|source| ExecServerError::WebSocketConnect { - url: diagnostic_url.clone(), - source, - })?; + }) + } + .await; + let duration = started_at.elapsed(); + record_remote_phase( + &tracing::Span::current(), + REMOTE_PHASE_HARNESS_WEBSOCKET, + remote_phase_result(&websocket), + duration, + ); + let (stream, _) = websocket?; + // This is the inclusive post-WebSocket critical-path interval: add it to + // bundle + WebSocket, but do not also add its Noise-handshake child. + // The ready and handshake completion events retain the timestamps needed + // to derive post-handshake initialize as ready_end - handshake_end. + let ready_span = tracing::info_span!( + "codex.exec_server.remote.harness.ready", + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.ready", + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + ); + let ready_phase = + RemotePhaseGuard::new(ready_span.clone(), REMOTE_PHASE_HARNESS_READY, "cancelled"); let connection_label = format!("Noise exec-server rendezvous websocket {diagnostic_url}"); let connection = noise_harness_connection_from_websocket( stream, @@ -297,6 +377,7 @@ impl ExecServerClient { identity: harness_identity, responder_public_key: executor_public_key, harness_key_authorization, + ready_span, }, ); Ok(( @@ -306,6 +387,7 @@ impl ExecServerClient { initialize_timeout, resume_session_id, }, + ready_phase, )) } @@ -350,6 +432,74 @@ impl ExecServerClient { } } +#[tracing::instrument( + name = "codex.exec_server.remote.harness.registry_connect_bundle", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.registry_connect_bundle", + rendezvous.cluster = tracing::field::Empty, + rendezvous.route_id = tracing::field::Empty, + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + ) +)] +async fn request_noise_connect_bundle( + provider: &Arc, + harness_public_key: crate::NoiseChannelPublicKey, +) -> Result { + let started_at = Instant::now(); + let result = provider.connect_bundle(harness_public_key).await; + if let Ok(bundle) = &result { + let span = tracing::Span::current(); + if let Some(cluster) = rendezvous_cluster(&bundle.websocket_url) { + span.record("rendezvous.cluster", cluster); + if let Some(route_id) = rendezvous_route_id(&bundle.websocket_url) { + span.record("rendezvous.route_id", route_id); + } + } + } + record_remote_phase( + &tracing::Span::current(), + REMOTE_PHASE_REGISTRY_CONNECT_BUNDLE, + remote_phase_result(&result), + started_at.elapsed(), + ); + result +} + +fn rendezvous_cluster(websocket_url: &str) -> Option<&str> { + let authority = websocket_url + .split_once("://")? + .1 + .split(['/', '?']) + .next()?; + authority + .strip_prefix("codex-cloud-rendezvous.gateway.")? + .strip_suffix(".internal.api.openai.org") + .filter(|cluster| !cluster.is_empty()) +} + +fn rendezvous_route_id(websocket_url: &str) -> Option<&str> { + let path = websocket_url.split_once("://")?.1.split('?').next()?; + let mut segments = path.split('/'); + while let Some(segment) = segments.next() { + if segment == "cloud-agent" { + return segments.next().filter(|route_id| !route_id.is_empty()); + } + } + None +} + +fn remote_phase_result(result: &Result) -> &'static str { + match result { + Ok(_) => "success", + Err(ExecServerError::EnvironmentRegistryRequest(error)) if error.is_timeout() => "timeout", + Err(ExecServerError::WebSocketConnectTimeout { .. }) => "timeout", + Err(_) => "error", + } +} + fn is_rendezvous_harness_url(websocket_url: &str) -> bool { let Some((_path, query)) = websocket_url.split_once('?') else { return false; diff --git a/codex-rs/exec-server/src/client_transport_tests.rs b/codex-rs/exec-server/src/client_transport_tests.rs index 7dab5fa21835..f7fbe0023f5e 100644 --- a/codex-rs/exec-server/src/client_transport_tests.rs +++ b/codex-rs/exec-server/src/client_transport_tests.rs @@ -11,6 +11,8 @@ use tokio::net::TcpListener; use tokio_tungstenite::accept_async; use super::ExecServerClient; +use super::rendezvous_cluster; +use super::rendezvous_route_id; use crate::ExecServerError; use crate::NoiseChannelIdentity; use crate::NoiseChannelPublicKey; @@ -22,6 +24,17 @@ struct SequenceNoiseConnectProvider { returned_urls: Mutex>, } +#[test] +fn rendezvous_span_dimensions_exclude_signed_url_data() { + let url = "wss://codex-cloud-rendezvous.gateway.unified-2s.internal.api.openai.org/cloud-agent/c2-u2s/ws/environment/env-1?role=harness&sig=secret"; + + assert_eq!(rendezvous_cluster(url), Some("unified-2s")); + assert_eq!(rendezvous_route_id(url), Some("c2-u2s")); + + let custom_url = "wss://customer.example/cloud-agent/unbounded/ws/environment/env-1?sig=secret"; + assert_eq!(rendezvous_cluster(custom_url), None); +} + impl SequenceNoiseConnectProvider { fn new(bundles: Vec) -> Self { Self { diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 7335bd1eeda1..958614618bc7 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -930,11 +930,15 @@ async fn maybe_emit_closed(process_id: ProcessId, inner: Arc) { let seq = process.next_seq; process.next_seq += 1; let _ = process.wake_tx.send(seq); - process.events.publish(ExecProcessEvent::Closed { seq }); + process.events.publish(ExecProcessEvent::Closed { + seq, + sandbox_denied: process.sandbox_denied, + }); ( ExecClosedNotification { process_id: process_id.clone(), seq, + sandbox_denied: process.sandbox_denied, }, Arc::clone(&process.output_notify), ) diff --git a/codex-rs/exec-server/src/noise_relay/harness.rs b/codex-rs/exec-server/src/noise_relay/harness.rs index 1989525224f6..2098dc5ee82a 100644 --- a/codex-rs/exec-server/src/noise_relay/harness.rs +++ b/codex-rs/exec-server/src/noise_relay/harness.rs @@ -39,6 +39,8 @@ use crate::relay::decode_relay_message_frame; use crate::relay::encode_relay_message_frame; use crate::relay_proto::RelayData; use crate::relay_proto::RelayMessageFrame; +use crate::telemetry::REMOTE_PHASE_HARNESS_NOISE_HANDSHAKE; +use crate::telemetry::RemotePhaseGuard; /// Values that bind one harness websocket to the intended executor registration. /// @@ -52,6 +54,7 @@ pub(crate) struct NoiseHarnessConnectionArgs { pub(crate) identity: NoiseChannelIdentity, pub(crate) responder_public_key: NoiseChannelPublicKey, pub(crate) harness_key_authorization: String, + pub(crate) ready_span: tracing::Span, } // Reset frames are cleartext relay control and are not authenticated by Noise. @@ -81,12 +84,22 @@ where identity, responder_public_key, harness_key_authorization, + ready_span, } = args; let stream_id = Uuid::new_v4().to_string(); let (outgoing_tx, mut outgoing_rx) = mpsc::channel(CHANNEL_CAPACITY); let (incoming_tx, incoming_rx) = mpsc::channel(CHANNEL_CAPACITY); let (disconnected_tx, disconnected_rx) = watch::channel(false); - let stream_span = tracing::debug_span!("noise_relay.stream", noise_side = "harness",); + let stream_span = + tracing::debug_span!(parent: &ready_span, "noise_relay.stream", noise_side = "harness",); + let handshake_span = tracing::info_span!( + parent: &ready_span, + "codex.exec_server.remote.harness.noise_handshake", + otel.kind = "client", + otel.name = "codex.exec_server.remote.harness.noise_handshake", + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + ); debug!( environment_id, executor_registration_id, stream_id, "Noise harness relay details" @@ -94,6 +107,11 @@ where let websocket_task = tokio::spawn(async move { let mut websocket = stream; + let handshake_phase = RemotePhaseGuard::new( + handshake_span, + REMOTE_PHASE_HARNESS_NOISE_HANDSHAKE, + "cancelled", + ); // Bind the Noise transcript to the exact environment registration and // virtual relay stream before emitting any handshake bytes. A captured @@ -114,6 +132,7 @@ where format!("failed to start Noise relay handshake: {error}"), ) .await; + handshake_phase.finish("error"); return; } }; @@ -135,6 +154,7 @@ where .is_err() { let _ = disconnected_tx.send(true); + handshake_phase.finish("error"); return; } @@ -149,6 +169,7 @@ where "Noise relay websocket ended during handshake".to_string(), ) .await; + handshake_phase.finish("error"); return; }; let message = match incoming_message { @@ -160,6 +181,7 @@ where "Noise relay websocket received close frame during handshake".to_string(), ) .await; + handshake_phase.finish("error"); return; } Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_)) => continue, @@ -170,6 +192,7 @@ where "Noise relay transport expects binary protobuf frames".to_string(), ) .await; + handshake_phase.finish("error"); return; } Err(error) => { @@ -181,6 +204,7 @@ where ), ) .await; + handshake_phase.finish("error"); return; } }; @@ -193,6 +217,7 @@ where format!("failed to parse Noise relay frame: {error}"), ) .await; + handshake_phase.finish("error"); return; } }; @@ -211,18 +236,12 @@ where format!("invalid Noise relay handshake response: {error}"), ) .await; + handshake_phase.finish("error"); return; } }; match initiator_handshake.finish(&response) { - Ok(transport) => { - info!( - noise_event = "handshake", - noise_outcome = "ok", - "Noise harness handshake completed" - ); - break transport; - } + Ok(transport) => break transport, Err(error) => { send_disconnected( &incoming_tx, @@ -230,6 +249,7 @@ where format!("Noise relay handshake failed: {error}"), ) .await; + handshake_phase.finish("error"); return; } } @@ -241,6 +261,7 @@ where NOISE_RELAY_RESET_DISCONNECT_REASON.to_string(), ) .await; + handshake_phase.finish("error"); return; } Ok( @@ -255,10 +276,17 @@ where "Noise relay received data before handshake completion".to_string(), ) .await; + handshake_phase.finish("error"); return; } } }; + handshake_phase.finish("success"); + info!( + noise_event = "handshake", + noise_outcome = "ok", + "Noise harness handshake completed" + ); // After the handshake, each relay sequence maps to exactly one Noise // transport record. Outbound records are encrypted once; inbound diff --git a/codex-rs/exec-server/src/process.rs b/codex-rs/exec-server/src/process.rs index b136ef4682fa..3dd93aaf1dae 100644 --- a/codex-rs/exec-server/src/process.rs +++ b/codex-rs/exec-server/src/process.rs @@ -31,7 +31,7 @@ pub struct StartedExecProcess { pub enum ExecProcessEvent { Output(ProcessOutputChunk), Exited { seq: u64, exit_code: i32 }, - Closed { seq: u64 }, + Closed { seq: u64, sandbox_denied: bool }, Failed(String), } @@ -67,7 +67,9 @@ impl ExecProcessEvent { pub(crate) fn seq(&self) -> Option { match self { ExecProcessEvent::Output(chunk) => Some(chunk.seq), - ExecProcessEvent::Exited { seq, .. } | ExecProcessEvent::Closed { seq } => Some(*seq), + ExecProcessEvent::Exited { seq, .. } | ExecProcessEvent::Closed { seq, .. } => { + Some(*seq) + } ExecProcessEvent::Failed(_) => None, } } @@ -125,13 +127,18 @@ impl ExecProcessEventLog { let live_rx = self.inner.live_tx.subscribe(); let replay = history.events.iter().cloned().collect(); - ExecProcessEventReceiver { replay, live_rx } + ExecProcessEventReceiver { + replay, + live_rx, + _keepalive: None, + } } } pub struct ExecProcessEventReceiver { replay: VecDeque, live_rx: broadcast::Receiver, + _keepalive: Option>, } impl ExecProcessEventReceiver { @@ -140,6 +147,17 @@ impl ExecProcessEventReceiver { Self { replay: VecDeque::new(), live_rx, + _keepalive: None, + } + } + + #[doc(hidden)] + pub fn from_events(events: Vec) -> Self { + let (live_tx, live_rx) = broadcast::channel(1); + Self { + replay: events.into(), + live_rx, + _keepalive: Some(live_tx), } } @@ -218,7 +236,10 @@ mod tests { seq: 2, exit_code: 0, }); - log.publish(ExecProcessEvent::Closed { seq: 3 }); + log.publish(ExecProcessEvent::Closed { + seq: 3, + sandbox_denied: false, + }); let mut events = log.subscribe(); let replay = vec![ @@ -239,7 +260,10 @@ mod tests { seq: 2, exit_code: 0, }, - ExecProcessEvent::Closed { seq: 3 }, + ExecProcessEvent::Closed { + seq: 3, + sandbox_denied: false, + }, ] ); } diff --git a/codex-rs/exec-server/src/relay.rs b/codex-rs/exec-server/src/relay.rs index d7e3d9c2a7fc..aac6a7195808 100644 --- a/codex-rs/exec-server/src/relay.rs +++ b/codex-rs/exec-server/src/relay.rs @@ -1,5 +1,6 @@ use std::collections::HashMap; use std::time::Duration; +use std::time::Instant; use codex_exec_server_protocol::JSONRPCMessage; use futures::Sink; @@ -15,6 +16,7 @@ use tokio::task::JoinSet; use tokio::time::timeout; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::tungstenite::Message; +use tracing::Instrument; use tracing::debug; use tracing::info; use tracing::warn; @@ -41,6 +43,8 @@ use crate::relay_proto::RelayReset; use crate::relay_proto::RelayResume; use crate::relay_proto::relay_message_frame; use crate::server::ConnectionProcessor; +use crate::telemetry::REMOTE_PHASE_EXECUTOR_AUTHORIZATION; +use crate::telemetry::record_remote_phase; const RELAY_MESSAGE_FRAME_VERSION: u32 = 1; const MAX_ACTIVE_NOISE_RELAY_STREAMS: usize = 128; @@ -434,6 +438,41 @@ pub(crate) trait HarnessKeyValidator: Send + Sync { harness_public_key: &NoiseChannelPublicKey, authorization: &str, ) -> impl std::future::Future> + Send; + + fn validation_completed(&self, _result: &'static str, _duration: Duration) {} +} + +async fn validate_harness_key_with_timeout( + validator: &V, + harness_public_key: &NoiseChannelPublicKey, + authorization: &str, + validation_timeout: Duration, +) -> Result<(), ExecServerError> { + let started_at = Instant::now(); + let (result, result_name) = match timeout( + validation_timeout, + validator.validate_harness_key(harness_public_key, authorization), + ) + .await + { + Ok(Ok(())) => (Ok(()), "success"), + Ok(Err(error)) => (Err(error), "error"), + Err(_) => ( + Err(ExecServerError::Protocol( + "timed out validating Noise relay harness key".to_string(), + )), + "timeout", + ), + }; + let duration = started_at.elapsed(); + validator.validation_completed(result_name, duration); + record_remote_phase( + &tracing::Span::current(), + REMOTE_PHASE_EXECUTOR_AUTHORIZATION, + result_name, + duration, + ); + result } /// Serve authenticated virtual JSON-RPC streams over one executor websocket. @@ -730,21 +769,25 @@ pub(crate) async fn run_multiplexed_environment( }, ); let validator = validator.clone(); + let validation_span = tracing::info_span!( + "codex.exec_server.remote.executor.authorization", + otel.kind = "client", + otel.name = "codex.exec_server.remote.executor.authorization", + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + ); // Failed validation leaves no transport state and sends only a // generic reset. validation_tasks.spawn(async move { - let result = match timeout( + let result = validate_harness_key_with_timeout( + &validator, + &harness_public_key, + &authorization, HARNESS_KEY_VALIDATION_TIMEOUT, - validator.validate_harness_key(&harness_public_key, &authorization), ) - .await - { - Ok(result) => result, - Err(_) => Err(ExecServerError::Protocol( - "timed out validating Noise relay harness key".to_string(), - )), - }; + .instrument(validation_span) + .await; HarnessKeyValidationResult { stream_id, validation_id, diff --git a/codex-rs/exec-server/src/relay_noise_tests.rs b/codex-rs/exec-server/src/relay_noise_tests.rs index 0020efe34766..4f9d49d92184 100644 --- a/codex-rs/exec-server/src/relay_noise_tests.rs +++ b/codex-rs/exec-server/src/relay_noise_tests.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::sync::Mutex; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::time::Duration; @@ -18,6 +19,7 @@ use super::HarnessKeyValidator; use super::MAX_FAILED_NOISE_HANDSHAKES; use super::MAX_HARNESS_KEY_AUTHORIZATION_BYTES; use super::run_multiplexed_environment; +use super::validate_harness_key_with_timeout; use crate::ExecServerError; use crate::ExecServerRuntimePaths; use crate::noise_channel::InitiatorHandshake; @@ -55,6 +57,53 @@ impl HarnessKeyValidator for BlockingValidator { } } +#[derive(Clone, Default)] +struct TimeoutObservingValidator { + outcomes: Arc>>, +} + +impl HarnessKeyValidator for TimeoutObservingValidator { + async fn validate_harness_key( + &self, + _harness_public_key: &NoiseChannelPublicKey, + _authorization: &str, + ) -> Result<(), ExecServerError> { + std::future::pending().await + } + + fn validation_completed(&self, result: &'static str, _duration: Duration) { + self.outcomes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .push(result); + } +} + +#[tokio::test] +async fn harness_key_validation_records_timeout_at_timeout_boundary() -> Result<()> { + let validator = TimeoutObservingValidator::default(); + let harness_public_key = NoiseChannelIdentity::generate()?.public_key(); + + let error = validate_harness_key_with_timeout( + &validator, + &harness_public_key, + "authorization", + Duration::from_millis(1), + ) + .await + .expect_err("validation should time out"); + + assert!(error.to_string().contains("timed out validating")); + assert_eq!( + *validator + .outcomes + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner), + vec!["timeout"] + ); + Ok(()) +} + #[tokio::test] async fn pending_harness_key_validation_does_not_block_new_handshakes() -> Result<()> { let listener = TcpListener::bind("127.0.0.1:0").await?; diff --git a/codex-rs/exec-server/src/remote.rs b/codex-rs/exec-server/src/remote.rs index f5e233ce9d75..fdb549427ae6 100644 --- a/codex-rs/exec-server/src/remote.rs +++ b/codex-rs/exec-server/src/remote.rs @@ -171,6 +171,7 @@ impl EnvironmentRegistryClient { &format!("/cloud/environment/{environment_id}/connect"), )) .headers(self.auth_provider.to_auth_headers()) + .headers(current_trace_context_headers()) .json(&EnvironmentRegistryConnectRequest { harness_public_key }) .timeout(self.connect_timeout) .send() @@ -251,6 +252,7 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { &format!("/cloud/environment/{environment_id}/validate"), )) .headers(self.client.auth_provider.to_auth_headers()) + .headers(current_trace_context_headers()) .json(&EnvironmentRegistryHarnessKeyValidationRequest { executor_registration_id: self.executor_registration_id.clone(), harness_public_key: harness_public_key.clone(), @@ -283,6 +285,12 @@ impl HarnessKeyValidator for RegistryHarnessKeyValidator { } Ok(()) } + + fn validation_completed(&self, result: &'static str, duration: Duration) { + self.client + .telemetry + .remote_executor_authorization_completed(result, duration); + } } /// Noise connection configuration for a Codex harness. @@ -561,7 +569,7 @@ async fn connect_rendezvous( connect_async_with_config( request, Some(noise_relay_websocket_config()), - /*disable_nagle*/ false, + /*disable_nagle*/ true, ) .await .map(|(websocket, _)| websocket) @@ -752,8 +760,14 @@ mod tests { ); } - #[tokio::test] + #[tokio::test(flavor = "current_thread")] async fn noise_connect_provider_requests_and_validates_a_full_bundle() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); let server = MockServer::start().await; let harness_public_key = NoiseChannelIdentity::generate() .expect("identity") @@ -765,6 +779,10 @@ mod tests { .and(path("/cloud/environment/environment-requested/connect")) .and(header("authorization", "Bearer registry-token")) .and(header("chatgpt-account-id", "workspace-123")) + .and(header_regex( + "traceparent", + "^00-[0-9a-f]{32}-[0-9a-f]{16}-0[01]$", + )) .and(body_partial_json(serde_json::json!({ "harness_public_key": harness_public_key.clone(), }))) @@ -789,6 +807,7 @@ mod tests { let bundle = config .connect_provider() .connect_bundle(harness_public_key) + .instrument(tracing::info_span!("remote-operation")) .await .expect("Noise connect bundle"); diff --git a/codex-rs/exec-server/src/remote/noise_tests.rs b/codex-rs/exec-server/src/remote/noise_tests.rs index 1c4525e63f9b..26759d78f35f 100644 --- a/codex-rs/exec-server/src/remote/noise_tests.rs +++ b/codex-rs/exec-server/src/remote/noise_tests.rs @@ -6,16 +6,21 @@ use codex_api::AuthProvider; use codex_api::SharedAuthProvider; use http::HeaderMap; use http::HeaderValue; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::trace::SdkTracerProvider; use tokio::io::AsyncReadExt; use tokio::io::AsyncWriteExt; use tokio::net::TcpListener; use tokio::time::timeout; use tokio_tungstenite::accept_async; +use tracing::Instrument; +use tracing_subscriber::prelude::*; use wiremock::Mock; use wiremock::MockServer; use wiremock::ResponseTemplate; use wiremock::matchers::body_partial_json; use wiremock::matchers::header; +use wiremock::matchers::header_regex; use wiremock::matchers::method; use wiremock::matchers::path; @@ -92,8 +97,14 @@ async fn reconnect_reuses_registration_until_url_is_rejected() -> Result<()> { Ok(()) } -#[tokio::test] +#[tokio::test(flavor = "current_thread")] async fn validate_harness_key_requires_explicit_valid_response() { + let provider = SdkTracerProvider::builder().build(); + let tracer = provider.tracer("exec-server-test"); + let subscriber = + tracing_subscriber::registry().with(tracing_opentelemetry::layer().with_tracer(tracer)); + let _guard = subscriber.set_default(); + tracing::callsite::rebuild_interest_cache(); let server = MockServer::start().await; let harness_public_key = NoiseChannelIdentity::generate() .expect("identity") @@ -101,6 +112,10 @@ async fn validate_harness_key_requires_explicit_valid_response() { Mock::given(method("POST")) .and(path("/cloud/environment/environment-requested/validate")) .and(header("authorization", "Bearer registry-token")) + .and(header_regex( + "traceparent", + "^00-[0-9a-f]{32}-[0-9a-f]{16}-0[01]$", + )) .and(body_partial_json(serde_json::json!({ "executor_registration_id": "registration-1", "harness_public_key": harness_public_key.clone(), @@ -120,6 +135,7 @@ async fn validate_harness_key_requires_explicit_valid_response() { executor_registration_id: "registration-1".to_string(), } .validate_harness_key(&harness_public_key, HARNESS_KEY_AUTHORIZATION) + .instrument(tracing::info_span!("remote-operation")) .await .expect_err("a false validation response must fail closed"); diff --git a/codex-rs/exec-server/src/rpc.rs b/codex-rs/exec-server/src/rpc.rs index a941dce41f34..df18b1065a9e 100644 --- a/codex-rs/exec-server/src/rpc.rs +++ b/codex-rs/exec-server/src/rpc.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicI64; use std::sync::atomic::Ordering; +use std::time::Instant; use codex_exec_server_protocol::JSONRPCError; use codex_exec_server_protocol::JSONRPCErrorError; @@ -333,59 +334,118 @@ impl RpcClient { drain_pending(&self.pending).await; } + #[tracing::instrument( + name = "codex.exec_server.rpc.client_call", + skip_all, + fields( + otel.kind = "client", + otel.name = "codex.exec_server.rpc.client_call", + rpc.method = method, + result = tracing::field::Empty, + duration_ms = tracing::field::Empty, + pending_registration_ms = tracing::field::Empty, + serialize_ms = tracing::field::Empty, + enqueue_ms = tracing::field::Empty, + response_wait_ms = tracing::field::Empty, + deserialize_ms = tracing::field::Empty, + ) + )] pub(crate) async fn call(&self, method: &str, params: &P) -> Result where P: Serialize, T: DeserializeOwned, { - let request_id = RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); - let (response_tx, response_rx) = oneshot::channel(); - { - let mut pending = self.pending.lock().await; - // Registering the pending request and checking disconnect must be - // atomic with the reader's drain_pending path. Otherwise a call - // can sneak in after the drain and wait forever. - if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { - return Err(RpcCallError::Closed); + let started_at = Instant::now(); + let span = tracing::Span::current(); + let result = async { + let request_id = + RequestId::Integer(self.next_request_id.fetch_add(1, Ordering::SeqCst)); + let (response_tx, response_rx) = oneshot::channel(); + let pending_registration_started_at = Instant::now(); + { + let mut pending = self.pending.lock().await; + // Registering the pending request and checking disconnect must be + // atomic with the reader's drain_pending path. Otherwise a call + // can sneak in after the drain and wait forever. + if self.closed.load(Ordering::Acquire) || *self.disconnected_rx.borrow() { + span.record( + "pending_registration_ms", + elapsed_ms(pending_registration_started_at), + ); + return Err(RpcCallError::Closed); + } + pending.insert(request_id.clone(), response_tx); } - pending.insert(request_id.clone(), response_tx); - } + span.record( + "pending_registration_ms", + elapsed_ms(pending_registration_started_at), + ); + + let serialize_started_at = Instant::now(); + let params = match serde_json::to_value(params) { + Ok(params) => params, + Err(err) => { + self.pending.lock().await.remove(&request_id); + span.record("serialize_ms", elapsed_ms(serialize_started_at)); + return Err(RpcCallError::Json(err)); + } + }; + span.record("serialize_ms", elapsed_ms(serialize_started_at)); - let params = match serde_json::to_value(params) { - Ok(params) => params, - Err(err) => { + let enqueue_started_at = Instant::now(); + if self + .write_tx + .send(JSONRPCMessage::Request(JSONRPCRequest { + id: request_id.clone(), + method: method.to_string(), + params: Some(params), + trace: codex_otel::current_span_w3c_trace_context(), + })) + .await + .is_err() + { self.pending.lock().await.remove(&request_id); - return Err(RpcCallError::Json(err)); + span.record("enqueue_ms", elapsed_ms(enqueue_started_at)); + return Err(RpcCallError::Closed); } - }; - if self - .write_tx - .send(JSONRPCMessage::Request(JSONRPCRequest { - id: request_id.clone(), - method: method.to_string(), - params: Some(params), - trace: codex_otel::current_span_w3c_trace_context(), - })) - .await - .is_err() - { - self.pending.lock().await.remove(&request_id); - return Err(RpcCallError::Closed); + span.record("enqueue_ms", elapsed_ms(enqueue_started_at)); + + // Do not race in-flight requests directly against the transport-close + // watch value. The connection reader receives JSON-RPC messages and + // the terminal disconnect event on one ordered queue, then drains any + // still-pending requests. Awaiting this receiver preserves that order: + // responses already read before EOF still win, and truly pending calls + // are failed once the reader observes the disconnect. + let response_wait_started_at = Instant::now(); + let response = match response_rx.await { + Ok(response) => response, + Err(_) => { + span.record("response_wait_ms", elapsed_ms(response_wait_started_at)); + return Err(RpcCallError::Closed); + } + }; + span.record("response_wait_ms", elapsed_ms(response_wait_started_at)); + let response = match response { + Ok(response) => response, + Err(error) => return Err(error), + }; + + let deserialize_started_at = Instant::now(); + let response = serde_json::from_value(response).map_err(RpcCallError::Json); + span.record("deserialize_ms", elapsed_ms(deserialize_started_at)); + response } + .await; - // Do not race in-flight requests directly against the transport-close - // watch value. The connection reader receives JSON-RPC messages and - // the terminal disconnect event on one ordered queue, then drains any - // still-pending requests. Awaiting this receiver preserves that order: - // responses already read before EOF still win, and truly pending calls - // are failed once the reader observes the disconnect. - let result: Result = - response_rx.await.map_err(|_| RpcCallError::Closed)?; - let response = match result { - Ok(response) => response, - Err(error) => return Err(error), + let result_name = match &result { + Ok(_) => "success", + Err(RpcCallError::Closed) => "closed", + Err(RpcCallError::Json(_)) => "json_error", + Err(RpcCallError::Server(_)) => "server_error", }; - serde_json::from_value(response).map_err(RpcCallError::Json) + span.record("result", result_name); + span.record("duration_ms", elapsed_ms(started_at)); + result } #[cfg(test)] @@ -394,6 +454,10 @@ impl RpcClient { } } +fn elapsed_ms(started_at: Instant) -> f64 { + started_at.elapsed().as_secs_f64() * 1_000.0 +} + impl Drop for RpcClient { fn drop(&mut self) { self.transport.terminate(); @@ -552,10 +616,12 @@ async fn drain_pending(pending: &Mutex>) { #[cfg(test)] mod tests { + use std::collections::BTreeMap; use std::time::Duration; use codex_exec_server_protocol::JSONRPCMessage; use codex_exec_server_protocol::JSONRPCResponse; + use opentelemetry::Value; use opentelemetry::trace::TracerProvider as _; use opentelemetry_sdk::trace::InMemorySpanExporter; use opentelemetry_sdk::trace::SdkTracerProvider; @@ -677,7 +743,7 @@ mod tests { async fn rpc_client_propagates_current_trace_context() { let span_exporter = InMemorySpanExporter::default(); let tracer_provider = SdkTracerProvider::builder() - .with_simple_exporter(span_exporter) + .with_simple_exporter(span_exporter.clone()) .build(); let tracer = tracer_provider.tracer("exec-server-test"); let subscriber = tracing_subscriber::registry().with( @@ -721,6 +787,51 @@ mod tests { .expect("RPC response"); assert_eq!(response, serde_json::json!({})); let trace = server.await.expect("server task").expect("trace context"); - assert_eq!(trace, expected_trace); + + tracer_provider.force_flush().expect("flush traces"); + let spans = span_exporter.get_finished_spans().expect("span export"); + let rpc_span = spans + .iter() + .find(|span| span.name.as_ref() == "codex.exec_server.rpc.client_call") + .unwrap_or_else(|| panic!("RPC client span missing: {spans:?}")); + let attributes = rpc_span + .attributes + .iter() + .map(|attribute| (attribute.key.as_str().to_string(), attribute.value.clone())) + .collect::>(); + assert_eq!( + attributes.get("rpc.method"), + Some(&Value::String("traced".into())) + ); + assert_eq!( + attributes.get("result"), + Some(&Value::String("success".into())) + ); + let parent_traceparent = expected_trace + .traceparent + .as_deref() + .expect("parent traceparent"); + let request_traceparent = trace.traceparent.as_deref().expect("request traceparent"); + let parent_parts = parent_traceparent.split('-').collect::>(); + let request_parts = request_traceparent.split('-').collect::>(); + assert_eq!(request_parts[1], parent_parts[1]); + assert_eq!( + request_parts[2], + rpc_span.span_context.span_id().to_string() + ); + assert_eq!(rpc_span.parent_span_id.to_string(), parent_parts[2]); + for field in [ + "duration_ms", + "pending_registration_ms", + "serialize_ms", + "enqueue_ms", + "response_wait_ms", + "deserialize_ms", + ] { + match attributes.get(field) { + Some(Value::F64(value)) if *value >= 0.0 => {} + value => panic!("expected non-negative {field}, got {value:?}"), + } + } } } diff --git a/codex-rs/exec-server/src/telemetry.rs b/codex-rs/exec-server/src/telemetry.rs index 650faacb2c63..f1dcb37c32ac 100644 --- a/codex-rs/exec-server/src/telemetry.rs +++ b/codex-rs/exec-server/src/telemetry.rs @@ -33,9 +33,22 @@ const REMOTE_RENDEZVOUS_METRICS: OperationMetrics = OperationMetrics { duration_name: "exec_server_remote_rendezvous_connect_duration_seconds", duration_description: "Duration of remote exec-server rendezvous connection attempts in seconds.", }; +const REMOTE_EXECUTOR_AUTHORIZATION_METRICS: OperationMetrics = OperationMetrics { + total_name: "exec_server_remote_executor_authorization_total", + total_description: "Total number of remote exec-server executor authorization attempts.", + duration_name: "exec_server_remote_executor_authorization_duration_seconds", + duration_description: "Duration of remote exec-server executor authorization attempts in seconds.", +}; const REMOTE_RECONNECTS_TOTAL_METRIC: &str = "exec_server_remote_reconnects_total"; const REMOTE_RECONNECTS_TOTAL_DESCRIPTION: &str = "Total number of remote exec-server reconnects."; +pub(crate) const REMOTE_PHASE_EXECUTOR_AUTHORIZATION: &str = "executor_authorization"; +pub(crate) const REMOTE_PHASE_HARNESS_INITIALIZE: &str = "harness_initialize"; +pub(crate) const REMOTE_PHASE_HARNESS_NOISE_HANDSHAKE: &str = "harness_noise_handshake"; +pub(crate) const REMOTE_PHASE_HARNESS_READY: &str = "harness_ready"; +pub(crate) const REMOTE_PHASE_HARNESS_WEBSOCKET: &str = "harness_websocket"; +pub(crate) const REMOTE_PHASE_REGISTRY_CONNECT_BUNDLE: &str = "registry_connect_bundle"; + #[derive(Clone, Copy)] struct OperationMetrics { total_name: &'static str, @@ -153,6 +166,14 @@ impl ExecServerTelemetry { self.record_operation(REMOTE_RENDEZVOUS_METRICS, result, duration); } + pub(crate) fn remote_executor_authorization_completed( + &self, + result: &'static str, + duration: Duration, + ) { + self.record_operation(REMOTE_EXECUTOR_AUTHORIZATION_METRICS, result, duration); + } + pub(crate) fn remote_reconnect(&self, reason: &'static str) { self.with_inner(|inner| { inner.counter( @@ -222,6 +243,65 @@ impl ExecServerTelemetry { } } +pub(crate) fn record_remote_phase( + span: &tracing::Span, + phase: &'static str, + result: &'static str, + duration: Duration, +) { + let duration_ms = duration.as_secs_f64() * 1_000.0; + span.record("result", result); + span.record("duration_ms", duration_ms); + tracing::info!( + parent: span, + remote_phase = phase, + result, + duration_ms, + "Remote exec-server phase completed" + ); +} + +pub(crate) struct RemotePhaseGuard { + span: tracing::Span, + phase: &'static str, + started_at: Instant, + result: &'static str, +} + +impl RemotePhaseGuard { + pub(crate) fn new( + span: tracing::Span, + phase: &'static str, + default_result: &'static str, + ) -> Self { + Self { + span, + phase, + started_at: Instant::now(), + result: default_result, + } + } + + pub(crate) fn span(&self) -> &tracing::Span { + &self.span + } + + pub(crate) fn finish(mut self, result: &'static str) { + self.result = result; + } +} + +impl Drop for RemotePhaseGuard { + fn drop(&mut self) { + record_remote_phase( + &self.span, + self.phase, + self.result, + self.started_at.elapsed(), + ); + } +} + impl Drop for ConnectionMetricGuard { fn drop(&mut self) { self.telemetry.connection_finished(self.transport); @@ -338,3 +418,59 @@ fn register_active_gauge( warn!(metric = name, "failed to register exec-server gauge"); } } + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use codex_otel::MetricsConfig; + use opentelemetry_sdk::metrics::InMemoryMetricExporter; + use opentelemetry_sdk::metrics::data::AggregatedMetrics; + use opentelemetry_sdk::metrics::data::MetricData; + use pretty_assertions::assert_eq; + + use super::*; + + #[test] + fn executor_authorization_duration_uses_only_bounded_result_tag() { + let exporter = InMemoryMetricExporter::default(); + let metrics = MetricsClient::new( + MetricsConfig::in_memory( + "test", + "codex-exec-server", + env!("CARGO_PKG_VERSION"), + exporter, + ) + .with_runtime_reader(), + ) + .expect("metrics"); + let telemetry = ExecServerTelemetry::new(metrics.clone()); + + telemetry.remote_executor_authorization_completed("timeout", Duration::from_millis(25)); + + let snapshot = metrics.snapshot().expect("metrics snapshot"); + let metric = snapshot + .scope_metrics() + .flat_map(opentelemetry_sdk::metrics::data::ScopeMetrics::metrics) + .find(|metric| metric.name() == REMOTE_EXECUTOR_AUTHORIZATION_METRICS.duration_name) + .expect("executor authorization duration metric"); + let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else { + panic!("executor authorization duration should be an f64 histogram"); + }; + let point = histogram.data_points().next().expect("histogram point"); + let attributes = point + .attributes() + .map(|attribute| { + ( + attribute.key.as_str().to_string(), + attribute.value.as_str().into_owned(), + ) + }) + .collect::>(); + + assert_eq!( + attributes, + BTreeMap::from([("result".to_string(), "timeout".to_string())]) + ); + } +} diff --git a/codex-rs/exec-server/tests/exec_process.rs b/codex-rs/exec-server/tests/exec_process.rs index 8fe5fe0be98e..c8d0e3e985ee 100644 --- a/codex-rs/exec-server/tests/exec_process.rs +++ b/codex-rs/exec-server/tests/exec_process.rs @@ -167,7 +167,7 @@ async fn collect_process_output_from_events( } => { exit_code = Some(code); } - ExecProcessEvent::Closed { seq: _ } => { + ExecProcessEvent::Closed { seq: _, .. } => { drop(session); return Ok((stdout, stderr, exit_code, true)); } @@ -193,7 +193,7 @@ async fn collect_process_event_snapshots( ExecProcessEvent::Exited { seq, exit_code } => { ProcessEventSnapshot::Exited { seq, exit_code } } - ExecProcessEvent::Closed { seq } => ProcessEventSnapshot::Closed { seq }, + ExecProcessEvent::Closed { seq, .. } => ProcessEventSnapshot::Closed { seq }, ExecProcessEvent::Failed(message) => { anyhow::bail!("process failed before closed state: {message}"); } @@ -798,7 +798,7 @@ async fn remote_exec_process_recovers_after_transport_disconnect() -> Result<()> last_seq = seq; saw_exit = true; } - ExecProcessEvent::Closed { seq } => { + ExecProcessEvent::Closed { seq, .. } => { assert!(saw_exit, "closed must be delivered after exit"); assert_eq!(seq, last_seq + 1); break; diff --git a/codex-rs/rmcp-client/src/executor_process_transport.rs b/codex-rs/rmcp-client/src/executor_process_transport.rs index cff4f517ed01..9b88b24450ec 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport.rs @@ -226,7 +226,7 @@ impl ExecutorProcessTransport { // output flushed during process shutdown can still be // decoded into JSON-RPC messages. } - Ok(ExecProcessEvent::Closed { seq }) => { + Ok(ExecProcessEvent::Closed { seq, .. }) => { self.note_seq(seq); self.closed = true; }