Extract OTEL trace WebSocket into a reusable crate (#42102)

## What changed

- Add `codex-otel-trace-websocket` with a `TraceWebSocket` API that owns the
  loopback OTLP receiver and WebSocket listener.
- Bind both listeners during startup, expose the exporter and bound listener
  addresses, and surface listener failures through a single lifecycle method.
- Update `codex-code-mode-host` to use the new crate and shut down the bridge
  after flushing its trace provider.

GitOrigin-RevId: ec5ca4c4369b6b3b1232875c2699730eb748ab1c

codex-rs/Cargo.lock

@@ -2627,6 +2627,7 @@ dependencies = [
  "codex-code-mode-protocol",
  "codex-code-mode-runtime",
  "codex-otel",
+ "codex-otel-trace-websocket",
  "codex-protocol",
  "codex-utils-cargo-bin",
  "futures",
@@ -3985,6 +3986,17 @@ dependencies = [
  "tracing-subscriber",
 ]
 
+[[package]]
+name = "codex-otel-trace-websocket"
+version = "0.0.0"
+dependencies = [
+ "anyhow",
+ "axum",
+ "futures",
+ "tokio",
+ "tracing",
+]
+
 [[package]]
 name = "codex-plugin"
 version = "0.0.0"

codex-rs/Cargo.toml

@@ -98,6 +98,7 @@ members = [
     "sandboxing",
     "stdio-to-uds",
     "otel",
+    "otel-trace-websocket",
     "tui",
     "tools",
     "v8-poc",
@@ -237,6 +238,7 @@ codex-models-manager = { path = "models-manager" }
 codex-network-proxy = { path = "network-proxy" }
 codex-ollama = { path = "ollama" }
 codex-otel = { path = "otel" }
+codex-otel-trace-websocket = { path = "otel-trace-websocket" }
 codex-plugin = { path = "plugin" }
 codex-model-provider = { path = "model-provider" }
 codex-process-hardening = { path = "process-hardening" }

codex-rs/code-mode-host/Cargo.toml

@@ -23,6 +23,7 @@ clap = { workspace = true, features = ["derive"] }
 codex-code-mode-protocol = { workspace = true }
 codex-code-mode-runtime = { workspace = true }
 codex-otel = { workspace = true }
+codex-otel-trace-websocket = { workspace = true }
 codex-protocol = { workspace = true }
 futures = { workspace = true }
 prost = "0.14.3"

codex-rs/code-mode-host/src/lib.rs

@@ -49,14 +49,8 @@ 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;

codex-rs/code-mode-host/src/main.rs

@@ -9,6 +9,7 @@ use codex_otel::OtelExporter;
 use codex_otel::OtelHttpProtocol;
 use codex_otel::OtelProvider;
 use codex_otel::OtelSettings;
+use codex_otel_trace_websocket::TraceWebSocket;
 use tracing_subscriber::Layer;
 use tracing_subscriber::layer::SubscriberExt;
 use tracing_subscriber::util::SubscriberInitExt;
@@ -38,22 +39,14 @@ struct Cli {
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> anyhow::Result<()> {
     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,
-        ))
+    let mut trace_transport = if let Some(trace_listen) = cli.otel_trace_listen.as_deref() {
+        Some(TraceWebSocket::start(trace_listen).await?)
     } else {
         None
     };
     let trace_exporter_endpoint = trace_transport
         .as_ref()
-        .map(|(_, _, _, endpoint)| endpoint.as_str())
+        .map(TraceWebSocket::exporter_endpoint)
         .or(cli.otel_trace_exporter.as_deref());
     let otel = trace_exporter_endpoint
         .map(build_trace_provider)
@@ -68,6 +61,10 @@ async fn main() -> anyhow::Result<()> {
         )
         .with(otel_layer)
         .init();
+    if let Some(trace_transport) = trace_transport.as_ref() {
+        let listen_addr = trace_transport.listen_addr();
+        tracing::info!("codex-code-mode-host OTEL trace websocket listening on ws://{listen_addr}");
+    }
     tracing::info_span!(
         "code_mode_host.startup",
         otel.name = "code_mode_host.startup"
@@ -75,33 +72,19 @@ async fn main() -> anyhow::Result<()> {
     .in_scope(|| {});
 
     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),
+    let result = match trace_transport.as_mut() {
+        Some(trace_transport) => tokio::select! {
+            result = main_transport => result,
+            result = trace_transport.wait_for_failure() => result,
+        },
+        None => main_transport.await,
     };
     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();
-    }
+    drop(trace_transport);
     result
 }
 

codex-rs/code-mode-host/src/trace_transport.rs

@@ -1,161 +0,0 @@
-//! 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,
-            },
-        }
-    }
-}

codex-rs/otel-trace-websocket/BUILD.bazel

@@ -0,0 +1,6 @@
+load("//:defs.bzl", "codex_rust_crate")
+
+codex_rust_crate(
+    name = "otel-trace-websocket",
+    crate_name = "codex_otel_trace_websocket",
+)

codex-rs/otel-trace-websocket/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "codex-otel-trace-websocket"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+[lib]
+test = false
+doctest = false
+
+[lints]
+workspace = true
+
+[dependencies]
+anyhow = { workspace = true }
+axum = { workspace = true, features = ["http1", "tokio", "ws"] }
+futures = { workspace = true }
+tokio = { workspace = true, features = ["macros", "net", "rt", "sync"] }
+tracing = { workspace = true }

codex-rs/otel-trace-websocket/src/lib.rs

@@ -0,0 +1,179 @@
+//! Forwards loopback OTLP trace batches to a separate WebSocket listener.
+//!
+//! Forwarding is best effort: disconnected or lagging clients lose batches.
+//! Callers must flush their trace provider before dropping the bridge.
+//! Dropping the bridge aborts its listeners without draining queued batches.
+
+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 tokio::task::JoinSet;
+use tracing::info;
+use tracing::warn;
+
+const TRACE_CHANNEL_CAPACITY: usize = 1_024;
+const MAX_TRACE_WEBSOCKET_FRAME_BYTES: usize = 4 * 1024 * 1024;
+
+/// Owns the loopback OTLP receiver and trace WebSocket listener tasks.
+pub struct TraceWebSocket {
+    exporter_endpoint: String,
+    listen_addr: SocketAddr,
+    tasks: JoinSet<Result<()>>,
+}
+
+impl TraceWebSocket {
+    /// Binds both listeners before returning, so startup reports address conflicts.
+    pub async fn start(listen_url: &str) -> Result<Self> {
+        let bind_address = listen_url
+            .strip_prefix("ws://")
+            .ok_or_else(|| anyhow::anyhow!("unsupported --otel-trace-listen URL '{listen_url}'"))?
+            .parse::<SocketAddr>()
+            .with_context(|| {
+                format!("invalid --otel-trace-listen URL '{listen_url}'; expected 'ws://IP:PORT'")
+            })?;
+        let websocket_listener = TcpListener::bind(bind_address)
+            .await
+            .with_context(|| format!("failed to bind OTEL trace websocket to {bind_address}"))?;
+        let listen_addr = websocket_listener
+            .local_addr()
+            .context("failed to read OTEL trace websocket listen address")?;
+        let otlp_listener = TcpListener::bind("127.0.0.1:0")
+            .await
+            .context("failed to bind loopback OTLP trace receiver")?;
+        let otlp_addr = otlp_listener
+            .local_addr()
+            .context("failed to read OTLP trace receiver listen address")?;
+        let exporter_endpoint = format!("http://{otlp_addr}/v1/traces");
+        let (sender, _) = broadcast::channel(TRACE_CHANNEL_CAPACITY);
+        let otlp_router = Router::new()
+            .route("/v1/traces", post(receive_otlp_trace_batch))
+            .with_state(sender.clone());
+        let websocket_router = Router::new()
+            .route("/", any(trace_websocket_upgrade_handler))
+            .layer(middleware::from_fn(reject_requests_with_origin_header))
+            .with_state(sender);
+        let mut tasks = JoinSet::new();
+        tasks.spawn(async move {
+            axum::serve(otlp_listener, otlp_router)
+                .await
+                .context("OTLP trace receiver failed")
+        });
+        tasks.spawn(async move {
+            axum::serve(websocket_listener, websocket_router)
+                .await
+                .context("OTEL trace websocket listener failed")
+        });
+        info!("OTLP trace receiver listening on http://{otlp_addr}");
+        info!("OTEL trace websocket listening on ws://{listen_addr}");
+        Ok(Self {
+            exporter_endpoint,
+            listen_addr,
+            tasks,
+        })
+    }
+
+    /// Returns the local OTLP/HTTP endpoint to configure with JSON trace export.
+    pub fn exporter_endpoint(&self) -> &str {
+        &self.exporter_endpoint
+    }
+
+    /// Returns the bound WebSocket address, including an assigned ephemeral port.
+    pub fn listen_addr(&self) -> SocketAddr {
+        self.listen_addr
+    }
+
+    /// Waits for either listener to stop, for use alongside the main server future.
+    pub async fn wait_for_failure(&mut self) -> Result<()> {
+        self.tasks
+            .join_next()
+            .await
+            .context("OTEL trace websocket server tasks stopped")?
+            .context("OTEL trace websocket server task 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 OTEL trace websocket request with Origin header"
+        );
+        Err(StatusCode::FORBIDDEN)
+    } else {
+        Ok(next.run(request).await)
+    }
+}
+
+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!("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, "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, "OTEL trace websocket writer failed");
+                        break;
+                    }
+                }
+                Err(broadcast::error::RecvError::Lagged(skipped)) => {
+                    warn!(skipped, "OTEL trace websocket client lagged");
+                }
+                Err(broadcast::error::RecvError::Closed) => break,
+            },
+        }
+    }
+}