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<Vec<u8>>` 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-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

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<u8>);
+pub struct ByteChunk(#[serde(with = "base64_bytes")] pub Arc<Vec<u8>>);
 
 impl ByteChunk {
     pub fn into_inner(self) -> Vec<u8> {
-        self.0
+        Arc::unwrap_or_clone(self.0)
     }
 }
 
 impl From<Vec<u8>> for ByteChunk {
     fn from(value: Vec<u8>) -> 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<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -1201,13 +1209,14 @@ mod base64_bytes {
         serializer.serialize_str(&BASE64_STANDARD.encode(bytes))
     }
 
-    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
+    pub fn deserialize<'de, D>(deserializer: D) -> Result<Arc<Vec<u8>>, 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::<super::ByteChunk>(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 =

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<u8>,
-}
-
 struct RunningProcess {
     session: ExecCommandSession,
     tty: bool,
     pipe_stdin: bool,
     accepted_stdin_write_ids: Arc<Mutex<AcceptedStdinWriteIds>>,
-    output: VecDeque<RetainedOutputChunk>,
+    output: VecDeque<ProcessOutputChunk>,
     retained_bytes: usize,
     next_seq: u64,
     exit_code: Option<i32>,
@@ -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;

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 => {

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));

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);
                 }

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)
             };