Reduce response serialization and rollout scan overhead (#36006)

## What changed

- Keep `ClientResponsePayload` typed through the app-server outgoing queue and serialize it at the transport boundary, avoiding an intermediate `serde_json::Value` while preserving the JSON-RPC wire shape.
- Return a JSON-RPC internal error when a response cannot be serialized, including over remote control, and skip unserializable analytics responses.
- Scan reverse JSONL input a chunk at a time with a larger buffer instead of processing each byte individually.

## Testing

- Cover typed response serialization and serialization failures for direct and remote-control transports.
- Cover filtering unserializable responses from analytics.

GitOrigin-RevId: cc4fd7a43461a273315bb95261bd1ae3ff65fbf7
This commit is contained in:
acrognale-oai
2026-07-29 18:29:26 +00:00
committed by copyberry
parent 7579a2b413
commit c41a38dd10
14 changed files with 293 additions and 132 deletions
+7 -4
View File
@@ -485,7 +485,7 @@ impl AnalyticsEventsClient {
&self,
connection_id: u64,
request_id: RequestId,
response: ClientResponsePayload,
response: &ClientResponsePayload,
) {
self.track_response_inner(
connection_id,
@@ -499,7 +499,7 @@ impl AnalyticsEventsClient {
&self,
connection_id: u64,
request_id: RequestId,
response: ClientResponsePayload,
response: &ClientResponsePayload,
thread_originator: String,
) {
self.track_response_inner(connection_id, request_id, response, Some(thread_originator));
@@ -509,7 +509,7 @@ impl AnalyticsEventsClient {
&self,
connection_id: u64,
request_id: RequestId,
response: ClientResponsePayload,
response: &ClientResponsePayload,
thread_originator: Option<String>,
) {
if !matches!(
@@ -522,10 +522,13 @@ impl AnalyticsEventsClient {
) {
return;
}
if serde_json::to_writer(std::io::sink(), response).is_err() {
return;
}
self.record_fact(AnalyticsFact::ClientResponse {
connection_id,
request_id,
response: Box::new(response),
response: Box::new(response.clone()),
thread_originator,
});
}
+23 -2
View File
@@ -614,7 +614,7 @@ fn track_response_only_enqueues_analytics_relevant_responses() {
(RequestId::Integer(4), sample_turn_start_response()),
(RequestId::Integer(5), sample_turn_steer_response()),
] {
client.track_response(/*connection_id*/ 7, request_id, response);
client.track_response(/*connection_id*/ 7, request_id, &response);
assert!(matches!(
receiver.try_recv(),
Ok(AnalyticsEventsQueueMessage::Fact(input))
@@ -625,11 +625,32 @@ fn track_response_only_enqueues_analytics_relevant_responses() {
client.track_response(
/*connection_id*/ 7,
RequestId::Integer(6),
ClientResponsePayload::ThreadArchive(ThreadArchiveResponse {}),
&ClientResponsePayload::ThreadArchive(ThreadArchiveResponse {}),
);
assert!(matches!(receiver.try_recv(), Err(TryRecvError::Empty)));
}
#[cfg(unix)]
#[test]
fn track_response_ignores_unserializable_thread_responses() {
use std::ffi::OsString;
use std::os::unix::ffi::OsStringExt;
let (client, mut receiver) = client_with_receiver();
let mut response = sample_thread_start_response();
let ClientResponsePayload::ThreadStart(thread_start) = &mut response else {
panic!("expected thread/start response");
};
thread_start.cwd = codex_utils_absolute_path::AbsolutePathBuf::from_absolute_path(
std::path::PathBuf::from(OsString::from_vec(vec![b'/', b'b', b'a', b'd', 0xff])),
)
.expect("non-UTF-8 Unix paths are valid absolute paths");
client.track_response(/*connection_id*/ 7, RequestId::Integer(1), &response);
assert!(matches!(receiver.try_recv(), Err(TryRecvError::Empty)));
}
#[tokio::test]
async fn flush_waits_for_preceding_fact_delivery() {
let (client, mut receiver) = client_with_receiver();
@@ -315,7 +315,8 @@ macro_rules! client_request_definitions {
}
}
#[derive(Debug, Clone)]
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
#[allow(clippy::large_enum_variant)]
pub enum ClientResponsePayload {
$( $variant($response), )*
@@ -323,26 +324,6 @@ macro_rules! client_request_definitions {
}
impl ClientResponsePayload {
pub fn into_jsonrpc_parts_and_payload(
self,
request_id: RequestId,
) -> std::result::Result<
(RequestId, crate::Result, Option<ClientResponsePayload>),
serde_json::Error,
> {
match self {
$(
Self::$variant(response) => {
let result = serde_json::to_value(&response)?;
Ok((request_id, result, Some(Self::$variant(response))))
}
)*
Self::InterruptConversation(response) => {
serde_json::to_value(response).map(|result| (request_id, result, None))
}
}
}
pub fn into_client_response(self, request_id: RequestId) -> Option<ClientResponse> {
match self {
$(
@@ -5,18 +5,13 @@ use pretty_assertions::assert_eq;
use serde_json::json;
#[test]
fn client_response_payload_returns_jsonrpc_parts_and_client_response() -> Result<()> {
let (request_id, result, payload) =
ClientResponsePayload::ThreadArchive(v2::ThreadArchiveResponse {})
.into_jsonrpc_parts_and_payload(RequestId::Integer(7))?;
assert_eq!(request_id, RequestId::Integer(7));
assert_eq!(result, json!({}));
fn client_response_payload_serializes_without_an_intermediate_json_value() -> Result<()> {
let payload = ClientResponsePayload::ThreadArchive(v2::ThreadArchiveResponse {});
assert_eq!(serde_json::to_string(&payload)?, "{}");
let Some(ClientResponse::ThreadArchive {
request_id,
response: _,
}) = payload.and_then(|payload| payload.into_client_response(RequestId::Integer(7)))
}) = payload.into_client_response(RequestId::Integer(7))
else {
panic!("expected thread/archive client response");
};
@@ -26,19 +21,19 @@ fn client_response_payload_returns_jsonrpc_parts_and_client_response() -> Result
#[test]
fn interrupt_conversation_payload_stays_jsonrpc_only() -> Result<()> {
let (request_id, result, payload) =
ClientResponsePayload::InterruptConversation(v1::InterruptConversationResponse {
abort_reason: TurnAbortReason::Interrupted,
})
.into_jsonrpc_parts_and_payload(RequestId::Integer(8))?;
assert_eq!(request_id, RequestId::Integer(8));
let payload = ClientResponsePayload::InterruptConversation(v1::InterruptConversationResponse {
abort_reason: TurnAbortReason::Interrupted,
});
assert_eq!(
result,
serde_json::to_value(&payload)?,
json!({
"abortReason": "interrupted",
})
);
assert!(payload.is_none());
assert!(
payload
.into_client_response(RequestId::Integer(8))
.is_none()
);
Ok(())
}
@@ -1,8 +1,8 @@
use std::fmt;
use codex_app_server_protocol::ClientResponsePayload;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::Result;
use codex_app_server_protocol::ServerNotificationEnvelope;
use codex_app_server_protocol::ServerRequest;
use serde::Serialize;
@@ -31,10 +31,10 @@ pub enum OutgoingMessage {
Error(OutgoingError),
}
#[derive(Debug, Clone, PartialEq, Serialize)]
#[derive(Debug, Clone, Serialize)]
pub struct OutgoingResponse {
pub id: RequestId,
pub result: Result,
pub result: Box<ClientResponsePayload>,
}
#[derive(Debug, Clone, PartialEq, Serialize)]
@@ -6,6 +6,7 @@ use crate::outgoing_message::OutgoingMessage;
use crate::outgoing_message::QueuedOutgoingMessage;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::RequestId;
use codex_core::config::find_codex_home;
use codex_utils_absolute_path::AbsolutePathBuf;
use std::net::SocketAddr;
@@ -47,6 +48,7 @@ pub use unix_socket::prepare_control_socket_path;
pub use unix_socket::start_control_socket_acceptor;
pub use websocket::start_websocket_acceptor;
const INTERNAL_ERROR_CODE: i64 = -32603;
const OVERLOADED_ERROR_CODE: i64 = -32001;
const APP_SERVER_CONTROL_SOCKET_DIR_NAME: &str = "app-server-control";
@@ -260,14 +262,35 @@ fn serialize_outgoing_message(outgoing_message: OutgoingMessage) -> Option<Strin
Ok(json) => Some(json),
Err(err) => {
error!("Failed to serialize JSONRPCMessage: {err}");
None
let OutgoingMessage::Response(response) = outgoing_message else {
return None;
};
serde_json::to_string(&response_serialization_error(response.id, err))
.inspect_err(|err| error!("Failed to serialize JSONRPC error: {err}"))
.ok()
}
}
}
fn response_serialization_error(
request_id: RequestId,
err: impl std::fmt::Display,
) -> OutgoingMessage {
OutgoingMessage::Error(OutgoingError {
id: request_id,
error: JSONRPCErrorError {
code: INTERNAL_ERROR_CODE,
message: format!("failed to serialize response: {err}"),
data: None,
},
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::outgoing_message::OutgoingResponse;
use codex_app_server_protocol::ClientResponsePayload;
use codex_app_server_protocol::ConfigWarningNotification;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
@@ -275,6 +298,7 @@ mod tests {
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::ServerNotificationEnvelope;
use codex_app_server_protocol::ThreadArchiveResponse;
use pretty_assertions::assert_eq;
use serde_json::json;
use tokio::time::Duration;
@@ -314,6 +338,60 @@ mod tests {
);
}
#[test]
fn serialize_typed_response_preserves_wire_shape() {
let message = OutgoingMessage::Response(OutgoingResponse {
id: RequestId::Integer(7),
result: Box::new(ClientResponsePayload::ThreadArchive(
ThreadArchiveResponse {},
)),
});
let json = serialize_outgoing_message(message).expect("message should serialize");
assert_eq!(
serde_json::from_str::<serde_json::Value>(&json).expect("message should be valid JSON"),
json!({ "id": 7, "result": {} })
);
}
#[cfg(unix)]
#[test]
fn serialize_invalid_typed_response_returns_jsonrpc_error() {
use std::ffi::OsString;
use std::os::unix::ffi::OsStringExt;
use std::path::PathBuf;
let codex_home =
AbsolutePathBuf::from_absolute_path(PathBuf::from(OsString::from_vec(vec![
b'/', b'b', b'a', b'd', 0xff,
])))
.expect("non-UTF-8 Unix paths are valid absolute paths");
let message = OutgoingMessage::Response(OutgoingResponse {
id: RequestId::Integer(7),
result: Box::new(ClientResponsePayload::Initialize(
codex_app_server_protocol::InitializeResponse {
user_agent: "codex-test-agent".to_string(),
codex_home,
platform_family: "unix".to_string(),
platform_os: "linux".to_string(),
},
)),
});
let json = serialize_outgoing_message(message)
.expect("invalid response should serialize as a JSON-RPC error");
assert_eq!(
serde_json::from_str::<serde_json::Value>(&json).expect("message should be valid JSON"),
json!({
"id": 7,
"error": {
"code": -32603,
"message": "failed to serialize response: path contains invalid UTF-8 characters",
}
})
);
}
#[tokio::test]
async fn enqueue_incoming_request_returns_overload_error_when_queue_is_full() {
let connection_id = ConnectionId(42);
@@ -4,6 +4,8 @@ use super::protocol::ClientId;
use super::protocol::ServerEnvelope;
use super::protocol::ServerEvent;
use super::protocol::StreamId;
use crate::outgoing_message::OutgoingMessage;
use crate::transport::response_serialization_error;
use base64::DecodeSliceError;
use base64::Engine;
use codex_app_server_protocol::JSONRPCMessage;
@@ -303,7 +305,25 @@ pub(super) fn split_server_envelope_for_transport(
return Ok(vec![envelope]);
}
let envelope_size_bytes = serialized_len(&envelope)?;
let envelope_size_bytes = match serialized_len(&envelope) {
Ok(envelope_size_bytes) => envelope_size_bytes,
Err(err) => {
let ServerEvent::ServerMessage { message } = envelope.event else {
unreachable!("server message variant checked above");
};
let OutgoingMessage::Response(response) = *message else {
return Err(err);
};
return Ok(vec![ServerEnvelope {
event: ServerEvent::ServerMessage {
message: Box::new(response_serialization_error(response.id, err)),
},
client_id: envelope.client_id,
stream_id: envelope.stream_id,
seq_id: envelope.seq_id,
}]);
}
};
if envelope_size_bytes <= REMOTE_CONTROL_SEGMENT_MAX_BYTES {
return Ok(vec![envelope]);
}
@@ -9,13 +9,25 @@ use super::segment::ClientSegmentReassembler;
use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES;
use super::segment::split_server_envelope_for_transport;
use crate::outgoing_message::OutgoingMessage;
#[cfg(unix)]
use crate::outgoing_message::OutgoingResponse;
use base64::Engine;
#[cfg(unix)]
use codex_app_server_protocol::ClientResponsePayload;
use codex_app_server_protocol::ConfigWarningNotification;
#[cfg(unix)]
use codex_app_server_protocol::InitializeResponse;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
#[cfg(unix)]
use codex_app_server_protocol::RequestId;
use codex_app_server_protocol::ServerNotification;
use codex_app_server_protocol::ServerNotificationEnvelope;
#[cfg(unix)]
use codex_utils_absolute_path::AbsolutePathBuf;
use pretty_assertions::assert_eq;
#[cfg(unix)]
use serde_json::json;
#[test]
fn reassembles_client_message_chunks() {
@@ -108,6 +120,54 @@ fn splits_large_server_messages_into_wire_chunks() {
}));
}
#[cfg(unix)]
#[test]
fn invalid_response_becomes_remote_control_jsonrpc_error() {
use std::ffi::OsString;
use std::os::unix::ffi::OsStringExt;
use std::path::PathBuf;
let codex_home = AbsolutePathBuf::from_absolute_path(PathBuf::from(OsString::from_vec(vec![
b'/', b'b', b'a', b'd', 0xff,
])))
.expect("non-UTF-8 Unix paths are valid absolute paths");
let envelope = ServerEnvelope {
event: ServerEvent::ServerMessage {
message: Box::new(OutgoingMessage::Response(OutgoingResponse {
id: RequestId::Integer(7),
result: Box::new(ClientResponsePayload::Initialize(InitializeResponse {
user_agent: "codex-test-agent".to_string(),
codex_home,
platform_family: "unix".to_string(),
platform_os: "linux".to_string(),
})),
})),
},
client_id: ClientId("client-1".to_string()),
stream_id: StreamId("stream-1".to_string()),
seq_id: 9,
};
let envelopes = split_server_envelope_for_transport(envelope)
.expect("invalid response should become a remote-control JSON-RPC error");
assert_eq!(
serde_json::to_value(envelopes).expect("error envelope should serialize"),
json!([{
"type": "server_message",
"client_id": "client-1",
"stream_id": "stream-1",
"seq_id": 9,
"message": {
"id": 7,
"error": {
"code": -32603,
"message": "failed to serialize response: path contains invalid UTF-8 characters",
}
}
}])
);
}
#[test]
fn invalidates_incomplete_stream_assemblies() {
let message = JSONRPCMessage::Notification(JSONRPCNotification {
@@ -1619,9 +1619,16 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
.send(QueuedOutgoingMessage::new(OutgoingMessage::Response(
crate::outgoing_message::OutgoingResponse {
id: codex_app_server_protocol::RequestId::Integer(11),
result: json!({
"userAgent": "codex-test-agent"
}),
result: Box::new(
codex_app_server_protocol::ClientResponsePayload::Initialize(
codex_app_server_protocol::InitializeResponse {
user_agent: "codex-test-agent".to_string(),
codex_home: codex_home.path().abs(),
platform_family: "test-family".to_string(),
platform_os: "test-os".to_string(),
},
),
),
},
)))
.await
@@ -1636,6 +1643,9 @@ async fn remote_control_http_mode_enrolls_before_connecting() {
"id": 11,
"result": {
"userAgent": "codex-test-agent",
"codexHome": codex_home.path(),
"platformFamily": "test-family",
"platformOs": "test-os",
}
}
})
+10 -4
View File
@@ -885,8 +885,11 @@ mod tests {
panic!("expected execution response after termination");
};
assert_eq!(response.id, request_id.request_id);
let response: CommandExecResponse =
serde_json::from_value(response.result).expect("deserialize command/exec response");
let codex_app_server_protocol::ClientResponsePayload::OneOffCommandExec(response) =
*response.result
else {
panic!("expected command/exec response");
};
assert_ne!(response.exit_code, 0);
assert_eq!(response.stdout, "");
// The deferred response now drains any already-emitted stderr before
@@ -961,8 +964,11 @@ mod tests {
panic!("expected execution response after cancellation");
};
assert_eq!(response.id, request_id.request_id);
let response: CommandExecResponse =
serde_json::from_value(response.result).expect("deserialize command/exec response");
let codex_app_server_protocol::ClientResponsePayload::OneOffCommandExec(response) =
*response.result
else {
panic!("expected command/exec response");
};
assert_ne!(response.exit_code, EXEC_TIMEOUT_EXIT_CODE);
}
+4 -1
View File
@@ -641,7 +641,10 @@ async fn start_uninitialized(args: InProcessStartArgs) -> IoResult<InProcessClie
match outgoing_message {
OutgoingMessage::Response(response) => {
if let Some(response_tx) = pending_request_responses.remove(&response.id) {
let _ = response_tx.send(Ok(response.result));
let result = serde_json::to_value(response.result).map_err(|err| {
internal_error(format!("failed to serialize response: {err}"))
});
let _ = response_tx.send(result);
} else {
warn!(
request_id = ?response.id,
@@ -450,8 +450,10 @@ async fn read_response<T: serde::de::DeserializeOwned>(
if response.id != RequestId::Integer(request_id) {
continue;
}
return serde_json::from_value(response.result)
.expect("response payload should deserialize");
return serde_json::from_value(
serde_json::to_value(response.result).expect("response payload should serialize"),
)
.expect("response payload should deserialize");
}
}
+32 -45
View File
@@ -538,53 +538,37 @@ impl OutgoingMessageSender {
) {
let connection_id = request_id.connection_id;
let request_id_for_analytics = request_id.request_id.clone();
let serialized_response = response
.into_jsonrpc_parts_and_payload(request_id.request_id.clone())
.map(|(id, result, response)| {
if let Some(response) = response {
match thread_originator {
Some(thread_originator) => {
self.analytics_events_client
.track_response_with_thread_originator(
connection_id.0,
request_id_for_analytics,
response,
thread_originator,
);
}
None => {
self.analytics_events_client.track_response(
connection_id.0,
request_id_for_analytics,
response,
);
}
}
}
(id, result)
});
let request_context = self.take_request_context(&request_id).await;
match serialized_response {
Ok((id, result)) => {
let outgoing_message = OutgoingMessage::Response(OutgoingResponse { id, result });
self.send_outgoing_message_to_connection(
request_context,
connection_id,
outgoing_message,
"response",
)
.await;
match thread_originator {
Some(thread_originator) => {
self.analytics_events_client
.track_response_with_thread_originator(
connection_id.0,
request_id_for_analytics,
&response,
thread_originator,
);
}
Err(err) => {
self.send_error_inner(
request_context,
request_id,
internal_error(format!("failed to serialize response: {err}")),
)
.await;
None => {
self.analytics_events_client.track_response(
connection_id.0,
request_id_for_analytics,
&response,
);
}
}
let response = Box::new(response);
let request_context = self.take_request_context(&request_id).await;
let outgoing_message = OutgoingMessage::Response(OutgoingResponse {
id: request_id.request_id,
result: response,
});
self.send_outgoing_message_to_connection(
request_context,
connection_id,
outgoing_message,
"response",
)
.await;
}
pub(crate) async fn send_server_notification(&self, notification: ServerNotification) {
@@ -1087,7 +1071,10 @@ mod tests {
panic!("expected response message");
};
assert_eq!(response.id, request_id.request_id);
assert_eq!(response.result, json!({}));
assert_eq!(
serde_json::to_value(response.result).expect("result should serialize"),
json!({})
);
}
other => panic!("expected targeted response envelope, got: {other:?}"),
}
+22 -27
View File
@@ -5,7 +5,7 @@ use std::io::SeekFrom;
use serde::de::DeserializeOwned;
const READ_CHUNK_SIZE: usize = 8 * 1024;
const READ_CHUNK_SIZE: usize = 64 * 1024;
#[derive(Debug)]
pub enum ScanOutcome<T> {
@@ -64,39 +64,34 @@ where
T: DeserializeOwned,
{
loop {
let Some(byte) = self.read_previous_byte()? else {
return Ok(self.finish_record());
};
if self.chunk_position == 0 {
if self.next_chunk_end == 0 {
return Ok(self.finish_record());
}
if byte != b'\n' {
self.record_reversed.push(byte);
continue;
let read_size = usize::try_from(self.next_chunk_end.min(READ_CHUNK_SIZE as u64))
.map_err(io::Error::other)?;
self.next_chunk_end -= read_size as u64;
self.reader.seek(SeekFrom::Start(self.next_chunk_end))?;
self.reader.read_exact(&mut self.chunk[..read_size])?;
self.chunk_position = read_size;
}
if let Some(outcome) = self.finish_record() {
return Ok(Some(outcome));
let chunk = &self.chunk[..self.chunk_position];
if let Some(newline_position) = chunk.iter().rposition(|byte| *byte == b'\n') {
self.record_reversed
.extend(chunk[newline_position + 1..].iter().rev().copied());
self.chunk_position = newline_position;
if let Some(outcome) = self.finish_record() {
return Ok(Some(outcome));
}
} else {
self.record_reversed.extend(chunk.iter().rev().copied());
self.chunk_position = 0;
}
}
}
fn read_previous_byte(&mut self) -> io::Result<Option<u8>> {
if self.chunk_position == 0 {
if self.next_chunk_end == 0 {
return Ok(None);
}
let read_size = usize::try_from(self.next_chunk_end.min(READ_CHUNK_SIZE as u64))
.map_err(io::Error::other)?;
self.next_chunk_end -= read_size as u64;
self.reader.seek(SeekFrom::Start(self.next_chunk_end))?;
self.reader.read_exact(&mut self.chunk[..read_size])?;
self.chunk_position = read_size;
}
self.chunk_position -= 1;
Ok(Some(self.chunk[self.chunk_position]))
}
fn finish_record<T>(&mut self) -> Option<ScanOutcome<T>>
where
T: DeserializeOwned,