From a933dd77dbe101d7bd746ea3c7d1f8174eca4a05 Mon Sep 17 00:00:00 2001 From: jif Date: Thu, 1 Oct 2026 12:32:22 +0000 Subject: [PATCH] Share byte buffers across exec-server output chunks (#49972) ## Why Retaining process output, publishing events, and serving reads cloned the underlying byte buffers. Sharing those buffers avoids copying output bytes when chunks are cloned. ## What changed Store `ByteChunk` bytes in `Arc>` and retain `ProcessOutputChunk` values directly. Add borrowed byte access for the MCP process transport while preserving base64 serialization and owned byte extraction through `into_inner()`. ## Testing Add a unit test covering base64 serialization, deserialization, and independent owned bytes extracted from shared chunks. GitOrigin-RevId: 6bba114af48424aa844487aa75163cfc604e0575 --- .../codex-mcp/src/trusted_access_tests.rs | 4 +- codex-rs/exec-server-protocol/src/protocol.rs | 35 ++++++++++++-- codex-rs/exec-server/src/local_process.rs | 46 ++++++------------- .../exec-server/tests/capability_discovery.rs | 9 +++- codex-rs/exec-server/tests/forward.rs | 2 +- .../src/executor_process_transport.rs | 6 +-- .../rmcp-client/src/http_headers_tests.rs | 2 +- 7 files changed, 60 insertions(+), 44 deletions(-) diff --git a/codex-rs/codex-mcp/src/trusted_access_tests.rs b/codex-rs/codex-mcp/src/trusted_access_tests.rs index 63bc1d9d4a..a9a2a88be6 100644 --- a/codex-rs/codex-mcp/src/trusted_access_tests.rs +++ b/codex-rs/codex-mcp/src/trusted_access_tests.rs @@ -58,7 +58,7 @@ impl HttpClient for RecordingHttpClient { let response = HttpRequestResponse { status: self.status, headers: Vec::new(), - body: ByteChunk(self.response.clone()), + body: ByteChunk::from(self.response.clone()), }; async move { if let Some((requested, release)) = &self.response_gate { @@ -78,7 +78,7 @@ impl HttpClient for RecordingHttpClient { let response = HttpRequestResponse { status: self.status, headers: Vec::new(), - body: ByteChunk(Vec::new()), + body: ByteChunk::from(Vec::new()), }; let chunks = self .response_chunks diff --git a/codex-rs/exec-server-protocol/src/protocol.rs b/codex-rs/exec-server-protocol/src/protocol.rs index 738cd26144..641f7e700a 100644 --- a/codex-rs/exec-server-protocol/src/protocol.rs +++ b/codex-rs/exec-server-protocol/src/protocol.rs @@ -60,19 +60,26 @@ pub const HTTP_REQUEST_BODY_DELTA_METHOD: &str = "http/request/bodyDelta"; /// Maximum decoded response-body bytes carried by one streamed HTTP notification. pub const MAX_HTTP_BODY_DELTA_BYTES: usize = 1024 * 1024; +/// Shared immutable bytes, encoded as a base64 string on the wire. #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] #[serde(transparent)] -pub struct ByteChunk(#[serde(with = "base64_bytes")] pub Vec); +pub struct ByteChunk(#[serde(with = "base64_bytes")] pub Arc>); impl ByteChunk { pub fn into_inner(self) -> Vec { - self.0 + Arc::unwrap_or_clone(self.0) } } impl From> for ByteChunk { fn from(value: Vec) -> Self { - Self(value) + Self(Arc::new(value)) + } +} + +impl AsRef<[u8]> for ByteChunk { + fn as_ref(&self) -> &[u8] { + self.0.as_slice() } } @@ -1193,6 +1200,7 @@ mod base64_bytes { use serde::Deserialize; use serde::Deserializer; use serde::Serializer; + use std::sync::Arc; pub fn serialize(bytes: &[u8], serializer: S) -> Result where @@ -1201,13 +1209,14 @@ mod base64_bytes { serializer.serialize_str(&BASE64_STANDARD.encode(bytes)) } - pub fn deserialize<'de, D>(deserializer: D) -> Result, D::Error> + pub fn deserialize<'de, D>(deserializer: D) -> Result>, D::Error> where D: Deserializer<'de>, { let encoded = String::deserialize(deserializer)?; BASE64_STANDARD .decode(encoded) + .map(Arc::new) .map_err(serde::de::Error::custom) } } @@ -1264,6 +1273,24 @@ mod tests { use pretty_assertions::assert_eq; use std::collections::HashMap; + #[test] + fn shared_byte_chunks_preserve_base64_and_owned_bytes() { + let chunk = super::ByteChunk::from(vec![0, 255, 10]); + let shared = chunk.clone(); + assert_eq!(serde_json::to_string(&shared).unwrap(), r#""AP8K""#); + assert_eq!( + serde_json::from_str::(r#""AP8K""#).unwrap(), + chunk, + ); + + let mut owned = shared.into_inner(); + owned[0] = 1; + assert_eq!( + (owned, chunk.into_inner()), + (vec![1, 255, 10], vec![0, 255, 10]) + ); + } + #[test] fn exec_params_keeps_proxy_launch_separate_from_sandbox_facts() { let cwd = diff --git a/codex-rs/exec-server/src/local_process.rs b/codex-rs/exec-server/src/local_process.rs index 85cbd9bd45..f384d84d88 100644 --- a/codex-rs/exec-server/src/local_process.rs +++ b/codex-rs/exec-server/src/local_process.rs @@ -47,6 +47,7 @@ use crate::network_policy_decisions::network_policy_decider; use crate::process::ExecProcessEventLog; use crate::process::sandbox_type_from_protocol; use crate::process_sandbox::prepare_exec_request_with_telemetry; +use crate::protocol::ByteChunk; use crate::protocol::EXEC_CLOSED_METHOD; use crate::protocol::ExecClosedNotification; use crate::protocol::ExecEnvPolicy; @@ -95,19 +96,12 @@ const EXITED_PROCESS_RETENTION: Duration = Duration::from_millis(25); #[cfg(not(test))] const EXITED_PROCESS_RETENTION: Duration = Duration::from_secs(30); -#[derive(Clone)] -struct RetainedOutputChunk { - seq: u64, - stream: ExecOutputStream, - chunk: Vec, -} - struct RunningProcess { session: ExecCommandSession, tty: bool, pipe_stdin: bool, accepted_stdin_write_ids: Arc>, - output: VecDeque, + output: VecDeque, retained_bytes: usize, next_seq: u64, exit_code: Option, @@ -578,16 +572,12 @@ impl LocalProcess { let mut total_bytes = 0; let mut next_seq = process.next_seq; for retained in process.output.iter().filter(|chunk| chunk.seq > after_seq) { - let chunk_len = retained.chunk.len(); + let chunk_len = retained.chunk.0.len(); if !chunks.is_empty() && total_bytes + chunk_len > max_bytes { break; } total_bytes += chunk_len; - chunks.push(ProcessOutputChunk { - seq: retained.seq, - stream: retained.stream, - chunk: retained.chunk.clone().into(), - }); + chunks.push(retained.clone()); next_seq = retained.seq + 1; if total_bytes >= max_bytes { break; @@ -984,6 +974,7 @@ async fn stream_output( ) { while let Some(chunk) = receiver.recv().await { let _chunk_len = chunk.len(); + let chunk = ByteChunk::from(chunk); let notification = { let mut processes = inner.processes.lock().await; let Some(entry) = processes.get_mut(&process_id) else { @@ -994,26 +985,19 @@ async fn stream_output( }; let seq = process.next_seq; process.next_seq += 1; - process.retained_bytes += chunk.len(); - process.output.push_back(RetainedOutputChunk { - seq, - stream, - chunk: chunk.clone(), - }); + process.retained_bytes += chunk.0.len(); + let output = ProcessOutputChunk { seq, stream, chunk }; + process.output.push_back(output.clone()); while process.retained_bytes > RETAINED_OUTPUT_BYTES_PER_PROCESS || process.output.len() > RETAINED_OUTPUT_CHUNKS_PER_PROCESS { let Some(evicted) = process.output.pop_front() else { break; }; - process.retained_bytes = process.retained_bytes.saturating_sub(evicted.chunk.len()); + process.retained_bytes = + process.retained_bytes.saturating_sub(evicted.chunk.0.len()); } let _ = process.wake_tx.send(seq); - let output = ProcessOutputChunk { - seq, - stream, - chunk: chunk.into(), - }; process .events .publish(ExecProcessEvent::Output(output.clone())); @@ -1087,11 +1071,11 @@ fn watch_exit( for chunk in &process.output { match chunk.stream { ExecOutputStream::Stdout | ExecOutputStream::Pty => { - stdout.extend_from_slice(&chunk.chunk); + stdout.extend_from_slice(&chunk.chunk.0); } - ExecOutputStream::Stderr => stderr.extend_from_slice(&chunk.chunk), + ExecOutputStream::Stderr => stderr.extend_from_slice(&chunk.chunk.0), } - aggregated.extend_from_slice(&chunk.chunk); + aggregated.extend_from_slice(&chunk.chunk.0); } let exec_output = ExecToolCallOutput { exit_code, @@ -1734,10 +1718,10 @@ mod tests { panic!("process should be running"); }; running.output = (1..=retained_chunk_count) - .map(|seq| RetainedOutputChunk { + .map(|seq| ProcessOutputChunk { seq, stream: ExecOutputStream::Stdout, - chunk: vec![b'x'], + chunk: vec![b'x'].into(), }) .collect(); running.retained_bytes = RETAINED_OUTPUT_CHUNKS_PER_PROCESS; diff --git a/codex-rs/exec-server/tests/capability_discovery.rs b/codex-rs/exec-server/tests/capability_discovery.rs index 9044555e0d..a5d777c24b 100644 --- a/codex-rs/exec-server/tests/capability_discovery.rs +++ b/codex-rs/exec-server/tests/capability_discovery.rs @@ -744,7 +744,12 @@ async fn executor_legacy_exec_uses_process_cwd_for_relative_denials() -> anyhow: anyhow::bail!("expected legacy process output, got {response:?}"); }; let response: ReadResponse = serde_json::from_value(response.result)?; - output.extend(response.chunks.into_iter().flat_map(|chunk| chunk.chunk.0)); + output.extend( + response + .chunks + .into_iter() + .flat_map(|chunk| chunk.chunk.into_inner()), + ); after_seq = response.next_seq.checked_sub(1).or(after_seq); if response.closed { break ( @@ -873,7 +878,7 @@ async fn executor_legacy_filesystem_cwd_keeps_absolute_read_allow_and_deny_rules anyhow::bail!("expected to read the legacy stream, got {block:?}"); }; let block: FsReadBlockResponse = serde_json::from_value(block.result)?; - assert_eq!(block.chunk.0, b"contents"); + assert_eq!(block.chunk.into_inner(), b"contents"); } } JSONRPCMessage::Error(error) if !permitted => { diff --git a/codex-rs/exec-server/tests/forward.rs b/codex-rs/exec-server/tests/forward.rs index 526bd12119..f7854db669 100644 --- a/codex-rs/exec-server/tests/forward.rs +++ b/codex-rs/exec-server/tests/forward.rs @@ -94,7 +94,7 @@ async fn forwarder_runs_commands_and_transfers_files() -> Result<()> { assert_eq!(response.failure, None); for chunk in response.chunks { assert_eq!(chunk.stream, ExecOutputStream::Stdout); - output.extend(chunk.chunk.0); + output.extend(chunk.chunk.into_inner()); } if response.closed { break Ok::<_, ExecServerError>((output, response.exit_code)); diff --git a/codex-rs/rmcp-client/src/executor_process_transport.rs b/codex-rs/rmcp-client/src/executor_process_transport.rs index b1fddbdf75..45d565b1cc 100644 --- a/codex-rs/rmcp-client/src/executor_process_transport.rs +++ b/codex-rs/rmcp-client/src/executor_process_transport.rs @@ -405,20 +405,20 @@ impl ExecutorProcessTransport { } fn push_process_output(&mut self, chunk: ProcessOutputChunk) { - let bytes = chunk.chunk.into_inner(); + let bytes = chunk.chunk.as_ref(); match chunk.stream { // MCP stdio uses stdout as the protocol stream. PTY output is // accepted defensively because the executor process API has a // unified stream enum, but remote MCP starts with `tty=false`. ExecOutputStream::Stdout | ExecOutputStream::Pty => { - if let Err(error) = self.stdout.extend_from_slice(&bytes) { + if let Err(error) = self.stdout.extend_from_slice(bytes) { self.close_for_oversized_line("stdout", error); } } // Stderr is intentionally out-of-band. It should help debug server // startup failures without entering rmcp framing. ExecOutputStream::Stderr => { - if let Err(error) = self.push_stderr(&bytes) { + if let Err(error) = self.push_stderr(bytes) { self.stdout.clear(); self.close_for_oversized_line("stderr", error); } diff --git a/codex-rs/rmcp-client/src/http_headers_tests.rs b/codex-rs/rmcp-client/src/http_headers_tests.rs index c27c5d1002..292e1b6410 100644 --- a/codex-rs/rmcp-client/src/http_headers_tests.rs +++ b/codex-rs/rmcp-client/src/http_headers_tests.rs @@ -440,7 +440,7 @@ async fn helper_refresh_preserves_oauth_challenges_and_retries_at_most_once() { .http_request(params) .await .expect("original OAuth response"); - let bytes = response.body.0.clone(); + let bytes = response.body.clone().into_inner(); (response, bytes) };