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 {