From b9635bd082b1310877a29bce039373d6f195cf8a Mon Sep 17 00:00:00 2001 From: Dylan Hurd Date: Wed, 18 Mar 2026 00:28:30 -0700 Subject: [PATCH] support websockets --- codex-rs/core/src/client.rs | 14 +- codex-rs/core/src/response_item_id_serde.rs | 136 +++++++++++++++++--- 2 files changed, 125 insertions(+), 25 deletions(-) diff --git a/codex-rs/core/src/client.rs b/codex-rs/core/src/client.rs index 1204a6c756..4ae1f79d82 100644 --- a/codex-rs/core/src/client.rs +++ b/codex-rs/core/src/client.rs @@ -105,8 +105,8 @@ use crate::response_debug_context::extract_response_debug_context_from_api_error use crate::response_debug_context::telemetry_api_error_message; use crate::response_debug_context::telemetry_transport_error_message; use crate::response_item_id_serde::ResponseItemIdSerialization; -use crate::response_item_id_serde::serialize_response_create_ws_request_body; use crate::response_item_id_serde::serialize_responses_request_body; +use crate::response_item_id_serde::serialize_responses_ws_request_body; use crate::tools::spec::create_tools_json_for_responses_api; use crate::util::FeedbackRequestTags; use crate::util::emit_feedback_auth_recovery_tags; @@ -1209,14 +1209,10 @@ impl ModelClientSession { )) })?; let stream_result = if self.client.state.response_item_ids.is_enabled() { - let request_body = match &ws_request { - ResponsesWsRequest::ResponseCreate(request) => { - serialize_response_create_ws_request_body( - request, - self.client.state.response_item_ids, - ) - } - } + let request_body = serialize_responses_ws_request_body( + &ws_request, + self.client.state.response_item_ids, + ) .map_err(|err| { map_api_error(ApiError::Stream(format!( "failed to encode websocket request: {err}" diff --git a/codex-rs/core/src/response_item_id_serde.rs b/codex-rs/core/src/response_item_id_serde.rs index 9cdaf109aa..44753684c3 100644 --- a/codex-rs/core/src/response_item_id_serde.rs +++ b/codex-rs/core/src/response_item_id_serde.rs @@ -1,6 +1,7 @@ -use codex_api::ResponseCreateWsRequest; use codex_api::ResponsesApiRequest; +use codex_api::common::ResponsesWsRequest; use codex_protocol::models::ResponseItem; +use codex_protocol::protocol::CompactedItem; use codex_protocol::protocol::RolloutItem; use serde::Serialize; use serde_json::Value; @@ -25,11 +26,19 @@ pub(crate) fn serialize_responses_request_body( serialize_input_payload(request, &request.input, response_item_ids) } -pub(crate) fn serialize_response_create_ws_request_body( - request: &ResponseCreateWsRequest, +pub(crate) fn serialize_responses_ws_request_body( + request: &ResponsesWsRequest, response_item_ids: ResponseItemIdSerialization, ) -> serde_json::Result { - serialize_input_payload(request, &request.input, response_item_ids) + let mut value = serde_json::to_value(request)?; + if response_item_ids.is_enabled() { + match request { + ResponsesWsRequest::ResponseCreate(request) => { + attach_response_item_ids(value.get_mut("input"), &request.input); + } + } + } + Ok(value) } pub(crate) fn serialize_rollout_line( @@ -51,29 +60,48 @@ fn serialize_input_payload( ) -> serde_json::Result { let mut value = serde_json::to_value(request)?; if response_item_ids.is_enabled() { - attach_input_item_ids(&mut value, input); + attach_response_item_ids(value.get_mut("input"), input); } Ok(value) } fn attach_rollout_item_id(value: &mut Value, item: &RolloutItem) { - let RolloutItem::ResponseItem(response_item) = item else { - return; - }; - let Some(id) = response_item_id(response_item) else { - return; - }; + match item { + RolloutItem::ResponseItem(response_item) => { + let Some(id) = response_item_id(response_item) else { + return; + }; - if let Some(payload) = value.get_mut("payload").and_then(Value::as_object_mut) { - payload.insert("id".to_string(), Value::String(id.to_string())); + if let Some(payload) = value.get_mut("payload").and_then(Value::as_object_mut) { + payload.insert("id".to_string(), Value::String(id.to_string())); + } + } + RolloutItem::Compacted(compacted_item) => { + attach_compacted_replacement_history_ids(value, compacted_item); + } + RolloutItem::SessionMeta(_) | RolloutItem::TurnContext(_) | RolloutItem::EventMsg(_) => {} } } -pub(crate) fn attach_input_item_ids(payload_json: &mut Value, original_items: &[ResponseItem]) { - let Some(input_value) = payload_json.get_mut("input") else { +fn attach_compacted_replacement_history_ids(value: &mut Value, compacted_item: &CompactedItem) { + let Some(payload) = value.get_mut("payload").and_then(Value::as_object_mut) else { return; }; - let Value::Array(items) = input_value else { + + attach_response_item_ids( + payload.get_mut("replacement_history"), + compacted_item + .replacement_history + .as_deref() + .unwrap_or_default(), + ); +} + +fn attach_response_item_ids(items_value: Option<&mut Value>, original_items: &[ResponseItem]) { + let Some(items_value) = items_value else { + return; + }; + let Value::Array(items) = items_value else { return; }; @@ -123,3 +151,79 @@ struct RolloutLineRef<'a> { #[serde(flatten)] item: &'a RolloutItem, } + +#[cfg(test)] +mod tests { + use super::ResponseItemIdSerialization; + use super::serialize_responses_ws_request_body; + use super::serialize_rollout_line; + use codex_api::ResponseCreateWsRequest; + use codex_api::common::ResponsesWsRequest; + use codex_protocol::models::ContentItem; + use codex_protocol::models::ResponseItem; + use codex_protocol::protocol::CompactedItem; + use codex_protocol::protocol::RolloutItem; + use pretty_assertions::assert_eq; + use serde_json::json; + + #[test] + fn serialize_responses_ws_request_body_preserves_response_create_wrapper() { + let request = ResponsesWsRequest::ResponseCreate(ResponseCreateWsRequest { + model: "gpt-test".to_string(), + instructions: "Be helpful".to_string(), + previous_response_id: None, + input: vec![assistant_message(Some("msg_123"))], + tools: Vec::new(), + tool_choice: "auto".to_string(), + parallel_tool_calls: false, + reasoning: None, + store: true, + stream: true, + include: Vec::new(), + service_tier: None, + prompt_cache_key: None, + text: None, + generate: None, + client_metadata: None, + }); + + let value = + serialize_responses_ws_request_body(&request, ResponseItemIdSerialization::Enabled) + .expect("websocket request should serialize"); + + assert_eq!(value["type"], json!("response.create")); + assert_eq!(value["input"][0]["id"], json!("msg_123")); + } + + #[test] + fn serialize_rollout_line_preserves_ids_in_compacted_replacement_history() { + let line = serialize_rollout_line( + "2026-03-18T00:00:00Z".to_string(), + &RolloutItem::Compacted(CompactedItem { + message: "compacted".to_string(), + replacement_history: Some(vec![assistant_message(Some("msg_123"))]), + }), + ResponseItemIdSerialization::Enabled, + ) + .expect("rollout line should serialize"); + + let value: serde_json::Value = + serde_json::from_str(&line).expect("serialized rollout line should be valid json"); + assert_eq!( + value["payload"]["replacement_history"][0]["id"], + json!("msg_123") + ); + } + + fn assistant_message(id: Option<&str>) -> ResponseItem { + ResponseItem::Message { + id: id.map(str::to_string), + role: "assistant".to_string(), + content: vec![ContentItem::OutputText { + text: "hello".to_string(), + }], + end_turn: None, + phase: None, + } + } +}