Share stored JSON payloads across code mode cells with Arc (#47036)

## Why

Creating a cell deep-copies every stored JSON value, including values the cell never loads. Storing a value also duplicates its payload between the cell's local state and pending writes.

## What changed

Use `Arc<JsonValue>` throughout the session store, cell runtime, and completion writes to share immutable payloads. Keep a separate map snapshot for each cell so later commits cannot change its view of stored values.

GitOrigin-RevId: 4535f0f2ed0c8f8cbfb8d24cf1cb44cbc74472b5

codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs

@@ -39,7 +39,7 @@ impl CellHost for PanickingCallbackHost {
 
     async fn commit_completion(
         &self,
-        _stored_value_writes: HashMap<String, JsonValue>,
+        _stored_value_writes: HashMap<String, Arc<JsonValue>>,
         _event: CellEvent,
         _pending_initial_yield_items: Option<Vec<crate::session_runtime::OutputItem>>,
         _cell_state: Arc<CellState>,

codex-rs/code-mode-runtime/src/cell_actor/mod.rs

@@ -47,7 +47,7 @@ pub(crate) struct CellActor;
 impl CellActor {
     pub(crate) fn prepare<H: CellHost>(
         request: CellRequest,
-        stored_values: HashMap<String, JsonValue>,
+        stored_values: HashMap<String, Arc<JsonValue>>,
         host: Arc<H>,
         initial_observe_mode: ObserveMode,
         cell_state: Arc<CellState>,

codex-rs/code-mode-runtime/src/cell_actor/tests.rs

@@ -43,7 +43,7 @@ impl CellHost for TestHost {
 
     async fn commit_completion(
         &self,
-        _stored_value_writes: HashMap<String, JsonValue>,
+        _stored_value_writes: HashMap<String, Arc<JsonValue>>,
         event: CellEvent,
         pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,
@@ -75,7 +75,7 @@ impl CellHost for RecordingHost {
 
     async fn commit_completion(
         &self,
-        _stored_value_writes: HashMap<String, JsonValue>,
+        _stored_value_writes: HashMap<String, Arc<JsonValue>>,
         event: CellEvent,
         pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,

codex-rs/code-mode-runtime/src/cell_actor/types.rs

@@ -53,7 +53,7 @@ pub(crate) trait CellHost: Send + Sync + 'static {
 
     fn commit_completion(
         &self,
-        stored_value_writes: HashMap<String, JsonValue>,
+        stored_value_writes: HashMap<String, Arc<JsonValue>>,
         event: CellEvent,
         pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,

codex-rs/code-mode-runtime/src/runtime/callbacks.rs

@@ -1,4 +1,5 @@
 use codex_code_mode_protocol::FunctionCallOutputContentItem;
+use std::sync::Arc;
 
 use super::EXIT_SENTINEL;
 use super::RuntimeEvent;
@@ -230,7 +231,10 @@ pub(super) fn store_callback(
         }
     };
     if let Some(state) = scope.get_slot_mut::<RuntimeState>() {
-        state.stored_values.insert(key.clone(), serialized.clone());
+        let serialized = Arc::new(serialized);
+        state
+            .stored_values
+            .insert(key.clone(), Arc::clone(&serialized));
         state.stored_value_writes.insert(key, serialized);
     }
 }

codex-rs/code-mode-runtime/src/runtime/mod.rs

@@ -8,6 +8,7 @@ mod value;
 use std::collections::HashMap;
 use std::panic::AssertUnwindSafe;
 use std::panic::catch_unwind;
+use std::sync::Arc;
 use std::sync::mpsc as std_mpsc;
 use std::thread;
 
@@ -65,14 +66,14 @@ pub(crate) enum RuntimeEvent {
         text: String,
     },
     Result {
-        stored_value_writes: HashMap<String, JsonValue>,
+        stored_value_writes: HashMap<String, Arc<JsonValue>>,
         error_text: Option<String>,
     },
     ThreadPanicked,
 }
 
 pub(crate) fn spawn_runtime(
-    stored_values: HashMap<String, JsonValue>,
+    stored_values: HashMap<String, Arc<JsonValue>>,
     request: ExecuteRequest,
     event_tx: mpsc::UnboundedSender<RuntimeEvent>,
     pending_mode: PendingRuntimeMode,
@@ -143,15 +144,15 @@ struct RuntimeConfig {
     tool_call_id: String,
     enabled_tools: Vec<EnabledToolMetadata>,
     source: String,
-    stored_values: HashMap<String, JsonValue>,
+    stored_values: HashMap<String, Arc<JsonValue>>,
 }
 
 pub(super) struct RuntimeState {
     event_tx: mpsc::UnboundedSender<RuntimeEvent>,
     pending_tool_calls: HashMap<String, v8::Global<v8::PromiseResolver>>,
     pending_timeouts: HashMap<u64, timers::ScheduledTimeout>,
-    stored_values: HashMap<String, JsonValue>,
-    stored_value_writes: HashMap<String, JsonValue>,
+    stored_values: HashMap<String, Arc<JsonValue>>,
+    stored_value_writes: HashMap<String, Arc<JsonValue>>,
     enabled_tools: Vec<EnabledToolMetadata>,
     next_tool_call_id: u64,
     next_timeout_id: u64,
@@ -163,7 +164,7 @@ pub(super) struct RuntimeState {
 pub(super) enum CompletionState {
     Pending,
     Completed {
-        stored_value_writes: HashMap<String, JsonValue>,
+        stored_value_writes: HashMap<String, Arc<JsonValue>>,
         error_text: Option<String>,
     },
 }
@@ -321,7 +322,7 @@ fn capture_scope_send_error(
 
 fn send_result(
     event_tx: &mpsc::UnboundedSender<RuntimeEvent>,
-    stored_value_writes: HashMap<String, JsonValue>,
+    stored_value_writes: HashMap<String, Arc<JsonValue>>,
     error_text: Option<String>,
 ) {
     let _ = event_tx.send(RuntimeEvent::Result {

codex-rs/code-mode-runtime/src/session_runtime/mod.rs

@@ -43,7 +43,9 @@ pub(crate) struct SessionRuntime {
 }
 
 struct Inner {
-    stored_values: Mutex<HashMap<String, JsonValue>>,
+    // Cells snapshot the keys while sharing immutable payloads, so later commits cannot
+    // change an existing cell's view and unused values do not need to be copied.
+    stored_values: Mutex<HashMap<String, Arc<JsonValue>>>,
     cells: Mutex<HashMap<CellId, CellHandle>>,
     cell_tasks: TaskTracker,
     shutdown_token: CancellationToken,
@@ -277,7 +279,7 @@ impl<D: SessionRuntimeDelegate> CellHost for RuntimeCellHost<D> {
 
     async fn commit_completion(
         &self,
-        stored_value_writes: HashMap<String, JsonValue>,
+        stored_value_writes: HashMap<String, Arc<JsonValue>>,
         event: CellEvent,
         pending_initial_yield_items: Option<Vec<OutputItem>>,
         cell_state: Arc<CellState>,

codex-rs/code-mode-runtime/src/session_runtime/tests.rs

@@ -114,7 +114,7 @@ async fn termination_rejects_a_waiting_store_commit_before_the_next_cell_can_loa
     let commit = host.commit_completion(
         HashMap::from([(
             "candidate".to_string(),
-            JsonValue::String("lost".to_string()),
+            Arc::new(JsonValue::String("lost".to_string())),
         )]),
         completion.clone(),
         /*pending_initial_yield_items*/ None,