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

@@ -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",

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"

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,

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 {

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<String>,
+
+    /// 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<String>,
 }
 
 #[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> {
+    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")
 }

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<Vec<u8>> {
+    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<Vec<u8>>,
+) -> 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<Vec<u8>>,
+) -> 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<Body>,
+    next: Next,
+) -> std::result::Result<Response, StatusCode> {
+    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<SocketAddr> {
+    listen_url
+        .strip_prefix(scheme)
+        .ok_or_else(|| anyhow::anyhow!("unsupported {flag} URL '{listen_url}'"))?
+        .parse::<SocketAddr>()
+        .with_context(|| format!("invalid {flag} URL '{listen_url}'; expected '{scheme}IP:PORT'"))
+}
+
+async fn receive_otlp_trace_batch(
+    State(sender): State<broadcast::Sender<Vec<u8>>>,
+    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<broadcast::Sender<Vec<u8>>>,
+) -> 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<Vec<u8>>) {
+    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,
+            },
+        }
+    }
+}