From dd6eff3954949d074ca45ee80da89499b955bf8e Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Mon, 21 Sep 2026 13:46:11 +0000 Subject: [PATCH] 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` 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 --- .../src/cell_actor/callbacks_tests.rs | 2 +- codex-rs/code-mode-runtime/src/cell_actor/mod.rs | 2 +- .../code-mode-runtime/src/cell_actor/tests.rs | 4 ++-- .../code-mode-runtime/src/cell_actor/types.rs | 2 +- .../code-mode-runtime/src/runtime/callbacks.rs | 6 +++++- codex-rs/code-mode-runtime/src/runtime/mod.rs | 15 ++++++++------- .../code-mode-runtime/src/session_runtime/mod.rs | 6 ++++-- .../src/session_runtime/tests.rs | 2 +- 8 files changed, 23 insertions(+), 16 deletions(-) diff --git a/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs index 585f7808a8..23c4bc9634 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/callbacks_tests.rs +++ b/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, + _stored_value_writes: HashMap>, _event: CellEvent, _pending_initial_yield_items: Option>, _cell_state: Arc, diff --git a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs index 7533f4b0cc..d403464746 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/mod.rs +++ b/codex-rs/code-mode-runtime/src/cell_actor/mod.rs @@ -47,7 +47,7 @@ pub(crate) struct CellActor; impl CellActor { pub(crate) fn prepare( request: CellRequest, - stored_values: HashMap, + stored_values: HashMap>, host: Arc, initial_observe_mode: ObserveMode, cell_state: Arc, diff --git a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs b/codex-rs/code-mode-runtime/src/cell_actor/tests.rs index 3612daff91..ca472cb3bc 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/tests.rs +++ b/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, + _stored_value_writes: HashMap>, event: CellEvent, pending_initial_yield_items: Option>, cell_state: Arc, @@ -75,7 +75,7 @@ impl CellHost for RecordingHost { async fn commit_completion( &self, - _stored_value_writes: HashMap, + _stored_value_writes: HashMap>, event: CellEvent, pending_initial_yield_items: Option>, cell_state: Arc, diff --git a/codex-rs/code-mode-runtime/src/cell_actor/types.rs b/codex-rs/code-mode-runtime/src/cell_actor/types.rs index 672820b93e..56762f96c9 100644 --- a/codex-rs/code-mode-runtime/src/cell_actor/types.rs +++ b/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, + stored_value_writes: HashMap>, event: CellEvent, pending_initial_yield_items: Option>, cell_state: Arc, diff --git a/codex-rs/code-mode-runtime/src/runtime/callbacks.rs b/codex-rs/code-mode-runtime/src/runtime/callbacks.rs index fcd61fd63c..7d18ece900 100644 --- a/codex-rs/code-mode-runtime/src/runtime/callbacks.rs +++ b/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::() { - 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); } } diff --git a/codex-rs/code-mode-runtime/src/runtime/mod.rs b/codex-rs/code-mode-runtime/src/runtime/mod.rs index f2e75c6f50..bdb667d943 100644 --- a/codex-rs/code-mode-runtime/src/runtime/mod.rs +++ b/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, + stored_value_writes: HashMap>, error_text: Option, }, ThreadPanicked, } pub(crate) fn spawn_runtime( - stored_values: HashMap, + stored_values: HashMap>, request: ExecuteRequest, event_tx: mpsc::UnboundedSender, pending_mode: PendingRuntimeMode, @@ -143,15 +144,15 @@ struct RuntimeConfig { tool_call_id: String, enabled_tools: Vec, source: String, - stored_values: HashMap, + stored_values: HashMap>, } pub(super) struct RuntimeState { event_tx: mpsc::UnboundedSender, pending_tool_calls: HashMap>, pending_timeouts: HashMap, - stored_values: HashMap, - stored_value_writes: HashMap, + stored_values: HashMap>, + stored_value_writes: HashMap>, enabled_tools: Vec, 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, + stored_value_writes: HashMap>, error_text: Option, }, } @@ -321,7 +322,7 @@ fn capture_scope_send_error( fn send_result( event_tx: &mpsc::UnboundedSender, - stored_value_writes: HashMap, + stored_value_writes: HashMap>, error_text: Option, ) { let _ = event_tx.send(RuntimeEvent::Result { diff --git a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs index f069b8af0f..0fa6abd58c 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/mod.rs +++ b/codex-rs/code-mode-runtime/src/session_runtime/mod.rs @@ -43,7 +43,9 @@ pub(crate) struct SessionRuntime { } struct Inner { - stored_values: Mutex>, + // 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>>, cells: Mutex>, cell_tasks: TaskTracker, shutdown_token: CancellationToken, @@ -277,7 +279,7 @@ impl CellHost for RuntimeCellHost { async fn commit_completion( &self, - stored_value_writes: HashMap, + stored_value_writes: HashMap>, event: CellEvent, pending_initial_yield_items: Option>, cell_state: Arc, diff --git a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs b/codex-rs/code-mode-runtime/src/session_runtime/tests.rs index 797cd00999..dc4395a975 100644 --- a/codex-rs/code-mode-runtime/src/session_runtime/tests.rs +++ b/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,