diff --git a/codex-rs/codex-api/src/common.rs b/codex-rs/codex-api/src/common.rs index e4600e26aa..cafea020f8 100644 --- a/codex-rs/codex-api/src/common.rs +++ b/codex-rs/codex-api/src/common.rs @@ -238,24 +238,24 @@ pub struct ResponsesApiRequest { pub client_metadata: Option>, } -impl From<&ResponsesApiRequest> for ResponseCreateWsRequest { - fn from(request: &ResponsesApiRequest) -> Self { +impl<'a> From<&'a ResponsesApiRequest> for ResponseCreateWsRequest<'a> { + fn from(request: &'a ResponsesApiRequest) -> Self { Self { - model: request.model.clone(), - instructions: request.instructions.clone(), + model: &request.model, + instructions: &request.instructions, previous_response_id: None, - input: request.input.clone(), - tools: request.tools.clone(), - tool_choice: request.tool_choice.clone(), + input: &request.input, + tools: request.tools.as_deref(), + tool_choice: &request.tool_choice, parallel_tool_calls: request.parallel_tool_calls, - reasoning: request.reasoning.clone(), + reasoning: request.reasoning.as_ref(), store: request.store, stream: request.stream, - stream_options: request.stream_options.clone(), - include: request.include.clone(), - service_tier: request.service_tier.clone(), - prompt_cache_key: request.prompt_cache_key.clone(), - text: request.text.clone(), + stream_options: request.stream_options.as_ref(), + include: &request.include, + service_tier: request.service_tier.as_deref(), + prompt_cache_key: request.prompt_cache_key.as_deref(), + text: request.text.as_ref(), generate: None, client_metadata: request.client_metadata.clone(), } @@ -263,29 +263,29 @@ impl From<&ResponsesApiRequest> for ResponseCreateWsRequest { } #[derive(Debug, Serialize)] -pub struct ResponseCreateWsRequest { - pub model: String, - #[serde(skip_serializing_if = "String::is_empty")] - pub instructions: String, +pub struct ResponseCreateWsRequest<'a> { + pub model: &'a str, + #[serde(skip_serializing_if = "str::is_empty")] + pub instructions: &'a str, #[serde(skip_serializing_if = "Option::is_none")] pub previous_response_id: Option, - pub input: Vec, + pub input: &'a [ResponseItem], #[serde(skip_serializing_if = "Option::is_none")] - pub tools: Option>, - pub tool_choice: String, + pub tools: Option<&'a [Value]>, + pub tool_choice: &'a str, pub parallel_tool_calls: bool, - pub reasoning: Option, + pub reasoning: Option<&'a Reasoning>, pub store: bool, pub stream: bool, #[serde(skip_serializing_if = "Option::is_none")] - pub stream_options: Option, - pub include: Vec, + pub stream_options: Option<&'a StreamOptions>, + pub include: &'a [String], #[serde(skip_serializing_if = "Option::is_none")] - pub service_tier: Option, + pub service_tier: Option<&'a str>, #[serde(skip_serializing_if = "Option::is_none")] - pub prompt_cache_key: Option, + pub prompt_cache_key: Option<&'a str>, #[serde(skip_serializing_if = "Option::is_none")] - pub text: Option, + pub text: Option<&'a TextControls>, #[serde(skip_serializing_if = "Option::is_none")] pub generate: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -317,9 +317,9 @@ pub fn response_create_client_metadata( #[derive(Debug, Serialize)] #[serde(tag = "type")] #[allow(clippy::large_enum_variant)] -pub enum ResponsesWsRequest { +pub enum ResponsesWsRequest<'a> { #[serde(rename = "response.create")] - ResponseCreate(ResponseCreateWsRequest), + ResponseCreate(ResponseCreateWsRequest<'a>), } pub fn create_text_param_for_request( diff --git a/codex-rs/codex-api/src/endpoint/responses_websocket.rs b/codex-rs/codex-api/src/endpoint/responses_websocket.rs index 369179665c..24d6de825e 100644 --- a/codex-rs/codex-api/src/endpoint/responses_websocket.rs +++ b/codex-rs/codex-api/src/endpoint/responses_websocket.rs @@ -230,7 +230,7 @@ impl ResponsesWebsocketConnection { )] pub async fn stream_request( &self, - request: ResponsesWsRequest, + request: ResponsesWsRequest<'_>, connection_reused: bool, turn_state: Option>>, ) -> Result { @@ -245,7 +245,7 @@ impl ResponsesWebsocketConnection { let ResponsesWsRequest::ResponseCreate(ws_request) = &request; let client_metadata = ws_request.client_metadata.as_ref(); let timing_log_context = ResponsesWebsocketTimingLogContext { - model: ws_request.model.clone(), + model: ws_request.model.to_string(), session_id: client_metadata .and_then(|metadata| metadata.get(SESSION_ID_CLIENT_METADATA_KEY)) .cloned(), @@ -883,7 +883,7 @@ async fn send_websocket_request( Ok(()) } -fn serialize_websocket_request(request: &ResponsesWsRequest) -> Result { +fn serialize_websocket_request(request: &ResponsesWsRequest<'_>) -> Result { serde_json::to_string(request) .map_err(|err| ApiError::Stream(format!("failed to encode websocket request: {err}"))) } @@ -892,6 +892,7 @@ fn serialize_websocket_request(request: &ResponsesWsRequest) -> Result(&request_text).expect("parse websocket request"); - assert_eq!(wire_payload, previous_payload); + assert_eq!(wire_payload, expected_payload); } #[test] diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 283506f25f..9b0592d35a 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -1222,11 +1222,10 @@ impl ModelClientSession { fn prepare_websocket_request( &mut self, - payload: ResponseCreateWsRequest, request: &ResponsesApiRequest, - ) -> (ResponsesWsRequest, bool) { + ) -> (Option<(String, Vec)>, bool) { let Some(last_response) = self.get_last_response() else { - return (ResponsesWsRequest::ResponseCreate(payload), false); + return (None, false); }; let previous_response_id_from_untraced_warmup = self.websocket_session.last_response_from_untraced_warmup; @@ -1235,20 +1234,16 @@ impl ModelClientSession { Some(&last_response), /*allow_empty_delta*/ true, ) else { - return (ResponsesWsRequest::ResponseCreate(payload), false); + return (None, false); }; if last_response.response_id.is_empty() { trace!("incremental request failed, no previous response id"); - return (ResponsesWsRequest::ResponseCreate(payload), false); + return (None, false); } ( - ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest { - previous_response_id: Some(last_response.response_id), - input: incremental_items, - ..payload - }), + Some((last_response.response_id, incremental_items)), previous_response_id_from_untraced_warmup, ) } @@ -1548,7 +1543,7 @@ impl ModelClientSession { client_setup.agent_identity_telemetry.clone(), pending_retry, ); - let request = self.client.build_responses_request( + let mut request = self.client.build_responses_request( &client_setup.api_provider, prompt, model_info, @@ -1569,17 +1564,6 @@ impl ModelClientSession { if let Some(turn_state) = self.turn_state.get() { client_metadata.insert(X_CODEX_TURN_STATE_HEADER.to_string(), turn_state.clone()); } - let mut ws_payload = ResponseCreateWsRequest { - client_metadata: response_create_client_metadata( - Some(client_metadata), - request_trace.as_ref(), - ), - ..ResponseCreateWsRequest::from(&request) - }; - if warmup { - ws_payload.generate = Some(false); - } - match self .websocket_connection(WebsocketConnectParams { session_telemetry, @@ -1616,8 +1600,8 @@ impl ModelClientSession { Err(err) => return Err(self.client.state.provider.map_api_error(err)), } - let (mut ws_request, previous_response_id_from_untraced_warmup) = - self.prepare_websocket_request(ws_payload, &request); + let (incremental_request, previous_response_id_from_untraced_warmup) = + self.prepare_websocket_request(&request); let inference_trace_attempt = if warmup { // Prewarm sends `generate=false`; it is connection setup, not a // model inference attempt that should appear in rollout traces. @@ -1625,21 +1609,48 @@ impl ModelClientSession { } else { inference_trace.start_attempt() }; - stamp_ws_stream_request_start_ms(&mut ws_request); - let ResponsesWsRequest::ResponseCreate(ws_payload) = &mut ws_request; - let store = ws_payload.store; - self.client - .prepare_response_items_for_request(&mut ws_payload.input, store); if previous_response_id_from_untraced_warmup { // The transport can reuse an untraced warmup response id and omit the // already-sent input, but rollout replay needs the logical model-visible // request rather than the compressed websocket delta. inference_trace_attempt.record_started(&request); + } + + let (previous_response_id, mut incremental_items) = match incremental_request { + Some((response_id, items)) => (Some(response_id), Some(items)), + None => (None, None), + }; + let store = request.store; + let original_item_ids = if let Some(incremental_items) = &mut incremental_items { + self.client + .prepare_response_items_for_request(incremental_items, store); + None } else { + let original_item_ids = request + .input + .iter() + .map(|item| item.id().cloned()) + .collect::>(); + self.client + .prepare_response_items_for_request(&mut request.input, store); + Some(original_item_ids) + }; + let ws_payload = ResponseCreateWsRequest { + previous_response_id, + input: incremental_items.as_deref().unwrap_or(&request.input), + generate: if warmup { Some(false) } else { None }, + client_metadata: response_create_client_metadata( + Some(client_metadata), + request_trace.as_ref(), + ), + ..ResponseCreateWsRequest::from(&request) + }; + let mut ws_request = ResponsesWsRequest::ResponseCreate(ws_payload); + stamp_ws_stream_request_start_ms(&mut ws_request); + if !previous_response_id_from_untraced_warmup { inference_trace_attempt.record_started(&ws_request); } - self.websocket_session.last_request = Some(request); - self.websocket_session.last_response_from_untraced_warmup = warmup; + let websocket_connection = self.websocket_session.connection.as_ref().ok_or_else(|| { self.client.state.provider.map_api_error(ApiError::Stream( @@ -1652,18 +1663,24 @@ impl ModelClientSession { self.websocket_session.connection_reused(), Some(Arc::clone(&self.turn_state)), ) - .await - .map_err(|err| { - let response_debug_context = - extract_response_debug_context_from_api_error(&err); - let err = self.client.state.provider.map_api_error(err); - inference_trace_attempt.record_failed( - &err, - response_debug_context.request_id.as_deref(), - /*output_items*/ &[], - ); - err - })?; + .await; + if let Some(original_item_ids) = original_item_ids { + for (item, original_item_id) in request.input.iter_mut().zip(original_item_ids) { + item.set_id(original_item_id); + } + } + self.websocket_session.last_request = Some(request); + self.websocket_session.last_response_from_untraced_warmup = warmup; + let stream_result = stream_result.map_err(|err| { + let response_debug_context = extract_response_debug_context_from_api_error(&err); + let err = self.client.state.provider.map_api_error(err); + inference_trace_attempt.record_failed( + &err, + response_debug_context.request_id.as_deref(), + /*output_items*/ &[], + ); + err + })?; let (stream, last_request_rx) = map_response_stream( stream_result, request_session_telemetry, @@ -1848,7 +1865,7 @@ impl ModelClientSession { /// /// Meant to be called just before sending the request over the socket, to capture realistic /// transport timing. -fn stamp_ws_stream_request_start_ms(request: &mut ResponsesWsRequest) { +fn stamp_ws_stream_request_start_ms(request: &mut ResponsesWsRequest<'_>) { let ResponsesWsRequest::ResponseCreate(payload) = request; payload .client_metadata