openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-18 17:18:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid cloning Responses WebSocket payloads (#34381)

Browse Source 
## Why

Building a Responses WebSocket request copied the model, instructions, input
history, tools, and other request settings into a second payload before
serialization.

## What changed

- Make `ResponseCreateWsRequest` borrow shared fields from
  `ResponsesApiRequest`.
- Keep only WebSocket-specific fields owned, and allocate a separate input
  vector only for incremental continuations.
- Restore response item IDs after preparing a full request so the request
  retained for continuation matching keeps its original state.

## Testing

Update the serialization test to verify that the borrowed WebSocket request
produces the same wire payload as the Responses API request plus its
WebSocket-specific fields.

GitOrigin-RevId: b63d746b7f5164b5204ab95b057c914a4aeb71e1
This commit is contained in:
Charlie Marsh
2026-07-20 17:18:25 +00:00
committed by copyberry
parent 8431dc590a
commit 6b9a5592a6
3 changed files with 105 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/codex-api/src/common.rs
+28 -28
View File
@@ -238,24 +238,24 @@ pub struct ResponsesApiRequest {
pub client_metadata: Option<HashMap<String, String>>,
}
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<String>,
pub input: Vec<ResponseItem>,
pub input: &'a [ResponseItem],
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<Value>>,
pub tool_choice: String,
pub tools: Option<&'a [Value]>,
pub tool_choice: &'a str,
pub parallel_tool_calls: bool,
pub reasoning: Option<Reasoning>,
pub reasoning: Option<&'a Reasoning>,
pub store: bool,
pub stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub stream_options: Option<StreamOptions>,
pub include: Vec<String>,
pub stream_options: Option<&'a StreamOptions>,
pub include: &'a [String],
#[serde(skip_serializing_if = "Option::is_none")]
pub service_tier: Option<String>,
pub service_tier: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_cache_key: Option<String>,
pub prompt_cache_key: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<TextControls>,
pub text: Option<&'a TextControls>,
#[serde(skip_serializing_if = "Option::is_none")]
pub generate: Option<bool>,
#[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(
codex-rs/codex-api/src/endpoint/responses_websocket.rs
+16 -8
View File
@@ -230,7 +230,7 @@ impl ResponsesWebsocketConnection {
)]
pub async fn stream_request(
&self,
request: ResponsesWsRequest,
request: ResponsesWsRequest<'_>,
connection_reused: bool,
turn_state: Option<Arc<OnceLock<String>>>,
) -> Result<ResponseStream, ApiError> {
@@ -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<String, ApiError> {
fn serialize_websocket_request(request: &ResponsesWsRequest<'_>) -> Result<String, ApiError> {
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<String, A
mod tests {
use super::*;
use crate::common::ResponseCreateWsRequest;
use crate::common::ResponsesApiRequest;
use codex_protocol::ResponseItemId;
use codex_protocol::models::ContentItem;
use codex_protocol::models::ResponseItem;
@@ -901,10 +902,9 @@ mod tests {
#[test]
fn direct_serialization_preserves_websocket_request_payload() {
let request = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest {
let api_request = ResponsesApiRequest {
model: "gpt-test".to_string(),
instructions: "Use the available tools.".to_string(),
previous_response_id: Some("resp-1".to_string()),
input: vec![ResponseItem::Message {
id: Some(ResponseItemId::with_suffix("msg", "1")),
role: "user".to_string(),
@@ -929,20 +929,28 @@ mod tests {
service_tier: Some("priority".to_string()),
prompt_cache_key: Some("cache-key".to_string()),
text: None,
generate: Some(false),
client_metadata: Some(HashMap::from([(
"traceparent".to_string(),
"00-0123456789abcdef0123456789abcdef-0123456789abcdef-01".to_string(),
)])),
};
let request = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest {
previous_response_id: Some("resp-1".to_string()),
generate: Some(false),
..ResponseCreateWsRequest::from(&api_request)
});
let previous_payload = serde_json::to_value(&request).expect("serialize previous payload");
let mut expected_payload =
serde_json::to_value(&api_request).expect("serialize responses API request");
expected_payload["type"] = json!("response.create");
expected_payload["previous_response_id"] = json!("resp-1");
expected_payload["generate"] = json!(false);
let request_text =
serialize_websocket_request(&request).expect("serialize websocket request");
let wire_payload =
serde_json::from_str::<Value>(&request_text).expect("parse websocket request");
assert_eq!(wire_payload, previous_payload);
assert_eq!(wire_payload, expected_payload);
}
#[test]
codex-rs/core/src/client.rs
+61 -44
View File
@@ -1222,11 +1222,10 @@ impl ModelClientSession {
fn prepare_websocket_request(
&mut self,
payload: ResponseCreateWsRequest,
request: &ResponsesApiRequest,
) -> (ResponsesWsRequest, bool) {
) -> (Option<(String, Vec<ResponseItem>)>, 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::<Vec<_>>();
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