codex: address PR review feedback (#31688)

This commit is contained in:
Tom Wiltzius
2026-07-09 10:12:37 -07:00
parent 709bc18391
commit 34abe798b6
2 changed files with 86 additions and 7 deletions
@@ -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<String>,
thread_id: Option<String>,
turn_id: Option<String>,
traceparent: Option<String>,
previous_response_id: Option<String>,
request_start_ms: Option<String>,
warmup: bool,
connection_reused: bool,
}
pub struct ResponsesWebsocketConnection {
stream: Arc<Mutex<Option<WsStream>>>,
@@ -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<Arc<dyn WebsocketTelemetry>>,
connection_reused: bool,
turn_state: Option<&OnceLock<String>>,
timing_log_context: &ResponsesWebsocketTimingLogContext,
) -> Result<(), ApiError> {
let mut last_server_model: Option<String> = 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"
);
+21 -1
View File
@@ -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();