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/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,
-            },
-        }
-    }
-}