mirror of
https://github.com/openai/codex.git
synced 2026-04-28 16:45:54 +00:00
feat(core, tracing): create turn spans over websockets (#14632)
## Description Dependent on: - [responsesapi] https://github.com/openai/openai/pull/760991 - [codex-backend] https://github.com/openai/openai/pull/760985 `codex app-server -> codex-backend -> responsesapi` now reuses a persistent websocket connection across many turns. This PR updates tracing when using websockets so that each `response.create` websocket request propagates the current tracing context, so we can get a holistic end-to-end trace for each turn. Tracing is propagated via special keys (`ws_request_header_traceparent`, `ws_request_header_tracestate`) set in the `client_metadata` param in Responses API. Currently tracing on websockets is a bit broken because we only set tracing context on ws connection time, so it's detached from a `turn/start` request.
This commit is contained in:
@@ -1,4 +1,6 @@
|
||||
#![allow(clippy::expect_used, clippy::unwrap_used)]
|
||||
use codex_api::WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY;
|
||||
use codex_api::WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY;
|
||||
use codex_core::CodexAuth;
|
||||
use codex_core::ModelClient;
|
||||
use codex_core::ModelClientSession;
|
||||
@@ -10,6 +12,7 @@ use codex_core::X_RESPONSESAPI_INCLUDE_TIMING_METRICS_HEADER;
|
||||
use codex_core::features::Feature;
|
||||
use codex_otel::SessionTelemetry;
|
||||
use codex_otel::TelemetryAuthMode;
|
||||
use codex_otel::current_span_w3c_trace_context;
|
||||
use codex_otel::metrics::MetricsClient;
|
||||
use codex_otel::metrics::MetricsConfig;
|
||||
use codex_protocol::ThreadId;
|
||||
@@ -24,6 +27,7 @@ use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig;
|
||||
use codex_protocol::protocol::EventMsg;
|
||||
use codex_protocol::protocol::Op;
|
||||
use codex_protocol::protocol::SessionSource;
|
||||
use codex_protocol::protocol::W3cTraceContext;
|
||||
use codex_protocol::user_input::UserInput;
|
||||
use core_test_support::load_default_config_for_test;
|
||||
use core_test_support::responses::WebSocketConnectionConfig;
|
||||
@@ -35,6 +39,7 @@ use core_test_support::responses::start_websocket_server;
|
||||
use core_test_support::responses::start_websocket_server_with_headers;
|
||||
use core_test_support::skip_if_no_network;
|
||||
use core_test_support::test_codex::test_codex;
|
||||
use core_test_support::tracing::install_test_tracing;
|
||||
use core_test_support::wait_for_event;
|
||||
use futures::StreamExt;
|
||||
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
|
||||
@@ -43,6 +48,7 @@ use serde_json::json;
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
use tempfile::TempDir;
|
||||
use tracing::Instrument;
|
||||
use tracing_test::traced_test;
|
||||
|
||||
const MODEL: &str = "gpt-5.2-codex";
|
||||
@@ -50,6 +56,32 @@ const OPENAI_BETA_HEADER: &str = "OpenAI-Beta";
|
||||
const WS_V2_BETA_HEADER_VALUE: &str = "responses_websockets=2026-02-06";
|
||||
const X_CLIENT_REQUEST_ID_HEADER: &str = "x-client-request-id";
|
||||
|
||||
fn assert_request_trace_matches(body: &serde_json::Value, expected_trace: &W3cTraceContext) {
|
||||
let client_metadata = body["client_metadata"]
|
||||
.as_object()
|
||||
.expect("missing client_metadata payload");
|
||||
let actual_traceparent = client_metadata
|
||||
.get(WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.expect("missing traceparent");
|
||||
let expected_traceparent = expected_trace
|
||||
.traceparent
|
||||
.as_deref()
|
||||
.expect("missing expected traceparent");
|
||||
|
||||
assert_eq!(actual_traceparent, expected_traceparent);
|
||||
assert_eq!(
|
||||
client_metadata
|
||||
.get(WS_REQUEST_HEADER_TRACESTATE_CLIENT_METADATA_KEY)
|
||||
.and_then(serde_json::Value::as_str),
|
||||
expected_trace.tracestate.as_deref()
|
||||
);
|
||||
assert!(
|
||||
body.get("trace").is_none(),
|
||||
"top-level trace should not be sent"
|
||||
);
|
||||
}
|
||||
|
||||
struct WebsocketTestHarness {
|
||||
_codex_home: TempDir,
|
||||
client: ModelClient,
|
||||
@@ -119,6 +151,112 @@ async fn responses_websocket_streams_without_feature_flag_when_provider_supports
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn responses_websocket_reuses_connection_with_per_turn_trace_payloads() {
|
||||
skip_if_no_network!();
|
||||
|
||||
let _trace_test_context = install_test_tracing("client-websocket-test");
|
||||
|
||||
let server = start_websocket_server(vec![vec![
|
||||
vec![ev_response_created("resp-1"), ev_completed("resp-1")],
|
||||
vec![ev_response_created("resp-2"), ev_completed("resp-2")],
|
||||
]])
|
||||
.await;
|
||||
|
||||
let harness = websocket_harness(&server).await;
|
||||
let prompt_one = prompt_with_input(vec![message_item("hello")]);
|
||||
let prompt_two = prompt_with_input(vec![message_item("again")]);
|
||||
|
||||
let first_trace = {
|
||||
let mut client_session = harness.client.new_session();
|
||||
async {
|
||||
let expected_trace =
|
||||
current_span_w3c_trace_context().expect("current span should have trace context");
|
||||
stream_until_complete(&mut client_session, &harness, &prompt_one).await;
|
||||
expected_trace
|
||||
}
|
||||
.instrument(tracing::info_span!("client.websocket.turn_one"))
|
||||
.await
|
||||
};
|
||||
|
||||
let second_trace = {
|
||||
let mut client_session = harness.client.new_session();
|
||||
async {
|
||||
let expected_trace =
|
||||
current_span_w3c_trace_context().expect("current span should have trace context");
|
||||
stream_until_complete(&mut client_session, &harness, &prompt_two).await;
|
||||
expected_trace
|
||||
}
|
||||
.instrument(tracing::info_span!("client.websocket.turn_two"))
|
||||
.await
|
||||
};
|
||||
|
||||
assert_eq!(server.handshakes().len(), 1);
|
||||
let connection = server.single_connection();
|
||||
assert_eq!(connection.len(), 2);
|
||||
|
||||
let first_request = connection
|
||||
.first()
|
||||
.expect("missing first request")
|
||||
.body_json();
|
||||
let second_request = connection
|
||||
.get(1)
|
||||
.expect("missing second request")
|
||||
.body_json();
|
||||
assert_request_trace_matches(&first_request, &first_trace);
|
||||
assert_request_trace_matches(&second_request, &second_trace);
|
||||
|
||||
let first_traceparent = first_request["client_metadata"]
|
||||
[WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY]
|
||||
.as_str()
|
||||
.expect("missing first traceparent");
|
||||
let second_traceparent = second_request["client_metadata"]
|
||||
[WS_REQUEST_HEADER_TRACEPARENT_CLIENT_METADATA_KEY]
|
||||
.as_str()
|
||||
.expect("missing second traceparent");
|
||||
assert_ne!(first_traceparent, second_traceparent);
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn responses_websocket_preconnect_does_not_replace_turn_trace_payload() {
|
||||
skip_if_no_network!();
|
||||
|
||||
let _trace_test_context = install_test_tracing("client-websocket-test");
|
||||
|
||||
let server = start_websocket_server(vec![vec![vec![
|
||||
ev_response_created("resp-1"),
|
||||
ev_completed("resp-1"),
|
||||
]]])
|
||||
.await;
|
||||
|
||||
let harness = websocket_harness(&server).await;
|
||||
let mut client_session = harness.client.new_session();
|
||||
client_session
|
||||
.preconnect_websocket(&harness.session_telemetry, &harness.model_info)
|
||||
.await
|
||||
.expect("websocket preconnect failed");
|
||||
let prompt = prompt_with_input(vec![message_item("hello")]);
|
||||
|
||||
let expected_trace = async {
|
||||
let expected_trace =
|
||||
current_span_w3c_trace_context().expect("current span should have trace context");
|
||||
stream_until_complete(&mut client_session, &harness, &prompt).await;
|
||||
expected_trace
|
||||
}
|
||||
.instrument(tracing::info_span!("client.websocket.request"))
|
||||
.await;
|
||||
|
||||
assert_eq!(server.handshakes().len(), 1);
|
||||
let connection = server.single_connection();
|
||||
assert_eq!(connection.len(), 1);
|
||||
let request = connection.first().expect("missing request").body_json();
|
||||
assert_request_trace_matches(&request, &expected_trace);
|
||||
|
||||
server.shutdown().await;
|
||||
}
|
||||
|
||||
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
|
||||
async fn responses_websocket_preconnect_reuses_connection() {
|
||||
skip_if_no_network!();
|
||||
|
||||
Reference in New Issue
Block a user