support websockets

This commit is contained in:
Dylan Hurd
2026-03-18 00:28:30 -07:00
parent 77c14132c9
commit b9635bd082
2 changed files with 125 additions and 25 deletions
+5 -9
View File
@@ -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}"
+120 -16
View File
@@ -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<Value> {
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<T: Serialize>(
) -> serde_json::Result<Value> {
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,
}
}
}