From 3ba7b6941d3caf6eec5b3c4e564988ee57d3f083 Mon Sep 17 00:00:00 2001 From: yqz-openai Date: Wed, 26 Aug 2026 02:44:08 +0000 Subject: [PATCH] Add OpenTelemetry tracing to the code-mode host (#40760) ## What changed - Add `--otel-trace-exporter` for sending OTLP/HTTP JSON traces directly to an exporter. - Add the mutually exclusive `--otel-trace-listen` option for streaming raw OTLP trace batches over a trace-only WebSocket, backed by a loopback OTLP receiver. - Instrument host startup, transport requests, and gRPC operations with request, session, execution, cell, wait, and tool-call identifiers where available. - Flush the trace provider with a bounded timeout during shutdown. GitOrigin-RevId: 4dcdb2bc3b15166dccd5041fb875caea74f3f6fe --- codex-rs/Cargo.lock | 1 + codex-rs/code-mode-host/Cargo.toml | 3 +- codex-rs/code-mode-host/src/grpc/mod.rs | 37 ++++ codex-rs/code-mode-host/src/lib.rs | 39 ++++- codex-rs/code-mode-host/src/main.rs | 115 ++++++++++++- .../code-mode-host/src/trace_transport.rs | 161 ++++++++++++++++++ 6 files changed, 349 insertions(+), 7 deletions(-) create mode 100644 codex-rs/code-mode-host/src/trace_transport.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index aa248f4633..2d3ebb535a 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -2597,6 +2597,7 @@ dependencies = [ "codex-code-mode", "codex-code-mode-protocol", "codex-code-mode-runtime", + "codex-otel", "codex-protocol", "codex-utils-cargo-bin", "futures", diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index 0045279df5..a064dd5759 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/codex-rs/code-mode-host/Cargo.toml @@ -18,10 +18,11 @@ workspace = true [dependencies] anyhow = { workspace = true } -axum = { workspace = true, features = ["http1", "tokio"] } +axum = { workspace = true, features = ["http1", "tokio", "ws"] } clap = { workspace = true, features = ["derive"] } codex-code-mode-protocol = { workspace = true } codex-code-mode-runtime = { workspace = true } +codex-otel = { workspace = true } codex-protocol = { workspace = true } futures = { workspace = true } prost = "0.14.3" diff --git a/codex-rs/code-mode-host/src/grpc/mod.rs b/codex-rs/code-mode-host/src/grpc/mod.rs index 5cb991af48..4647f6cdce 100644 --- a/codex-rs/code-mode-host/src/grpc/mod.rs +++ b/codex-rs/code-mode-host/src/grpc/mod.rs @@ -43,6 +43,12 @@ impl GrpcCodeModeHost { } } + #[tracing::instrument( + name = "code_mode_host.grpc.open_session", + level = "info", + skip_all, + fields(otel.name = "code_mode_host.grpc.open_session") + )] async fn open_session_request( &self, request: proto::OpenSessionRequest, @@ -52,6 +58,15 @@ impl GrpcCodeModeHost { Ok(Response::new(self.state.open_session(limits)?)) } + #[tracing::instrument( + name = "code_mode_host.grpc.close_session", + level = "info", + skip_all, + fields( + otel.name = "code_mode_host.grpc.close_session", + session.id = %request.session_id, + ) + )] async fn close_session_request( &self, request: proto::CloseSessionRequest, @@ -104,6 +119,17 @@ impl GrpcCodeModeHost { Ok(Response::new(proto::AcknowledgeNotificationResponse {})) } + #[tracing::instrument( + name = "code_mode_host.grpc.execute", + level = "info", + skip_all, + fields( + otel.name = "code_mode_host.grpc.execute", + session.id = %request.session_id, + execution.id = %request.execution_id, + call_id = %request.tool_call_id, + ) + )] async fn execute_request( &self, request: proto::ExecuteRequest, @@ -171,6 +197,17 @@ impl GrpcCodeModeHost { Ok(Response::new(Box::pin(stream))) } + #[tracing::instrument( + name = "code_mode_host.grpc.wait", + level = "info", + skip_all, + fields( + otel.name = "code_mode_host.grpc.wait", + session.id = %request.session_id, + cell.id = %request.cell_id, + wait.id = %request.wait_id, + ) + )] async fn wait_request( &self, request: proto::WaitRequest, diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index d61df805fa..6a4557cddd 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/codex-rs/code-mode-host/src/lib.rs @@ -47,8 +47,14 @@ mod delegate; mod grpc; mod grpc_transport; mod peer; +mod trace_transport; mod transport; +pub use self::trace_transport::bind_otlp_trace_receiver; +pub use self::trace_transport::run_otel_trace_listener; +pub use self::trace_transport::run_otlp_trace_receiver; +pub use self::trace_transport::trace_batch_channel; + const MAX_IN_FLIGHT_REQUESTS: usize = 256; const MAX_ACTIVE_CELLS: usize = 128; const MAX_RECENT_REQUEST_IDS: usize = 4096; @@ -84,6 +90,12 @@ impl HostLimits { } /// Runs the code-mode host on its configured stdio or gRPC transport. +#[tracing::instrument( + name = "code_mode_host.run_main", + level = "info", + skip_all, + fields(otel.name = "code_mode_host.run_main") +)] pub async fn run_main(listen_url: &str) -> Result<()> { transport::run_transport(listen_url).await } @@ -294,11 +306,12 @@ impl HostState { request_id: RequestId, request: HostRequest, ) -> Result<(), anyhow::Error> { + let request_kind = RequestKind::from(&request); let cancellation = self .requests .lock() .unwrap_or_else(PoisonError::into_inner) - .start(request_id, RequestKind::from(&request))?; + .start(request_id, request_kind)?; let Ok(permit) = self.limits.request_permit() else { self.respond( request_id, @@ -328,12 +341,26 @@ impl HostState { }); } + #[tracing::instrument( + name = "code_mode_host.request", + level = "info", + skip_all, + fields( + otel.name = "code_mode_host.request", + request.kind = RequestKind::from(&request).as_str(), + request.id = ?request_id, + call_id = tracing::field::Empty, + ) + )] async fn handle_request( &self, request_id: RequestId, request: HostRequest, cancellation: CancellationToken, ) { + if let HostRequest::Execute { request, .. } = &request { + tracing::Span::current().record("call_id", request.tool_call_id.as_str()); + } if self.closing.load(Ordering::Acquire) { self.respond( request_id, @@ -580,6 +607,16 @@ impl RequestKind { fn is_cancellable(self) -> bool { matches!(self, Self::Execute | Self::Wait) } + + fn as_str(self) -> &'static str { + match self { + Self::OpenSession => "open_session", + Self::Execute => "execute", + Self::Wait => "wait", + Self::Terminate => "terminate", + Self::ShutdownSession => "shutdown_session", + } + } } struct ActiveRequest { diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index 52e2346513..81ffb3ae78 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/codex-rs/code-mode-host/src/main.rs @@ -1,4 +1,19 @@ +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::path::PathBuf; +use std::time::Duration; + +use anyhow::Context; use clap::Parser; +use codex_otel::OtelExporter; +use codex_otel::OtelHttpProtocol; +use codex_otel::OtelProvider; +use codex_otel::OtelSettings; +use tracing_subscriber::Layer; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; + +const OTEL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 5); #[derive(Debug, Parser)] struct Cli { @@ -9,15 +24,105 @@ struct Cli { default_value = codex_code_mode_host::DEFAULT_LISTEN_URL )] listen: String, + + /// Optional WebSocket endpoint that streams only raw OTLP trace batches. + #[arg(long, value_name = "URL")] + otel_trace_listen: Option, + + /// Optional OTLP/HTTP JSON trace exporter endpoint, analogous to + /// `otel.trace_exporter` in app-server configuration. + #[arg(long, value_name = "URL", conflicts_with = "otel_trace_listen")] + otel_trace_exporter: Option, } #[tokio::main(flavor = "current_thread")] async fn main() -> anyhow::Result<()> { - tracing_subscriber::fmt() - .with_max_level(tracing::Level::INFO) - .with_writer(std::io::stderr) - .with_ansi(false) + let cli = Cli::parse(); + let trace_transport = if let Some(trace_listen) = cli.otel_trace_listen.as_deref() { + let sender = codex_code_mode_host::trace_batch_channel(); + let (receiver, exporter_endpoint) = + codex_code_mode_host::bind_otlp_trace_receiver().await?; + Some(( + trace_listen.to_string(), + receiver, + sender, + exporter_endpoint, + )) + } else { + None + }; + let trace_exporter_endpoint = trace_transport + .as_ref() + .map(|(_, _, _, endpoint)| endpoint.as_str()) + .or(cli.otel_trace_exporter.as_deref()); + let otel = trace_exporter_endpoint + .map(build_trace_provider) + .transpose()?; + let otel_layer = otel.as_ref().and_then(OtelProvider::tracing_layer); + tracing_subscriber::registry() + .with( + tracing_subscriber::fmt::layer() + .with_writer(std::io::stderr) + .with_ansi(false) + .with_filter(tracing_subscriber::filter::LevelFilter::INFO), + ) + .with(otel_layer) .init(); + tracing::info_span!( + "code_mode_host.startup", + otel.name = "code_mode_host.startup" + ) + .in_scope(|| {}); - codex_code_mode_host::run_main(&Cli::parse().listen).await + let main_transport = codex_code_mode_host::run_main(&cli.listen); + let (result, trace_tasks) = match trace_transport { + Some((trace_listen, receiver, sender, _)) => { + let mut trace_listener_task = tokio::spawn({ + let sender = sender.clone(); + async move { codex_code_mode_host::run_otel_trace_listener(&trace_listen, sender).await } + }); + let mut trace_receiver_task = tokio::spawn( + codex_code_mode_host::run_otlp_trace_receiver(receiver, sender), + ); + let result = tokio::select! { + result = main_transport => result, + result = &mut trace_listener_task => result.context("OTEL trace websocket task failed").flatten(), + result = &mut trace_receiver_task => result.context("OTLP trace receiver task failed").flatten(), + }; + (result, Some((trace_listener_task, trace_receiver_task))) + } + None => (main_transport.await, None), + }; + if let Some(otel) = otel + && let Err(error) = otel.shutdown_with_timeout(OTEL_SHUTDOWN_TIMEOUT).await + { + tracing::warn!(%error, "failed to finish code-mode host telemetry shutdown"); + } + if let Some((trace_listener_task, trace_receiver_task)) = trace_tasks { + trace_listener_task.abort(); + trace_receiver_task.abort(); + } + result +} + +fn build_trace_provider(endpoint: &str) -> anyhow::Result { + OtelProvider::try_new(&OtelSettings { + environment: "code-mode-host".to_string(), + service_name: "codex-code-mode-host".to_string(), + service_version: env!("CARGO_PKG_VERSION").to_string(), + codex_home: PathBuf::from("/tmp"), + exporter: OtelExporter::None, + trace_exporter: OtelExporter::OtlpHttp { + endpoint: endpoint.to_string(), + headers: HashMap::new(), + protocol: OtelHttpProtocol::Json, + tls: None, + }, + metrics_exporter: OtelExporter::None, + runtime_metrics: false, + span_attributes: BTreeMap::new(), + tracestate: BTreeMap::new(), + }) + .map_err(|error| anyhow::anyhow!("failed to build code-mode host OTEL provider: {error}"))? + .context("code-mode host OTEL trace provider was unexpectedly disabled") } diff --git a/codex-rs/code-mode-host/src/trace_transport.rs b/codex-rs/code-mode-host/src/trace_transport.rs new file mode 100644 index 0000000000..971b289a62 --- /dev/null +++ b/codex-rs/code-mode-host/src/trace_transport.rs @@ -0,0 +1,161 @@ +//! OTLP trace-only WebSocket transport for forwarding code-mode host spans. + +use std::net::SocketAddr; + +use anyhow::Context; +use anyhow::Result; +use axum::Router; +use axum::body::Body; +use axum::body::Bytes; +use axum::extract::State; +use axum::extract::ws::Message; +use axum::extract::ws::WebSocket; +use axum::extract::ws::WebSocketUpgrade; +use axum::http::Request; +use axum::http::StatusCode; +use axum::http::header::CONTENT_TYPE; +use axum::http::header::ORIGIN; +use axum::middleware; +use axum::middleware::Next; +use axum::response::IntoResponse; +use axum::response::Response; +use axum::routing::any; +use axum::routing::post; +use futures::SinkExt; +use futures::StreamExt; +use tokio::net::TcpListener; +use tokio::sync::broadcast; +use tracing::info; +use tracing::warn; + +const TRACE_CHANNEL_CAPACITY: usize = 1_024; +const MAX_TRACE_WEBSOCKET_FRAME_BYTES: usize = 4 * 1024 * 1024; + +/// Creates the channel shared by the loopback OTLP receiver and trace WebSocket. +pub fn trace_batch_channel() -> broadcast::Sender> { + let (sender, _) = broadcast::channel(TRACE_CHANNEL_CAPACITY); + sender +} + +/// Binds the loopback OTLP receiver used when trace WebSocket forwarding is enabled. +pub async fn bind_otlp_trace_receiver() -> Result<(TcpListener, String)> { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .context("failed to bind loopback OTLP trace receiver")?; + let local_address = listener + .local_addr() + .context("failed to read OTLP trace receiver listen address")?; + let endpoint = format!("http://{local_address}/v1/traces"); + Ok((listener, endpoint)) +} + +/// Receives OTLP/HTTP JSON trace batches from this process's local exporter. +pub async fn run_otlp_trace_receiver( + listener: TcpListener, + sender: broadcast::Sender>, +) -> Result<()> { + let local_address = listener + .local_addr() + .context("failed to read OTLP trace receiver listen address")?; + info!("codex-code-mode-host OTLP trace receiver listening on http://{local_address}"); + let router = Router::new() + .route("/v1/traces", post(receive_otlp_trace_batch)) + .with_state(sender); + axum::serve(listener, router) + .await + .context("code-mode host OTLP trace receiver failed") +} + +/// Serves a second WebSocket that streams only raw OTLP trace batches. +pub async fn run_otel_trace_listener( + listen_url: &str, + sender: broadcast::Sender>, +) -> Result<()> { + let bind_address = parse_listen_url(listen_url, "--otel-trace-listen", "ws://")?; + let listener = TcpListener::bind(bind_address) + .await + .with_context(|| format!("failed to bind OTEL trace websocket to {bind_address}"))?; + let local_address = listener + .local_addr() + .context("failed to read OTEL trace websocket listen address")?; + info!("codex-code-mode-host OTEL trace websocket listening on ws://{local_address}"); + let router = Router::new() + .route("/", any(trace_websocket_upgrade_handler)) + .layer(middleware::from_fn(reject_requests_with_origin_header)) + .with_state(sender); + axum::serve(listener, router) + .await + .context("code-mode host OTEL trace websocket listener failed") +} + +async fn reject_requests_with_origin_header( + request: Request, + next: Next, +) -> std::result::Result { + if request.headers().contains_key(ORIGIN) { + warn!( + method = %request.method(), + uri = %request.uri(), + "rejecting code-mode host OTEL trace websocket request with Origin header" + ); + Err(StatusCode::FORBIDDEN) + } else { + Ok(next.run(request).await) + } +} + +fn parse_listen_url(listen_url: &str, flag: &str, scheme: &str) -> Result { + listen_url + .strip_prefix(scheme) + .ok_or_else(|| anyhow::anyhow!("unsupported {flag} URL '{listen_url}'"))? + .parse::() + .with_context(|| format!("invalid {flag} URL '{listen_url}'; expected '{scheme}IP:PORT'")) +} + +async fn receive_otlp_trace_batch( + State(sender): State>>, + body: Bytes, +) -> impl IntoResponse { + let _ = sender.send(body.to_vec()); + ([(CONTENT_TYPE, "application/json")], "{}") +} + +async fn trace_websocket_upgrade_handler( + websocket: WebSocketUpgrade, + State(sender): State>>, +) -> impl IntoResponse { + websocket + .max_frame_size(MAX_TRACE_WEBSOCKET_FRAME_BYTES) + .max_message_size(MAX_TRACE_WEBSOCKET_FRAME_BYTES) + .on_upgrade(move |stream| stream_trace_batches(stream, sender)) +} + +async fn stream_trace_batches(stream: WebSocket, sender: broadcast::Sender>) { + let (mut writer, mut reader) = stream.split(); + let mut receiver = sender.subscribe(); + info!("code-mode host OTEL trace websocket client connected"); + loop { + tokio::select! { + incoming = reader.next() => match incoming { + Some(Ok(Message::Close(_))) | None => break, + Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Binary(_) | Message::Text(_))) => {} + Some(Err(error)) => { + warn!(%error, "code-mode host OTEL trace websocket reader failed"); + break; + } + }, + outgoing = receiver.recv() => match outgoing { + Ok(batch) => { + if let Err(error) = writer.send(Message::Binary(batch.into())).await { + warn!(%error, "code-mode host OTEL trace websocket writer failed"); + break; + } + } + Err(broadcast::error::RecvError::Lagged(skipped)) => { + warn!(skipped, "code-mode host OTEL trace websocket client lagged"); + } + Err(broadcast::error::RecvError::Closed) => break, + }, + } + } +}