From 89698ad1c303aa6f923e810c95415597e94bfe37 Mon Sep 17 00:00:00 2001 From: cassirer-openai Date: Tue, 28 Apr 2026 14:11:17 -0700 Subject: [PATCH] [rollout-trace] Include x-request-id in rollout trace. (#20066) ## Why Rollout traces need an identifier that can be used to correlate a Codex inference with upstream Responses API, proxy, and engine logs. The reduced trace model already exposed `upstream_request_id`, but it was being populated from the Responses API `response.id`. That value is useful for `previous_response_id` chaining, but it is not the transport request id that upstream systems key on. This PR separates those concepts so trace consumers can reliably answer both questions: - which Responses API response did this inference produce? - which upstream request handled it? ## Structure The change keeps the upstream request id at the same lifecycle level as the provider stream: - `codex-api` captures the `x-request-id` HTTP response header when the SSE stream is created and exposes it on `ResponseStream`. Fixture and websocket streams set the field to `None` because they do not have that HTTP response header. - `codex-core` carries that stream-level id into `InferenceTraceAttempt` when recording terminal stream outcomes. Completed, failed, cancelled, dropped-stream, and pre-response error paths all record the id when it is available. - `rollout-trace` now records both identifiers in raw terminal inference events and response payloads: `response_id` for the Responses API `response.id`, and `upstream_request_id` for `x-request-id`. - The reducer stores both fields on `InferenceCall`. It also uses `response_id` for `previous_response_id` conversation linking, which removes the old accidental dependency on the misnamed `upstream_request_id` field. - Terminal inference reduction now consumes the full terminal payload (`InferenceCompleted`, `InferenceFailed`, or `InferenceCancelled`) in one place. That keeps status, partial payloads, response ids, and upstream request ids consistent across success, failure, cancellation, and late stream-mapper events. ## Why This Shape `x-request-id` is a property of the HTTP/provider response envelope, not an SSE event. Capturing it once in `codex-api` and plumbing it through terminal trace recording avoids trying to infer the value from stream contents, and it preserves the id even when the stream fails or is cancelled after only partial output. Keeping `response_id` separate from `upstream_request_id` also makes the reduced trace model less surprising: `response_id` remains the conversation-continuation id, while `upstream_request_id` is the operational correlation id for upstream debugging. ## Validation The PR updates trace and reducer coverage for: - reading `x-request-id` from SSE response headers; - storing the true upstream request id on completed inference calls; - preserving upstream request ids for cancelled and late-cancelled inference streams; - keeping `previous_response_id` reconstruction tied to `response_id` rather than transport request ids. --- codex-rs/codex-api/src/common.rs | 2 + .../src/endpoint/responses_websocket.rs | 5 +- codex-rs/codex-api/src/sse/responses.rs | 21 ++++- codex-rs/core/src/client.rs | 88 +++++++++++++++--- codex-rs/core/src/client_tests.rs | 14 ++- codex-rs/rollout-trace/src/inference.rs | 28 +++++- .../rollout-trace/src/model/conversation.rs | 4 +- codex-rs/rollout-trace/src/raw_event.rs | 9 ++ .../src/reducer/code_cell_tests.rs | 4 + .../rollout-trace/src/reducer/conversation.rs | 2 +- .../rollout-trace/src/reducer/inference.rs | 90 ++++++++++++++----- .../src/reducer/inference_tests.rs | 10 +++ codex-rs/rollout-trace/src/reducer/mod.rs | 45 +--------- .../rollout-trace/src/reducer/test_support.rs | 2 + .../src/reducer/tool/terminal_tests.rs | 1 + codex-rs/rollout-trace/src/thread_tests.rs | 4 +- codex-rs/rollout-trace/src/writer.rs | 1 + 17 files changed, 240 insertions(+), 90 deletions(-) diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 4b150b55f..e2d2ed3c3 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -287,6 +287,8 @@ pub fn create_text_param_for_request( pub struct ResponseStream { pub rx_event: mpsc::Receiver>, + /// Server-assigned `x-request-id` response header, when present. + pub upstream_request_id: Option, } impl Stream for ResponseStream { diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 364dc1e9b..4e97ecef9 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -279,7 +279,10 @@ impl ResponsesWebsocketConnection { .instrument(current_span), ); - Ok(ResponseStream { rx_event }) + Ok(ResponseStream { + rx_event, + upstream_request_id: None, + }) } } diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index fb1742463..4e949e43a 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -28,6 +28,7 @@ use tracing::trace; const X_REASONING_INCLUDED_HEADER: &str = "x-reasoning-included"; const OPENAI_MODEL_HEADER: &str = "openai-model"; +const REQUEST_ID_HEADER: &str = "x-request-id"; const TRUSTED_ACCESS_FOR_CYBER_VERIFICATION: &str = "trusted_access_for_cyber"; /// Streams SSE events from an on-disk fixture for tests. @@ -53,7 +54,10 @@ pub fn stream_from_fixture( idle_timeout, /*telemetry*/ None, )); - Ok(ResponseStream { rx_event }) + Ok(ResponseStream { + rx_event, + upstream_request_id: None, + }) } pub fn spawn_response_stream( @@ -77,6 +81,11 @@ pub fn spawn_response_stream( .headers .get(X_REASONING_INCLUDED_HEADER) .is_some(); + let upstream_request_id = stream_response + .headers + .get(REQUEST_ID_HEADER) + .and_then(|value| value.to_str().ok()) + .map(str::to_string); if let Some(turn_state) = turn_state.as_ref() && let Some(header_value) = stream_response .headers @@ -104,7 +113,10 @@ pub fn spawn_response_stream( process_sse(stream_response.bytes, tx_event, idle_timeout, telemetry).await; }); - ResponseStream { rx_event } + ResponseStream { + rx_event, + upstream_request_id, + } } #[derive(Debug, Deserialize)] @@ -1058,8 +1070,9 @@ mod tests { } #[tokio::test] - async fn spawn_response_stream_emits_server_model_header() { + async fn spawn_response_stream_emits_header_events() { let mut headers = HeaderMap::new(); + headers.insert(REQUEST_ID_HEADER, HeaderValue::from_static("req-1")); headers.insert( OPENAI_MODEL_HEADER, HeaderValue::from_static(CYBER_RESTRICTED_MODEL_FOR_TESTS), @@ -1077,13 +1090,13 @@ mod tests { /*telemetry*/ None, /*turn_state*/ None, ); + assert_eq!(stream.upstream_request_id.as_deref(), Some("req-1")); let event = stream .rx_event .recv() .await .expect("expected server model event") .expect("expected ok event"); - match event { ResponseEvent::ServerModel(model) => { assert_eq!(model, CYBER_RESTRICTED_MODEL_FOR_TESTS); diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index fe64099cc..ba81b451a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1234,8 +1234,13 @@ impl ModelClientSession { Err(ApiError::Transport( unauthorized_transport @ TransportError::Http { status, .. }, )) if status == StatusCode::UNAUTHORIZED => { - inference_trace_attempt - .record_failed(&unauthorized_transport, /*output_items*/ &[]); + let response_debug_context = + extract_response_debug_context(&unauthorized_transport); + inference_trace_attempt.record_failed( + &unauthorized_transport, + response_debug_context.request_id.as_deref(), + /*output_items*/ &[], + ); pending_retry = PendingUnauthorizedRetry::from_recovery( handle_unauthorized( unauthorized_transport, @@ -1247,8 +1252,14 @@ impl ModelClientSession { continue; } Err(err) => { + let response_debug_context = + extract_response_debug_context_from_api_error(&err); let err = map_api_error(err); - inference_trace_attempt.record_failed(&err, /*output_items*/ &[]); + inference_trace_attempt.record_failed( + &err, + response_debug_context.request_id.as_deref(), + /*output_items*/ &[], + ); return Err(err); } } @@ -1374,8 +1385,14 @@ impl ModelClientSession { .stream_request(ws_request, self.websocket_session.connection_reused()) .await .map_err(|err| { + let response_debug_context = + extract_response_debug_context_from_api_error(&err); let err = map_api_error(err); - inference_trace_attempt.record_failed(&err, /*output_items*/ &[]); + inference_trace_attempt.record_failed( + &err, + response_debug_context.request_id.as_deref(), + /*output_items*/ &[], + ); err })?; let (stream, last_request_rx) = map_response_stream( @@ -1636,7 +1653,29 @@ fn parent_thread_id_header_value(session_source: &SessionSource) -> Option( +fn map_response_stream( + api_stream: codex_api::ResponseStream, + session_telemetry: SessionTelemetry, + inference_trace_attempt: InferenceTraceAttempt, +) -> (ResponseStream, oneshot::Receiver) { + let codex_api::ResponseStream { + rx_event, + upstream_request_id, + } = api_stream; + let api_stream = codex_api::ResponseStream { + rx_event, + upstream_request_id: None, + }; + map_response_events( + upstream_request_id, + api_stream, + session_telemetry, + inference_trace_attempt, + ) +} + +fn map_response_events( + upstream_request_id: Option, api_stream: S, session_telemetry: SessionTelemetry, inference_trace_attempt: InferenceTraceAttempt, @@ -1658,10 +1697,15 @@ where let mut tx_last_response = Some(tx_last_response); let mut items_added: Vec = Vec::new(); let mut api_stream = api_stream; + let upstream_request_id = upstream_request_id.as_deref(); loop { let event = tokio::select! { _ = consumer_dropped.cancelled() => { - inference_trace_attempt.record_cancelled(STREAM_DROPPED_REASON, &items_added); + inference_trace_attempt.record_cancelled( + STREAM_DROPPED_REASON, + upstream_request_id, + &items_added, + ); return; } event = api_stream.next() => event, @@ -1677,8 +1721,11 @@ where .await .is_err() { - inference_trace_attempt - .record_cancelled(STREAM_DROPPED_REASON, &items_added); + inference_trace_attempt.record_cancelled( + STREAM_DROPPED_REASON, + upstream_request_id, + &items_added, + ); return; } } @@ -1698,6 +1745,7 @@ where } inference_trace_attempt.record_completed( &response_id, + upstream_request_id, &token_usage, &items_added, ); @@ -1721,14 +1769,25 @@ where } Ok(event) => { if tx_event.send(Ok(event)).await.is_err() { - inference_trace_attempt - .record_cancelled(STREAM_DROPPED_REASON, &items_added); + inference_trace_attempt.record_cancelled( + STREAM_DROPPED_REASON, + upstream_request_id, + &items_added, + ); return; } } Err(err) => { + let response_debug_context = + extract_response_debug_context_from_api_error(&err); + let upstream_request_id = + upstream_request_id.or(response_debug_context.request_id.as_deref()); let mapped = map_api_error(err); - inference_trace_attempt.record_failed(&mapped, &items_added); + inference_trace_attempt.record_failed( + &mapped, + upstream_request_id, + &items_added, + ); if !logged_error { session_telemetry.see_event_completed_failed(&mapped); logged_error = true; @@ -1739,8 +1798,11 @@ where } } } - inference_trace_attempt - .record_failed("stream closed before response.completed", &items_added); + inference_trace_attempt.record_failed( + "stream closed before response.completed", + upstream_request_id, + &items_added, + ); }); ( diff --git a/codex-rs/core/src/client_tests.rs b/codex-rs/core/src/client_tests.rs index f7d263c0e..e56500ba5 100644 --- a/codex-rs/core/src/client_tests.rs +++ b/codex-rs/core/src/client_tests.rs @@ -282,7 +282,12 @@ async fn dropped_response_stream_traces_cancelled_partial_output() -> anyhow::Re let item = output_message("msg-1", "partial answer"); let api_stream = futures::stream::iter([Ok(ResponseEvent::OutputItemDone(item))]) .chain(futures::stream::pending()); - let (mut stream, _) = super::map_response_stream(api_stream, test_session_telemetry(), attempt); + let (mut stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + attempt, + ); let observed = stream .next() @@ -332,7 +337,12 @@ async fn dropped_backpressured_response_stream_traces_cancelled_partial_output() notify: Arc::clone(&backpressured_item_yielded), }; - let (stream, _) = super::map_response_stream(api_stream, test_session_telemetry(), attempt); + let (stream, _) = super::map_response_events( + /*upstream_request_id*/ None, + api_stream, + test_session_telemetry(), + attempt, + ); // Fill the mapper channel with non-terminal events, then yield one output // item. The mapper has observed that item and is blocked trying to send it diff --git a/codex-rs/rollout-trace/src/inference.rs b/codex-rs/rollout-trace/src/inference.rs index 3162308fb..b05f891e7 100644 --- a/codex-rs/rollout-trace/src/inference.rs +++ b/codex-rs/rollout-trace/src/inference.rs @@ -84,6 +84,7 @@ struct EnabledInferenceTraceAttempt { #[derive(Serialize)] struct TracedResponseStreamOutput<'a> { response_id: Option<&'a str>, + upstream_request_id: Option<&'a str>, token_usage: Option<&'a TokenUsage>, output_items: Vec, } @@ -174,6 +175,7 @@ impl InferenceTraceAttempt { pub fn record_completed( &self, response_id: &str, + upstream_request_id: Option<&str>, token_usage: &Option, output_items: &[ResponseItem], ) { @@ -183,6 +185,7 @@ impl InferenceTraceAttempt { let Some(response_payload) = write_response_payload_best_effort( attempt, Some(response_id), + upstream_request_id, token_usage.as_ref(), output_items, ) else { @@ -194,13 +197,19 @@ impl InferenceTraceAttempt { RawTraceEventPayload::InferenceCompleted { inference_call_id: attempt.inference_call_id.clone(), response_id: Some(response_id.to_string()), + upstream_request_id: upstream_request_id.map(str::to_string), response_payload, }, ); } /// Records pre-response and mid-stream failures. - pub fn record_failed(&self, error: impl Display, output_items: &[ResponseItem]) { + pub fn record_failed( + &self, + error: impl Display, + upstream_request_id: Option<&str>, + output_items: &[ResponseItem], + ) { let Some(attempt) = self.take_terminal_attempt() else { return; }; @@ -210,6 +219,7 @@ impl InferenceTraceAttempt { write_response_payload_best_effort( attempt, /*response_id*/ None, + upstream_request_id, /*token_usage*/ None, output_items, ) @@ -218,6 +228,7 @@ impl InferenceTraceAttempt { &attempt.context, RawTraceEventPayload::InferenceFailed { inference_call_id: attempt.inference_call_id.clone(), + upstream_request_id: upstream_request_id.map(str::to_string), error: error.to_string(), partial_response_payload, }, @@ -229,7 +240,12 @@ impl InferenceTraceAttempt { /// This happens when the turn is interrupted or when mailbox delivery /// preempts the current sampling request. Complete output items observed /// before that point are retained as partial response evidence. - pub fn record_cancelled(&self, reason: impl Display, output_items: &[ResponseItem]) { + pub fn record_cancelled( + &self, + reason: impl Display, + upstream_request_id: Option<&str>, + output_items: &[ResponseItem], + ) { let Some(attempt) = self.take_terminal_attempt() else { return; }; @@ -239,6 +255,7 @@ impl InferenceTraceAttempt { write_response_payload_best_effort( attempt, /*response_id*/ None, + upstream_request_id, /*token_usage*/ None, output_items, ) @@ -247,6 +264,7 @@ impl InferenceTraceAttempt { &attempt.context, RawTraceEventPayload::InferenceCancelled { inference_call_id: attempt.inference_call_id.clone(), + upstream_request_id: upstream_request_id.map(str::to_string), reason: reason.to_string(), partial_response_payload, }, @@ -312,11 +330,13 @@ fn write_json_payload_best_effort( fn write_response_payload_best_effort( attempt: &EnabledInferenceTraceAttempt, response_id: Option<&str>, + upstream_request_id: Option<&str>, token_usage: Option<&TokenUsage>, output_items: &[ResponseItem], ) -> Option { let response_payload = TracedResponseStreamOutput { response_id, + upstream_request_id, token_usage, output_items: output_items.iter().map(trace_response_item_json).collect(), }; @@ -387,7 +407,7 @@ mod tests { "content": [{"type": "input_text", "text": "hello"}] }], })); - attempt.record_completed("resp-1", &None, &[]); + attempt.record_completed("resp-1", Some("req-1"), &None, &[]); let rollout = replay_bundle(temp.path())?; let inference = rollout @@ -400,7 +420,7 @@ mod tests { assert_eq!(inference.thread_id, "thread-root"); assert_eq!(inference.codex_turn_id, "turn-1"); assert_eq!(inference.execution.status, ExecutionStatus::Completed); - assert_eq!(inference.upstream_request_id, Some("resp-1".to_string())); + assert_eq!(inference.upstream_request_id, Some("req-1".to_string())); assert_eq!(rollout.raw_payloads.len(), 2); Ok(()) diff --git a/codex-rs/rollout-trace/src/model/conversation.rs b/codex-rs/rollout-trace/src/model/conversation.rs index e248b72d4..1a8af0270 100644 --- a/codex-rs/rollout-trace/src/model/conversation.rs +++ b/codex-rs/rollout-trace/src/model/conversation.rs @@ -152,7 +152,9 @@ pub struct InferenceCall { pub execution: ExecutionWindow, pub model: String, pub provider_name: String, - /// Upstream request ID returned by HTTP/proxy/engine infrastructure. + /// Responses API response id, used by follow-up `previous_response_id` requests. + pub response_id: Option, + /// Request id returned by HTTP/proxy/engine infrastructure. pub upstream_request_id: Option, /// Complete ordered input snapshot sent with this request. pub request_item_ids: Vec, diff --git a/codex-rs/rollout-trace/src/raw_event.rs b/codex-rs/rollout-trace/src/raw_event.rs index ad8f436da..d60a2b6c4 100644 --- a/codex-rs/rollout-trace/src/raw_event.rs +++ b/codex-rs/rollout-trace/src/raw_event.rs @@ -101,17 +101,26 @@ pub enum RawTraceEventPayload { }, InferenceCompleted { inference_call_id: InferenceCallId, + /// Responses API `response.id`; used by `previous_response_id`. response_id: Option, + /// Provider transport request id, such as `x-request-id`. + upstream_request_id: Option, response_payload: RawPayloadRef, }, InferenceFailed { inference_call_id: InferenceCallId, + /// Provider transport request id, such as `x-request-id`, when the + /// provider returned one before the stream failed. + upstream_request_id: Option, error: String, /// Partial response payload, when stream events arrived before failure. partial_response_payload: Option, }, InferenceCancelled { inference_call_id: InferenceCallId, + /// Provider transport request id, such as `x-request-id`, when observed + /// before Codex stopped consuming the stream. + upstream_request_id: Option, /// Why Codex stopped consuming the provider stream before a terminal response event. reason: String, /// Completed output items observed before cancellation, if any. diff --git a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs index f77921352..1aabdb843 100644 --- a/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/code_cell_tests.rs @@ -64,6 +64,7 @@ fn code_cell_lifecycle_links_nested_tools_waits_and_outputs() -> anyhow::Result< writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: "inference-1".to_string(), response_id: Some("resp-1".to_string()), + upstream_request_id: None, response_payload: response, })?; writer.append_with_context( @@ -247,6 +248,7 @@ fn fast_code_cell_lifecycle_waits_for_source_item() -> anyhow::Result<()> { writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: "inference-1".to_string(), response_id: Some("resp-1".to_string()), + upstream_request_id: None, response_payload: response, })?; @@ -301,6 +303,7 @@ fn cancelled_turn_terminates_unfinished_code_cell() -> anyhow::Result<()> { writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: "inference-1".to_string(), response_id: Some("resp-1".to_string()), + upstream_request_id: None, response_payload: response, })?; writer.append_with_context( @@ -388,6 +391,7 @@ fn runtime_code_cell_ids_can_repeat_across_threads() -> anyhow::Result<()> { writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: inference_call_id.to_string(), response_id: Some(format!("resp-{thread_id}")), + upstream_request_id: None, response_payload: response, })?; writer.append_with_context( diff --git a/codex-rs/rollout-trace/src/reducer/conversation.rs b/codex-rs/rollout-trace/src/reducer/conversation.rs index 53ae60824..2502e770b 100644 --- a/codex-rs/rollout-trace/src/reducer/conversation.rs +++ b/codex-rs/rollout-trace/src/reducer/conversation.rs @@ -76,7 +76,7 @@ impl TraceReducer { .values() .find(|inference| { inference.thread_id == thread_id - && inference.upstream_request_id.as_deref() == Some(previous_response_id) + && inference.response_id.as_deref() == Some(previous_response_id) }) .map(|inference| { let mut ids = inference.request_item_ids.clone(); diff --git a/codex-rs/rollout-trace/src/reducer/inference.rs b/codex-rs/rollout-trace/src/reducer/inference.rs index c65b4236a..5becb59ca 100644 --- a/codex-rs/rollout-trace/src/reducer/inference.rs +++ b/codex-rs/rollout-trace/src/reducer/inference.rs @@ -13,6 +13,7 @@ use crate::model::InferenceCall; use crate::model::InferenceCallId; use crate::payload::RawPayloadRef; use crate::raw_event::RawEventSeq; +use crate::raw_event::RawTraceEventPayload; /// Raw inference-start fields after dispatch has stripped the common event envelope. /// @@ -91,6 +92,7 @@ impl TraceReducer { }, model: started.model, provider_name: started.provider_name, + response_id: None, upstream_request_id: None, request_item_ids, response_item_ids: Vec::new(), @@ -137,11 +139,49 @@ impl TraceReducer { &mut self, seq: RawEventSeq, wall_time_unix_ms: i64, - inference_call_id: InferenceCallId, - status: ExecutionStatus, - response_id: Option, - response_payload: Option, + payload: RawTraceEventPayload, ) -> Result<()> { + let (inference_call_id, status, response_id, upstream_request_id, response_payload) = + match payload { + RawTraceEventPayload::InferenceCompleted { + inference_call_id, + response_id, + upstream_request_id, + response_payload, + } => ( + inference_call_id, + ExecutionStatus::Completed, + response_id, + upstream_request_id, + Some(response_payload), + ), + RawTraceEventPayload::InferenceFailed { + inference_call_id, + upstream_request_id, + partial_response_payload, + .. + } => ( + inference_call_id, + ExecutionStatus::Failed, + None, + upstream_request_id, + partial_response_payload, + ), + RawTraceEventPayload::InferenceCancelled { + inference_call_id, + upstream_request_id, + partial_response_payload, + .. + } => ( + inference_call_id, + ExecutionStatus::Cancelled, + None, + upstream_request_id, + partial_response_payload, + ), + _ => bail!("complete_inference_call received a non-terminal inference event"), + }; + if !self .rollout .inference_calls @@ -156,23 +196,31 @@ impl TraceReducer { self.reduce_inference_response(wall_time_unix_ms, &inference_call_id, payload) }) .transpose()?; - let Some(inference) = self.rollout.inference_calls.get_mut(&inference_call_id) else { - bail!("inference call {inference_call_id} disappeared during response reduction"); - }; - // Turn-end cleanup can close a stream before the async mapper observes - // cancellation. Preserve that terminal status while still retaining any - // late partial response evidence from the mapper. - if inference.execution.status == ExecutionStatus::Running { - inference.execution.ended_at_unix_ms = Some(wall_time_unix_ms); - inference.execution.ended_seq = Some(seq); - inference.execution.status = status; - inference.upstream_request_id = response_id; - } - if let Some(response_payload) = response_payload { - inference.raw_response_payload_id = Some(response_payload.raw_payload_id); - } - if let Some(response_item_ids) = response_item_ids { - inference.response_item_ids = response_item_ids; + { + let Some(inference) = self.rollout.inference_calls.get_mut(&inference_call_id) else { + bail!("inference call {inference_call_id} disappeared during response reduction"); + }; + inference.response_id = response_id; + // Turn-end cleanup can close a stream before the async mapper observes + // cancellation. Preserve that terminal status while still retaining any + // late partial response evidence from the mapper. + if inference.execution.status == ExecutionStatus::Running { + inference.execution.ended_at_unix_ms = Some(wall_time_unix_ms); + inference.execution.ended_seq = Some(seq); + inference.execution.status = status; + } + // Turn-end cleanup can mark an inference terminal before the stream + // mapper records its late partial payload. Keep the server request + // id from that late payload even when the status is already closed. + if let Some(upstream_request_id) = upstream_request_id { + inference.upstream_request_id = Some(upstream_request_id); + } + if let Some(response_payload) = response_payload { + inference.raw_response_payload_id = Some(response_payload.raw_payload_id); + } + if let Some(response_item_ids) = response_item_ids { + inference.response_item_ids = response_item_ids; + } } Ok(()) } diff --git a/codex-rs/rollout-trace/src/reducer/inference_tests.rs b/codex-rs/rollout-trace/src/reducer/inference_tests.rs index d5ba0ba9d..77e0f39a9 100644 --- a/codex-rs/rollout-trace/src/reducer/inference_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/inference_tests.rs @@ -40,6 +40,7 @@ fn cancelled_inference_reduces_partial_response_items() -> anyhow::Result<()> { )?; writer.append(RawTraceEventPayload::InferenceCancelled { inference_call_id: "inference-1".to_string(), + upstream_request_id: Some("req-cancelled".to_string()), reason: "test interruption".to_string(), partial_response_payload: Some(partial_response), })?; @@ -49,6 +50,10 @@ fn cancelled_inference_reduces_partial_response_items() -> anyhow::Result<()> { let response_item_id = &inference.response_item_ids[0]; assert_eq!(inference.execution.status, ExecutionStatus::Cancelled); + assert_eq!( + inference.upstream_request_id, + Some("req-cancelled".to_string()), + ); assert_eq!(inference.response_item_ids.len(), 1); assert_eq!( rollout.conversation_items[response_item_id].kind, @@ -123,6 +128,7 @@ fn late_cancelled_inference_preserves_turn_end_status() -> anyhow::Result<()> { )?; writer.append(RawTraceEventPayload::InferenceCancelled { inference_call_id: "inference-1".to_string(), + upstream_request_id: Some("req-late-cancelled".to_string()), reason: "stream mapper noticed cancellation after turn end".to_string(), partial_response_payload: Some(partial_response.clone()), })?; @@ -135,6 +141,10 @@ fn late_cancelled_inference_preserves_turn_end_status() -> anyhow::Result<()> { inference.raw_response_payload_id, Some(partial_response.raw_payload_id), ); + assert_eq!( + inference.upstream_request_id, + Some("req-late-cancelled".to_string()), + ); assert_eq!(inference.response_item_ids.len(), 1); let response_item_id = &inference.response_item_ids[0]; assert_eq!( diff --git a/codex-rs/rollout-trace/src/reducer/mod.rs b/codex-rs/rollout-trace/src/reducer/mod.rs index 34a4e5dcf..3e8f6f135 100644 --- a/codex-rs/rollout-trace/src/reducer/mod.rs +++ b/codex-rs/rollout-trace/src/reducer/mod.rs @@ -226,47 +226,10 @@ impl TraceReducer { }, )?; } - RawTraceEventPayload::InferenceCompleted { - inference_call_id, - response_id, - response_payload, - } => { - self.complete_inference_call( - event.seq, - event.wall_time_unix_ms, - inference_call_id, - ExecutionStatus::Completed, - response_id, - Some(response_payload), - )?; - } - RawTraceEventPayload::InferenceFailed { - inference_call_id, - partial_response_payload, - .. - } => { - self.complete_inference_call( - event.seq, - event.wall_time_unix_ms, - inference_call_id, - ExecutionStatus::Failed, - /*response_id*/ None, - partial_response_payload, - )?; - } - RawTraceEventPayload::InferenceCancelled { - inference_call_id, - partial_response_payload, - .. - } => { - self.complete_inference_call( - event.seq, - event.wall_time_unix_ms, - inference_call_id, - ExecutionStatus::Cancelled, - /*response_id*/ None, - partial_response_payload, - )?; + payload @ (RawTraceEventPayload::InferenceCompleted { .. } + | RawTraceEventPayload::InferenceFailed { .. } + | RawTraceEventPayload::InferenceCancelled { .. }) => { + self.complete_inference_call(event.seq, event.wall_time_unix_ms, payload)?; } RawTraceEventPayload::ProtocolEventObserved { .. } => { // Protocol wrappers are raw debug breadcrumbs. Typed hooks own diff --git a/codex-rs/rollout-trace/src/reducer/test_support.rs b/codex-rs/rollout-trace/src/reducer/test_support.rs index eae93eff6..bd12e9d65 100644 --- a/codex-rs/rollout-trace/src/reducer/test_support.rs +++ b/codex-rs/rollout-trace/src/reducer/test_support.rs @@ -146,6 +146,7 @@ pub(crate) fn append_inference_completion( writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: inference_call_id.to_string(), response_id: Some(response_id.to_string()), + upstream_request_id: None, response_payload, })?; Ok(()) @@ -184,6 +185,7 @@ pub(crate) fn append_completed_inference( RawTraceEventPayload::InferenceCompleted { inference_call_id: inference_id.to_string(), response_id: Some(format!("resp-{inference_id}")), + upstream_request_id: None, response_payload: response, }, )?; diff --git a/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs b/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs index ddf76fe5a..e51e54a74 100644 --- a/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs +++ b/codex-rs/rollout-trace/src/reducer/tool/terminal_tests.rs @@ -551,6 +551,7 @@ fn append_inference_with_tool_call(writer: &TraceWriter) -> anyhow::Result<()> { writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: "inference-1".to_string(), response_id: Some("resp-1".to_string()), + upstream_request_id: None, response_payload: response, })?; Ok(()) diff --git a/codex-rs/rollout-trace/src/thread_tests.rs b/codex-rs/rollout-trace/src/thread_tests.rs index a27026f28..6421b29af 100644 --- a/codex-rs/rollout-trace/src/thread_tests.rs +++ b/codex-rs/rollout-trace/src/thread_tests.rs @@ -133,8 +133,8 @@ fn disabled_thread_context_accepts_trace_calls_without_writing() -> anyhow::Resu let inference_attempt = inference_trace.start_attempt(); inference_attempt.record_started(&serde_json::json!({ "kind": "inference" })); let token_usage: Option = None; - inference_attempt.record_completed("response-1", &token_usage, &[]); - inference_attempt.record_failed("inference failed", &[]); + inference_attempt.record_completed("response-1", Some("req-1"), &token_usage, &[]); + inference_attempt.record_failed("inference failed", /*upstream_request_id*/ None, &[]); let compaction_trace = thread_trace.compaction_trace_context( "turn-1", diff --git a/codex-rs/rollout-trace/src/writer.rs b/codex-rs/rollout-trace/src/writer.rs index 1676ff2e6..925f7bab7 100644 --- a/codex-rs/rollout-trace/src/writer.rs +++ b/codex-rs/rollout-trace/src/writer.rs @@ -226,6 +226,7 @@ mod tests { writer.append(RawTraceEventPayload::InferenceCompleted { inference_call_id: "inference-1".to_string(), response_id: Some("resp-1".to_string()), + upstream_request_id: Some("req-1".to_string()), response_payload: inference_response.clone(), })?; writer.append(RawTraceEventPayload::CodexTurnEnded {