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,