code-mode: merge stored values by key (#24159)

## Summary

Change code-mode stored value updates to merge writes by key instead of
replacing the session's complete stored-value map after each cell
completes.

Previously, each cell received a snapshot of stored values and returned
the complete resulting map. When multiple cells ran concurrently, a
later completion could overwrite values written by another cell because
it committed an older snapshot.

This change moves stored-value ownership into `CodeModeService`:

- Each runtime starts from the service's current stored values.
- Runtime completion reports only keys written by that cell.
- The service merges those writes into the current stored-value map on
successful completion.
- Core no longer replaces its stored-value state from a cell result.

As a result, concurrently executing cells can update different stored
keys without clobbering one another.

The move into CodeModeService is motivated by a desire to have this
lifetime tied to a new lifetime object on that side in a subsequent PR.
This commit is contained in:
Channing Conger
2026-05-22 19:09:02 -07:00
committed by GitHub
Unverified
parent 0febb1100f
commit f94157a4b2
7 changed files with 188 additions and 77 deletions
+2 -1
View File
@@ -157,7 +157,8 @@ pub(super) fn store_callback(
}
};
if let Some(state) = scope.get_slot_mut::<RuntimeState>() {
state.stored_values.insert(key, serialized);
state.stored_values.insert(key.clone(), serialized.clone());
state.stored_value_writes.insert(key, serialized);
}
}
+18 -21
View File
@@ -35,7 +35,6 @@ pub struct ExecuteRequest {
pub tool_call_id: String,
pub enabled_tools: Vec<ToolDefinition>,
pub source: String,
pub stored_values: HashMap<String, JsonValue>,
pub yield_time_ms: Option<u64>,
pub max_output_tokens: Option<usize>,
}
@@ -116,7 +115,6 @@ pub enum RuntimeResponse {
Result {
cell_id: String,
content_items: Vec<FunctionCallOutputContentItem>,
stored_values: HashMap<String, JsonValue>,
error_text: Option<String>,
},
}
@@ -182,12 +180,13 @@ pub(crate) enum RuntimeEvent {
text: String,
},
Result {
stored_values: HashMap<String, JsonValue>,
stored_value_writes: HashMap<String, JsonValue>,
error_text: Option<String>,
},
}
pub(crate) fn spawn_runtime(
stored_values: HashMap<String, JsonValue>,
request: ExecuteRequest,
event_tx: mpsc::UnboundedSender<RuntimeEvent>,
pending_mode: PendingRuntimeMode,
@@ -214,7 +213,7 @@ pub(crate) fn spawn_runtime(
tool_call_id: request.tool_call_id,
enabled_tools,
source: request.source,
stored_values: request.stored_values,
stored_values,
};
thread::spawn(move || {
@@ -248,6 +247,7 @@ pub(super) struct RuntimeState {
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>,
enabled_tools: Vec<EnabledToolMetadata>,
next_tool_call_id: u64,
next_timeout_id: u64,
@@ -259,7 +259,7 @@ pub(super) struct RuntimeState {
pub(super) enum CompletionState {
Pending,
Completed {
stored_values: HashMap<String, JsonValue>,
stored_value_writes: HashMap<String, JsonValue>,
error_text: Option<String>,
},
}
@@ -305,6 +305,7 @@ fn run_runtime(
pending_tool_calls: HashMap::new(),
pending_timeouts: HashMap::new(),
stored_values: config.stored_values,
stored_value_writes: HashMap::new(),
enabled_tools: config.enabled_tools,
next_tool_call_id: 1,
next_timeout_id: 1,
@@ -330,10 +331,10 @@ fn run_runtime(
match module_loader::completion_state(scope, pending_promise.as_ref()) {
CompletionState::Completed {
stored_values,
stored_value_writes,
error_text,
} => {
send_result(&event_tx, stored_values, error_text);
send_result(&event_tx, stored_value_writes, error_text);
return;
}
CompletionState::Pending => {}
@@ -375,10 +376,10 @@ fn run_runtime(
scope.perform_microtask_checkpoint();
match module_loader::completion_state(scope, pending_promise.as_ref()) {
CompletionState::Completed {
stored_values,
stored_value_writes,
error_text,
} => {
send_result(&event_tx, stored_values, error_text);
send_result(&event_tx, stored_value_writes, error_text);
return;
}
CompletionState::Pending => {}
@@ -422,21 +423,21 @@ fn capture_scope_send_error(
event_tx: &mpsc::UnboundedSender<RuntimeEvent>,
error_text: Option<String>,
) {
let stored_values = scope
let stored_value_writes = scope
.get_slot::<RuntimeState>()
.map(|state| state.stored_values.clone())
.map(|state| state.stored_value_writes.clone())
.unwrap_or_default();
send_result(event_tx, stored_values, error_text);
send_result(event_tx, stored_value_writes, error_text);
}
fn send_result(
event_tx: &mpsc::UnboundedSender<RuntimeEvent>,
stored_values: HashMap<String, JsonValue>,
stored_value_writes: HashMap<String, JsonValue>,
error_text: Option<String>,
) {
let _ = event_tx.send(RuntimeEvent::Result {
stored_values,
stored_value_writes,
error_text,
});
}
@@ -463,7 +464,6 @@ mod tests {
tool_call_id: "call_1".to_string(),
enabled_tools: Vec::new(),
source: source.to_string(),
stored_values: HashMap::new(),
yield_time_ms: Some(1),
max_output_tokens: None,
}
@@ -473,6 +473,7 @@ mod tests {
async fn terminate_execution_stops_cpu_bound_module() {
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
let (_runtime_tx, _runtime_control_tx, runtime_terminate_handle) = spawn_runtime(
HashMap::new(),
execute_request("while (true) {}"),
event_tx,
PendingRuntimeMode::Continue,
@@ -491,14 +492,9 @@ mod tests {
.await
.unwrap()
.unwrap();
let RuntimeEvent::Result {
stored_values,
error_text,
} = result_event
else {
let RuntimeEvent::Result { error_text, .. } = result_event else {
panic!("expected runtime result after termination");
};
assert_eq!(stored_values, HashMap::new());
assert!(error_text.is_some());
assert!(
@@ -513,6 +509,7 @@ mod tests {
async fn pending_mode_freezes_runtime_commands_until_resume() {
let (event_tx, mut event_rx) = mpsc::unbounded_channel();
let (runtime_tx, runtime_control_tx, _runtime_terminate_handle) = spawn_runtime(
HashMap::new(),
execute_request(
r#"
await new Promise((resolve) => setTimeout(resolve, 60_000));
@@ -104,14 +104,14 @@ pub(super) fn completion_state(
scope: &mut v8::PinScope<'_, '_>,
pending_promise: Option<&v8::Global<v8::Promise>>,
) -> CompletionState {
let stored_values = scope
let stored_value_writes = scope
.get_slot::<RuntimeState>()
.map(|state| state.stored_values.clone())
.map(|state| state.stored_value_writes.clone())
.unwrap_or_default();
let Some(pending_promise) = pending_promise else {
return CompletionState::Completed {
stored_values,
stored_value_writes,
error_text: None,
};
};
@@ -120,7 +120,7 @@ pub(super) fn completion_state(
match promise.state() {
v8::PromiseState::Pending => CompletionState::Pending,
v8::PromiseState::Fulfilled => CompletionState::Completed {
stored_values,
stored_value_writes,
error_text: None,
},
v8::PromiseState::Rejected => {
@@ -131,7 +131,7 @@ pub(super) fn completion_state(
Some(value_to_error_text(scope, result))
};
CompletionState::Completed {
stored_values,
stored_value_writes,
error_text,
}
}