From bb05429f8dde951cc647a8ee4aaedcb745393853 Mon Sep 17 00:00:00 2001 From: Charles Cunningham Date: Thu, 5 Mar 2026 21:03:08 -0800 Subject: [PATCH] Persist Responses API IDs in rollouts Write request IDs as soon as headers arrive, including error paths, and persist explicit null fields in rollout response metadata. Add rollout coverage for both successful and error responses. Co-authored-by: Codex --- codex-rs/codex-api/src/common.rs | 14 +- .../src/endpoint/responses_websocket.rs | 27 ++- codex-rs/codex-api/src/sse/responses.rs | 47 +++++- codex-rs/core/src/client.rs | 21 +-- codex-rs/core/src/client_common.rs | 1 + codex-rs/core/src/codex.rs | 46 ++++- codex-rs/core/src/error.rs | 37 ++++ codex-rs/core/src/rollout/policy.rs | 1 + codex-rs/core/tests/suite/client.rs | 158 ++++++++++++++++++ codex-rs/otel/src/traces/otel_manager.rs | 2 +- codex-rs/protocol/src/protocol.rs | 10 ++ 11 files changed, 340 insertions(+), 24 deletions(-) diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index 31b4dcdb44..f94c775159 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -6,6 +6,7 @@ use codex_protocol::openai_models::ReasoningEffort as ReasoningEffortConfig; use codex_protocol::protocol::RateLimitSnapshot; use codex_protocol::protocol::TokenUsage; use futures::Stream; +use http::HeaderMap; use serde::Deserialize; use serde::Serialize; use serde_json::Value; @@ -54,7 +55,9 @@ pub struct MemorySummarizeOutput { #[derive(Debug)] pub enum ResponseEvent { - Created, + Created { + response_id: Option, + }, OutputItemDone(ResponseItem), OutputItemAdded(ResponseItem), /// Emitted when the server includes `OpenAI-Model` on the stream response. @@ -84,6 +87,14 @@ pub enum ResponseEvent { ModelsEtag(String), } +pub(crate) fn extract_request_id(headers: &HeaderMap) -> Option { + headers + .get("x-request-id") + .or_else(|| headers.get("x-oai-request-id")) + .and_then(|value| value.to_str().ok()) + .map(str::to_string) +} + #[derive(Debug, Serialize, Clone, PartialEq)] pub struct Reasoning { #[serde(skip_serializing_if = "Option::is_none")] @@ -238,6 +249,7 @@ pub fn create_text_param_for_request( pub struct ResponseStream { pub rx_event: mpsc::Receiver>, + pub initial_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 925f7d52d0..c0643219c0 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::auth::add_auth_headers_to_header_map; use crate::common::ResponseEvent; use crate::common::ResponseStream; use crate::common::ResponsesWsRequest; +use crate::common::extract_request_id; use crate::error::ApiError; use crate::provider::Provider; use crate::rate_limits::parse_rate_limit_event; @@ -174,6 +175,7 @@ pub struct ResponsesWebsocketConnection { server_reasoning_included: bool, models_etag: Option, server_model: Option, + request_id: Option, telemetry: Option>, } @@ -185,6 +187,7 @@ impl std::fmt::Debug for ResponsesWebsocketConnection { .field("server_reasoning_included", &self.server_reasoning_included) .field("models_etag", &self.models_etag) .field("server_model", &self.server_model) + .field("request_id", &self.request_id) .field("telemetry", &self.telemetry.as_ref().map(|_| "")) .finish() } @@ -197,6 +200,7 @@ impl ResponsesWebsocketConnection { server_reasoning_included: bool, models_etag: Option, server_model: Option, + request_id: Option, telemetry: Option>, ) -> Self { Self { @@ -205,6 +209,7 @@ impl ResponsesWebsocketConnection { server_reasoning_included, models_etag, server_model, + request_id, telemetry, } } @@ -224,6 +229,7 @@ impl ResponsesWebsocketConnection { let server_reasoning_included = self.server_reasoning_included; let models_etag = self.models_etag.clone(); let server_model = self.server_model.clone(); + let request_id = self.request_id.clone(); let telemetry = self.telemetry.clone(); let request_body = serde_json::to_value(&request).map_err(|err| { ApiError::Stream(format!("failed to encode websocket request: {err}")) @@ -266,7 +272,10 @@ impl ResponsesWebsocketConnection { } }); - Ok(ResponseStream { rx_event }) + Ok(ResponseStream { + rx_event, + initial_request_id: request_id, + }) } } @@ -296,7 +305,7 @@ impl ResponsesWebsocketClient { merge_request_headers(&self.provider.headers, extra_headers, default_headers); add_auth_headers_to_header_map(&self.auth, &mut headers); - let (stream, server_reasoning_included, models_etag, server_model) = + let (stream, server_reasoning_included, models_etag, server_model, request_id) = connect_websocket(ws_url, headers, turn_state.clone()).await?; Ok(ResponsesWebsocketConnection::new( stream, @@ -304,6 +313,7 @@ impl ResponsesWebsocketClient { server_reasoning_included, models_etag, server_model, + request_id, telemetry, )) } @@ -328,7 +338,16 @@ async fn connect_websocket( url: Url, headers: HeaderMap, turn_state: Option>>, -) -> Result<(WsStream, bool, Option, Option), ApiError> { +) -> Result< + ( + WsStream, + bool, + Option, + Option, + Option, + ), + ApiError, +> { ensure_rustls_crypto_provider(); info!("connecting to websocket: {url}"); @@ -370,6 +389,7 @@ async fn connect_websocket( .get(OPENAI_MODEL_HEADER) .and_then(|value| value.to_str().ok()) .map(ToString::to_string); + let request_id = extract_request_id(response.headers()); if let Some(turn_state) = turn_state && let Some(header_value) = response .headers() @@ -383,6 +403,7 @@ async fn connect_websocket( reasoning_included, models_etag, server_model, + request_id, )) } diff --git a/codex-rs/codex-api/src/sse/responses.rs b/codex-rs/codex-api/src/sse/responses.rs index d01e1d918a..77286df104 100644 --- a/codex-rs/codex-api/src/sse/responses.rs +++ b/codex-rs/codex-api/src/sse/responses.rs @@ -1,5 +1,6 @@ use crate::common::ResponseEvent; use crate::common::ResponseStream; +use crate::common::extract_request_id; use crate::error::ApiError; use crate::rate_limits::parse_all_rate_limits; use crate::telemetry::SseTelemetry; @@ -46,7 +47,10 @@ pub fn stream_from_fixture( let stream = ReaderStream::new(reader).map_err(|err| TransportError::Network(err.to_string())); let (tx_event, rx_event) = mpsc::channel::>(1600); tokio::spawn(process_sse(Box::pin(stream), tx_event, idle_timeout, None)); - Ok(ResponseStream { rx_event }) + Ok(ResponseStream { + rx_event, + initial_request_id: None, + }) } pub fn spawn_response_stream( @@ -55,6 +59,7 @@ pub fn spawn_response_stream( telemetry: Option>, turn_state: Option>>, ) -> ResponseStream { + let request_id = extract_request_id(&stream_response.headers); let rate_limit_snapshots = parse_all_rate_limits(&stream_response.headers); let models_etag = stream_response .headers @@ -97,7 +102,10 @@ pub fn spawn_response_stream( process_sse(stream_response.bytes, tx_event, idle_timeout, telemetry).await; }); - ResponseStream { rx_event } + ResponseStream { + rx_event, + initial_request_id: request_id, + } } #[derive(Debug, Deserialize)] @@ -263,7 +271,13 @@ pub fn process_responses_event( } "response.created" => { if event.response.is_some() { - return Ok(Some(ResponseEvent::Created {})); + let response_id = event + .response + .as_ref() + .and_then(|response| response.get("id")) + .and_then(Value::as_str) + .map(str::to_string); + return Ok(Some(ResponseEvent::Created { response_id })); } } "response.failed" => { @@ -771,7 +785,7 @@ mod tests { } fn is_created(ev: &ResponseEvent) -> bool { - matches!(ev, ResponseEvent::Created) + matches!(ev, ResponseEvent::Created { .. }) } fn is_output(ev: &ResponseEvent) -> bool { matches!(ev, ResponseEvent::OutputItemDone(_)) @@ -869,6 +883,21 @@ mod tests { } } + #[tokio::test] + async fn spawn_response_stream_emits_request_id_header() { + let mut headers = HeaderMap::new(); + headers.insert("x-request-id", HeaderValue::from_static("req_123")); + let bytes = stream::iter(Vec::>::new()); + let stream_response = StreamResponse { + status: StatusCode::OK, + headers, + bytes: Box::pin(bytes), + }; + + let stream = spawn_response_stream(stream_response, idle_timeout(), None, None); + assert_eq!(stream.initial_request_id.as_deref(), Some("req_123")); + } + #[tokio::test] async fn process_sse_ignores_response_model_field_in_payload() { let events = run_sse(vec![ @@ -890,7 +919,10 @@ mod tests { .await; assert_eq!(events.len(), 2); - assert_matches!(&events[0], ResponseEvent::Created); + assert_matches!( + &events[0], + ResponseEvent::Created { response_id } if response_id.as_deref() == Some("resp-1") + ); assert_matches!( &events[1], ResponseEvent::Completed { @@ -926,7 +958,10 @@ mod tests { &events[0], ResponseEvent::ServerModel(model) if model == CYBER_RESTRICTED_MODEL_FOR_TESTS ); - assert_matches!(&events[1], ResponseEvent::Created); + assert_matches!( + &events[1], + ResponseEvent::Created { response_id } if response_id.as_deref() == Some("resp-1") + ); assert_matches!( &events[2], ResponseEvent::Completed { diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index ce354d9b04..8c97ffbe8e 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1094,18 +1094,13 @@ fn build_responses_headers( headers } -fn map_response_stream( - api_stream: S, +fn map_response_stream( + api_stream: codex_api::ResponseStream, otel_manager: OtelManager, -) -> (ResponseStream, oneshot::Receiver) -where - S: futures::Stream> - + Unpin - + Send - + 'static, -{ +) -> (ResponseStream, oneshot::Receiver) { let (tx_event, rx_event) = mpsc::channel::>(1600); let (tx_last_response, rx_last_response) = oneshot::channel::(); + let initial_request_id = api_stream.initial_request_id.clone(); tokio::spawn(async move { let mut logged_error = false; @@ -1173,7 +1168,13 @@ where } }); - (ResponseStream { rx_event }, rx_last_response) + ( + ResponseStream { + rx_event, + initial_request_id, + }, + rx_last_response, + ) } /// Handles a 401 response by optionally refreshing ChatGPT tokens once. diff --git a/codex-rs/core/src/client_common.rs b/codex-rs/core/src/client_common.rs index d775c9a376..9ae2be1e99 100644 --- a/codex-rs/core/src/client_common.rs +++ b/codex-rs/core/src/client_common.rs @@ -223,6 +223,7 @@ pub(crate) mod tools { pub struct ResponseStream { pub(crate) rx_event: mpsc::Receiver>, + pub(crate) initial_request_id: Option, } impl Stream for ResponseStream { diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index f7f90eeaa8..e5e2b91907 100644 --- a/codex-rs/core/src/codex.rs +++ b/codex-rs/core/src/codex.rs @@ -89,6 +89,7 @@ use codex_protocol::protocol::HasLegacyEvent; use codex_protocol::protocol::ItemCompletedEvent; use codex_protocol::protocol::ItemStartedEvent; use codex_protocol::protocol::RawResponseItemEvent; +use codex_protocol::protocol::ResponseMetadataEvent; use codex_protocol::protocol::ReviewRequest; use codex_protocol::protocol::RolloutItem; use codex_protocol::protocol::SessionSource; @@ -5986,6 +5987,7 @@ fn realtime_text_for_event(msg: &EventMsg) -> Option { | EventMsg::ShutdownComplete | EventMsg::EnteredReviewMode(_) | EventMsg::ExitedReviewMode(_) + | EventMsg::ResponseMetadata(_) | EventMsg::RawResponseItem(_) | EventMsg::ItemStarted(_) | EventMsg::AgentMessageContentDelta(_) @@ -6308,7 +6310,7 @@ async fn try_run_sampling_request( auth_mode = sess.services.auth_manager.auth_mode(), features = sess.features.enabled_features(), ); - let mut stream = client_session + let stream_result = client_session .stream( prompt, &turn_context.model_info, @@ -6320,7 +6322,34 @@ async fn try_run_sampling_request( ) .instrument(trace_span!("stream_request")) .or_cancel(&cancellation_token) - .await??; + .await; + let mut stream = match stream_result { + Ok(Ok(stream)) => stream, + Ok(Err(err)) => { + if let Some(request_id) = err.request_id() { + sess.send_event( + &turn_context, + EventMsg::ResponseMetadata(ResponseMetadataEvent { + request_id: Some(request_id.to_string()), + response_id: None, + }), + ) + .await; + } + return Err(err); + } + Err(codex_async_utils::CancelErr::Cancelled) => return Err(CodexErr::TurnAborted), + }; + if let Some(request_id) = stream.initial_request_id.clone() { + sess.send_event( + &turn_context, + EventMsg::ResponseMetadata(ResponseMetadataEvent { + request_id: Some(request_id), + response_id: None, + }), + ) + .await; + } let tool_runtime = ToolCallRuntime::new( Arc::clone(&router), @@ -6372,7 +6401,18 @@ async fn try_run_sampling_request( .record_responses(&handle_responses, &event); match event { - ResponseEvent::Created => {} + ResponseEvent::Created { response_id } => { + if let Some(response_id) = response_id { + sess.send_event( + &turn_context, + EventMsg::ResponseMetadata(ResponseMetadataEvent { + request_id: None, + response_id: Some(response_id), + }), + ) + .await; + } + } ResponseEvent::OutputItemDone(item) => { let previously_active_item = active_item.take(); if let Some(previous) = previously_active_item.as_ref() diff --git a/codex-rs/core/src/error.rs b/codex-rs/core/src/error.rs index ad49c611fc..84a91fe47b 100644 --- a/codex-rs/core/src/error.rs +++ b/codex-rs/core/src/error.rs @@ -228,6 +228,43 @@ impl CodexErr { CodexErr::LandlockRuleset(_) | CodexErr::LandlockPathFd(_) => false, } } + + pub fn request_id(&self) -> Option<&str> { + match self { + CodexErr::UnexpectedStatus(err) => err.request_id.as_deref(), + CodexErr::RetryLimit(err) => err.request_id.as_deref(), + CodexErr::ResponseStreamFailed(err) => err.request_id.as_deref(), + CodexErr::TurnAborted + | CodexErr::Interrupted + | CodexErr::Stream(..) + | CodexErr::Timeout + | CodexErr::ConnectionFailed(_) + | CodexErr::InternalServerError + | CodexErr::InternalAgentDied + | CodexErr::Io(_) + | CodexErr::Json(_) + | CodexErr::TokioJoin(_) + | CodexErr::EnvVar(_) + | CodexErr::Fatal(_) + | CodexErr::UsageNotIncluded + | CodexErr::QuotaExceeded + | CodexErr::InvalidImageRequest() + | CodexErr::InvalidRequest(_) + | CodexErr::RefreshTokenFailed(_) + | CodexErr::UnsupportedOperation(_) + | CodexErr::Sandbox(_) + | CodexErr::LandlockSandboxExecutableNotProvided + | CodexErr::ContextWindowExceeded + | CodexErr::ThreadNotFound(_) + | CodexErr::AgentLimitReached { .. } + | CodexErr::Spawn + | CodexErr::SessionConfiguredNotFirstEvent + | CodexErr::UsageLimitReached(_) + | CodexErr::ServerOverloaded => None, + #[cfg(target_os = "linux")] + CodexErr::LandlockRuleset(_) | CodexErr::LandlockPathFd(_) => None, + } + } } #[derive(Debug)] diff --git a/codex-rs/core/src/rollout/policy.rs b/codex-rs/core/src/rollout/policy.rs index 92b6076812..8bf8964644 100644 --- a/codex-rs/core/src/rollout/policy.rs +++ b/codex-rs/core/src/rollout/policy.rs @@ -91,6 +91,7 @@ fn event_msg_persistence_mode(ev: &EventMsg) -> Option { match ev { EventMsg::UserMessage(_) | EventMsg::AgentMessage(_) + | EventMsg::ResponseMetadata(_) | EventMsg::AgentReasoning(_) | EventMsg::AgentReasoningRawContent(_) | EventMsg::TokenCount(_) diff --git a/codex-rs/core/tests/suite/client.rs b/codex-rs/core/tests/suite/client.rs index 3ac9a72c3c..2c5333b09b 100644 --- a/codex-rs/core/tests/suite/client.rs +++ b/codex-rs/core/tests/suite/client.rs @@ -142,6 +142,164 @@ fn write_auth_json( fake_jwt } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rollout_persists_responses_request_and_response_ids() { + skip_if_no_network!(); + + let server = MockServer::start().await; + let response = ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .insert_header("x-request-id", "req_123") + .set_body_string(sse(vec![ + ev_response_created("resp_123"), + ev_completed("resp_123"), + ])); + let response_mock = core_test_support::responses::mount_response_once(&server, response).await; + + let test = test_codex().build(&server).await.expect("build test codex"); + test.submit_turn("hello?") + .await + .expect("submit turn should succeed"); + + let request = response_mock.single_request(); + assert_eq!(request.path(), "/v1/responses"); + + let rollout_path = test.codex.rollout_path().expect("rollout path"); + let rollout = std::fs::read_to_string(&rollout_path).expect("read rollout"); + let json_lines: Vec = rollout + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| serde_json::from_str(line).expect("rollout json line")) + .collect(); + let lines: Vec = rollout + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| serde_json::from_str(line).expect("rollout line")) + .collect(); + + assert!(lines.iter().any(|line| { + matches!( + &line.item, + RolloutItem::EventMsg(EventMsg::ResponseMetadata(event)) + if event.request_id.as_deref() == Some("req_123") + && event.response_id.is_none() + ) + })); + assert!(lines.iter().any(|line| { + matches!( + &line.item, + RolloutItem::EventMsg(EventMsg::ResponseMetadata(event)) + if event.request_id.is_none() + && event.response_id.as_deref() == Some("resp_123") + ) + })); + assert!(json_lines.iter().any(|line| { + line["type"] == "event_msg" + && line["payload"]["type"] == "response_metadata" + && line["payload"]["request_id"] == "req_123" + && line["payload"]["response_id"].is_null() + })); + assert!(json_lines.iter().any(|line| { + line["type"] == "event_msg" + && line["payload"]["type"] == "response_metadata" + && line["payload"]["request_id"].is_null() + && line["payload"]["response_id"] == "resp_123" + })); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rollout_persists_request_id_when_responses_request_errors() -> anyhow::Result<()> { + skip_if_no_network!(Ok(())); + + let server = MockServer::start().await; + let response = ResponseTemplate::new(418) + .insert_header("x-request-id", "req_error_123") + .set_body_json(json!({ + "error": { + "message": "teapot" + } + })); + Mock::given(method("POST")) + .and(path("/v1/responses")) + .respond_with(response) + .expect(1) + .mount(&server) + .await; + + let TestCodex { codex, .. } = test_codex() + .with_config(|config| { + config.model_provider.stream_max_retries = Some(0); + }) + .build(&server) + .await?; + codex + .submit(Op::UserInput { + items: vec![UserInput::Text { + text: "trigger error".into(), + text_elements: Vec::new(), + }], + final_output_json_schema: None, + }) + .await?; + + let response_metadata = wait_for_event(&codex, |event| { + matches!( + event, + EventMsg::ResponseMetadata(metadata) + if metadata.request_id.as_deref() == Some("req_error_123") + && metadata.response_id.is_none() + ) + }) + .await; + assert!( + matches!( + &response_metadata, + EventMsg::ResponseMetadata(metadata) + if metadata.request_id.as_deref() == Some("req_error_123") + && metadata.response_id.is_none() + ), + "expected response metadata with request id; got {response_metadata:?}" + ); + + let error_event = wait_for_event(&codex, |event| matches!(event, EventMsg::Error(_))).await; + assert!( + matches!(error_event, EventMsg::Error(_)), + "expected error event; got {error_event:?}" + ); + + wait_for_event(&codex, |event| matches!(event, EventMsg::TurnComplete(_))).await; + + let rollout_path = codex.rollout_path().expect("rollout path"); + let rollout = std::fs::read_to_string(&rollout_path).expect("read rollout"); + let json_lines: Vec = rollout + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| serde_json::from_str(line).expect("rollout json line")) + .collect(); + let lines: Vec = rollout + .lines() + .filter(|line| !line.trim().is_empty()) + .map(|line| serde_json::from_str(line).expect("rollout line")) + .collect(); + + assert!(lines.iter().any(|line| { + matches!( + &line.item, + RolloutItem::EventMsg(EventMsg::ResponseMetadata(event)) + if event.request_id.as_deref() == Some("req_error_123") + && event.response_id.is_none() + ) + })); + assert!(json_lines.iter().any(|line| { + line["type"] == "event_msg" + && line["payload"]["type"] == "response_metadata" + && line["payload"]["request_id"] == "req_error_123" + && line["payload"]["response_id"].is_null() + })); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn resume_includes_initial_messages_and_sends_prior_items() { skip_if_no_network!(); diff --git a/codex-rs/otel/src/traces/otel_manager.rs b/codex-rs/otel/src/traces/otel_manager.rs index aa11fce3b0..e9c71c3c21 100644 --- a/codex-rs/otel/src/traces/otel_manager.rs +++ b/codex-rs/otel/src/traces/otel_manager.rs @@ -744,7 +744,7 @@ impl OtelManager { fn responses_type(event: &ResponseEvent) -> String { match event { - ResponseEvent::Created => "created".into(), + ResponseEvent::Created { .. } => "created".into(), ResponseEvent::OutputItemDone(item) => OtelManager::responses_item_type(item), ResponseEvent::OutputItemAdded(item) => OtelManager::responses_item_type(item), ResponseEvent::Completed { .. } => "completed".into(), diff --git a/codex-rs/protocol/src/protocol.rs b/codex-rs/protocol/src/protocol.rs index 8e7e9881ec..c857c820b1 100644 --- a/codex-rs/protocol/src/protocol.rs +++ b/codex-rs/protocol/src/protocol.rs @@ -1176,6 +1176,8 @@ pub enum EventMsg { /// Exited review mode with an optional final result to apply. ExitedReviewMode(ExitedReviewModeEvent), + ResponseMetadata(ResponseMetadataEvent), + RawResponseItem(RawResponseItemEvent), ItemStarted(ItemStartedEvent), @@ -1356,6 +1358,14 @@ impl CodexErrorInfo { } } +#[derive(Debug, Clone, Deserialize, Serialize, TS, JsonSchema)] +pub struct ResponseMetadataEvent { + #[serde(default)] + pub request_id: Option, + #[serde(default)] + pub response_id: Option, +} + #[derive(Debug, Clone, Deserialize, Serialize, TS, JsonSchema)] pub struct RawResponseItemEvent { pub item: ResponseItem,