From 9969043b95339ec008b2f2f317f664d763fa4825 Mon Sep 17 00:00:00 2001 From: "Adam Perry @ OpenAI" Date: Tue, 1 Sep 2026 16:57:58 +0000 Subject: [PATCH] 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 | 12 ++ codex-rs/Cargo.toml | 2 + codex-rs/code-mode-host/Cargo.toml | 1 + codex-rs/code-mode-host/src/lib.rs | 6 - codex-rs/code-mode-host/src/main.rs | 47 ++--- .../code-mode-host/src/trace_transport.rs | 161 ---------------- codex-rs/otel-trace-websocket/BUILD.bazel | 6 + codex-rs/otel-trace-websocket/Cargo.toml | 19 ++ codex-rs/otel-trace-websocket/src/lib.rs | 179 ++++++++++++++++++ 9 files changed, 234 insertions(+), 199 deletions(-) delete mode 100644 codex-rs/code-mode-host/src/trace_transport.rs create mode 100644 codex-rs/otel-trace-websocket/BUILD.bazel create mode 100644 codex-rs/otel-trace-websocket/Cargo.toml create mode 100644 codex-rs/otel-trace-websocket/src/lib.rs diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index ec12db9d93..298a2845bf 100644 --- a/codex-rs/Cargo.lock +++ b/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" diff --git a/codex-rs/Cargo.toml b/codex-rs/Cargo.toml index 23e9bdbb10..f7aa05a842 100644 --- a/codex-rs/Cargo.toml +++ b/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" } diff --git a/codex-rs/code-mode-host/Cargo.toml b/codex-rs/code-mode-host/Cargo.toml index a064dd5759..194d05dff6 100644 --- a/codex-rs/code-mode-host/Cargo.toml +++ b/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" diff --git a/codex-rs/code-mode-host/src/lib.rs b/codex-rs/code-mode-host/src/lib.rs index 50dd409b70..b147983415 100644 --- a/codex-rs/code-mode-host/src/lib.rs +++ b/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; diff --git a/codex-rs/code-mode-host/src/main.rs b/codex-rs/code-mode-host/src/main.rs index 81ffb3ae78..d7be1f60cb 100644 --- a/codex-rs/code-mode-host/src/main.rs +++ b/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 } diff --git a/codex-rs/code-mode-host/src/trace_transport.rs b/codex-rs/code-mode-host/src/trace_transport.rs deleted file mode 100644 index 971b289a62..0000000000 --- a/codex-rs/code-mode-host/src/trace_transport.rs +++ /dev/null @@ -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> { - 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, - }, - } - } -} diff --git a/codex-rs/otel-trace-websocket/BUILD.bazel b/codex-rs/otel-trace-websocket/BUILD.bazel new file mode 100644 index 0000000000..ada1b13f71 --- /dev/null +++ b/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", +) diff --git a/codex-rs/otel-trace-websocket/Cargo.toml b/codex-rs/otel-trace-websocket/Cargo.toml new file mode 100644 index 0000000000..538666ec72 --- /dev/null +++ b/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 } diff --git a/codex-rs/otel-trace-websocket/src/lib.rs b/codex-rs/otel-trace-websocket/src/lib.rs new file mode 100644 index 0000000000..33afa67db7 --- /dev/null +++ b/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>, +} + +impl TraceWebSocket { + /// Binds both listeners before returning, so startup reports address conflicts. + pub async fn start(listen_url: &str) -> Result { + let bind_address = listen_url + .strip_prefix("ws://") + .ok_or_else(|| anyhow::anyhow!("unsupported --otel-trace-listen URL '{listen_url}'"))? + .parse::() + .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, + next: Next, +) -> std::result::Result { + 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>>, + 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!("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, + }, + } + } +}