diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 539543978b..2eeb91bd66 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -3,6 +3,7 @@ use crate::common::ResponseEvent; use crate::common::ResponseStream; use crate::common::ResponsesWsRequest; use crate::common::SafetyBufferingTreatment; +use crate::common::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY; use crate::error::ApiError; use crate::provider::Provider; use crate::rate_limits::parse_rate_limit_event; @@ -158,6 +159,22 @@ const WEBSOCKET_CONNECTION_LIMIT_REACHED_CODE: &str = "websocket_connection_limi const WEBSOCKET_CONNECTION_LIMIT_REACHED_MESSAGE: &str = "Responses websocket connection limit reached (60 minutes). Create a new websocket connection to continue."; const RESPONSES_WEBSOCKET_TIMING_KIND: &str = "responsesapi.websocket_timing"; const RESPONSES_WEBSOCKET_TIMING_EVENT_TARGET: &str = "codex_api::responses_websocket_timing"; +const SESSION_ID_CLIENT_METADATA_KEY: &str = "session_id"; +const THREAD_ID_CLIENT_METADATA_KEY: &str = "thread_id"; +const TURN_ID_CLIENT_METADATA_KEY: &str = "turn_id"; +const WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY: &str = "x-codex-ws-stream-request-start-ms"; + +struct ResponsesWebsocketTimingLogContext { + model: String, + session_id: Option, + thread_id: Option, + turn_id: Option, + traceparent: Option, + previous_response_id: Option, + request_start_ms: Option, + warmup: bool, + connection_reused: bool, +} pub struct ResponsesWebsocketConnection { stream: Arc>>, @@ -225,6 +242,31 @@ impl ResponsesWebsocketConnection { let models_etag = self.models_etag.clone(); let server_model = self.server_model.clone(); let telemetry = self.telemetry.clone(); + let ResponsesWsRequest::ResponseCreate(ws_request) = &request; + let client_metadata = ws_request.client_metadata.as_ref(); + let timing_log_context = ResponsesWebsocketTimingLogContext { + model: ws_request.model.clone(), + session_id: client_metadata + .and_then(|metadata| metadata.get(SESSION_ID_CLIENT_METADATA_KEY)) + .cloned(), + thread_id: client_metadata + .and_then(|metadata| metadata.get(THREAD_ID_CLIENT_METADATA_KEY)) + .cloned(), + turn_id: client_metadata + .and_then(|metadata| metadata.get(TURN_ID_CLIENT_METADATA_KEY)) + .cloned(), + traceparent: client_metadata + .and_then(|metadata| { + metadata.get(WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY) + }) + .cloned(), + previous_response_id: ws_request.previous_response_id.clone(), + request_start_ms: client_metadata + .and_then(|metadata| metadata.get(WS_STREAM_REQUEST_START_MS_CLIENT_METADATA_KEY)) + .cloned(), + warmup: ws_request.generate == Some(false), + connection_reused, + }; let request_text = serialize_websocket_request(&request)?; let current_span = Span::current(); @@ -262,8 +304,8 @@ impl ResponsesWebsocketConnection { request_text, idle_timeout, telemetry, - connection_reused, turn_state.as_deref(), + &timing_log_context, ) .await }; @@ -627,8 +669,8 @@ async fn run_websocket_response_stream( request_text: String, idle_timeout: Duration, telemetry: Option>, - connection_reused: bool, turn_state: Option<&OnceLock>, + timing_log_context: &ResponsesWebsocketTimingLogContext, ) -> Result<(), ApiError> { let mut last_server_model: Option = None; let mut safety_buffering_treatment = SafetyBufferingTreatment::default(); @@ -637,7 +679,7 @@ async fn run_websocket_response_stream( request_text, idle_timeout, telemetry.as_ref(), - connection_reused, + timing_log_context.connection_reused, ) .await?; @@ -680,7 +722,11 @@ async fn run_websocket_response_stream( continue; } }; - emit_responses_websocket_timing_event(event.kind(), text.as_str()); + emit_responses_websocket_timing_event( + event.kind(), + text.as_str(), + timing_log_context, + ); if let Some(response_turn_state) = event.turn_state() && let Some(turn_state) = turn_state { @@ -764,17 +810,30 @@ async fn run_websocket_response_stream( Ok(()) } -fn emit_responses_websocket_timing_event(kind: &str, payload: &str) { +fn emit_responses_websocket_timing_event( + kind: &str, + payload: &str, + context: &ResponsesWebsocketTimingLogContext, +) { if kind != RESPONSES_WEBSOCKET_TIMING_KIND { return; } - // This full payload is excluded from always-on sinks. Opt in to the request span's event with + // This full payload is excluded from always-on sinks. Opt in with // `RUST_LOG='codex_api::responses_websocket_timing=trace'`. tracing::event!( name: RESPONSES_WEBSOCKET_TIMING_KIND, target: RESPONSES_WEBSOCKET_TIMING_EVENT_TARGET, tracing::Level::TRACE, + model = context.model.as_str(), + session_id = context.session_id.as_deref().unwrap_or_default(), + thread_id = context.thread_id.as_deref().unwrap_or_default(), + turn_id = context.turn_id.as_deref().unwrap_or_default(), + traceparent = context.traceparent.as_deref().unwrap_or_default(), + previous_response_id = context.previous_response_id.as_deref().unwrap_or_default(), + request_start_ms = context.request_start_ms.as_deref().unwrap_or_default(), + warmup = context.warmup, + connection_reused = context.connection_reused, payload, "responses websocket timing" ); diff --git a/codex-rs/feedback/src/lib.rs b/codex-rs/feedback/src/lib.rs index cd97349c48..a8864690bd 100644 --- a/codex-rs/feedback/src/lib.rs +++ b/codex-rs/feedback/src/lib.rs @@ -18,6 +18,7 @@ use codex_protocol::protocol::SessionSource; use tracing::Event; use tracing::Level; use tracing::field::Visit; +use tracing::level_filters::LevelFilter; use tracing_subscriber::Layer; use tracing_subscriber::filter::Targets; use tracing_subscriber::fmt::writer::MakeWriter; @@ -205,7 +206,11 @@ impl CodexFeedback { .with_target(false) // Capture everything, regardless of the caller's `RUST_LOG`, so feedback includes the // full trace when the user uploads a report. - .with_filter(Targets::new().with_default(Level::TRACE)) + .with_filter( + Targets::new() + .with_default(Level::TRACE) + .with_target("codex_api::responses_websocket_timing", LevelFilter::OFF), + ) } /// Returns a [`tracing_subscriber`] layer that collects structured metadata for feedback. @@ -722,6 +727,21 @@ mod tests { pretty_assertions::assert_eq!(std::str::from_utf8(snap.as_bytes()).unwrap(), "cdefghij"); } + #[test] + fn logger_layer_excludes_responses_websocket_timing_payloads() { + let fb = CodexFeedback::new(); + let _guard = tracing_subscriber::registry() + .with(fb.logger_layer()) + .set_default(); + + tracing::trace!(target: "codex_api::responses_websocket_timing", payload = "secret"); + tracing::trace!(target: "codex_feedback_test", "retained"); + + let logs = String::from_utf8(fb.snapshot(/*session_id*/ None).bytes).unwrap(); + assert!(!logs.contains("secret")); + assert!(logs.contains("retained")); + } + #[test] fn metadata_layer_records_tags_from_feedback_target() { let fb = CodexFeedback::new();