From 531a79baf381eb2a4504576fb867260908ad8bb5 Mon Sep 17 00:00:00 2001 From: Channing Conger Date: Wed, 3 Jun 2026 23:50:00 +0000 Subject: [PATCH] feat(code-mode): allow disabling session store and load --- codex-rs/code-mode/src/config.rs | 16 ++++++ codex-rs/code-mode/src/description.rs | 53 ++++++++++++++++-- codex-rs/code-mode/src/lib.rs | 5 ++ codex-rs/code-mode/src/runtime/globals.rs | 13 +++-- codex-rs/code-mode/src/runtime/mod.rs | 26 +++++---- codex-rs/code-mode/src/service.rs | 68 +++++++++++++++++++++-- codex-rs/code-mode/src/stored_values.rs | 55 ++++++++++++++++++ 7 files changed, 210 insertions(+), 26 deletions(-) create mode 100644 codex-rs/code-mode/src/config.rs create mode 100644 codex-rs/code-mode/src/stored_values.rs diff --git a/codex-rs/code-mode/src/config.rs b/codex-rs/code-mode/src/config.rs new file mode 100644 index 0000000000..168192bd0c --- /dev/null +++ b/codex-rs/code-mode/src/config.rs @@ -0,0 +1,16 @@ +/// Configuration for one durable code-mode session. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct CodeModeSessionConfig { + /// Controls whether cells can persist JSON values for later cells. + pub store_load: StoreLoadMode, +} + +/// Controls whether the `store` and `load` globals are available in a session. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum StoreLoadMode { + /// Expose `store` and `load` and retain their JSON values between cells. + #[default] + Enabled, + /// Do not expose `store` or `load`, and do not retain JSON values. + Disabled, +} diff --git a/codex-rs/code-mode/src/description.rs b/codex-rs/code-mode/src/description.rs index 7e0801ed25..ccb6b95fd5 100644 --- a/codex-rs/code-mode/src/description.rs +++ b/codex-rs/code-mode/src/description.rs @@ -4,11 +4,17 @@ use serde::Serialize; use serde_json::Value as JsonValue; use std::collections::BTreeMap; +use crate::CodeModeSessionConfig; use crate::PUBLIC_TOOL_NAME; +use crate::StoreLoadMode; const MAX_JS_SAFE_INTEGER: u64 = (1_u64 << 53) - 1; const DEFERRED_NESTED_TOOLS_GUIDANCE: &str = r#"Some deferred nested tools may be omitted from this description. They are still available on the global `tools` object and listed in `ALL_TOOLS`. To find one, filter `ALL_TOOLS` by `name` and `description`."#; +const STORE_LOAD_HELPERS_PLACEHOLDER: &str = "{store_load_helpers}"; +const STORE_LOAD_HELPERS_DESCRIPTION: &str = r#"- `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session. +- `load(key: string)`: returns the stored value for a string key, or `undefined` if it is missing. +"#; const EXEC_DESCRIPTION_TEMPLATE: &str = r#"Run JavaScript code to orchestrate/compose tool calls - Evaluates the provided JavaScript code in a fresh V8 isolate as an async module. - All nested tools are available on the global `tools` object, for example `await tools.exec_command(...)`. Tool names are exposed as normalized JavaScript identifiers, for example `await tools.mcp__ologs__get_profile(...)`. @@ -26,9 +32,7 @@ const EXEC_DESCRIPTION_TEMPLATE: &str = r#"Run JavaScript code to orchestrate/co - `text(value: string | number | boolean | undefined | null)`: Appends a text item. Non-string values are stringified with `JSON.stringify(...)` when possible. - `image(imageUrlOrItem: string | { image_url: string; detail?: "auto" | "low" | "high" | "original" | null } | ImageContent, detail?: "auto" | "low" | "high" | "original" | null)`: Appends an image item. `image_url` can be an HTTPS URL or a base64-encoded `data:` URL. To forward an MCP tool image, pass an individual `ImageContent` block from `result.content`, for example `image(result.content[0])`. MCP image blocks may request detail with `_meta: { "codex/imageDetail": "original" }`. When provided, the second `detail` argument overrides any detail embedded in the first argument. - `generatedImage(result: { image_url: string; output_hint?: string })`: Appends an image-generation result and its optional output hint. -- `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session. -- `load(key: string)`: returns the stored value for a string key, or `undefined` if it is missing. -- `notify(value: string | number | boolean | undefined | null)`: immediately injects an extra `custom_tool_call_output` for the current `exec` call. Values are stringified like `text(...)`. +{store_load_helpers}- `notify(value: string | number | boolean | undefined | null)`: immediately injects an extra `custom_tool_call_output` for the current `exec` call. Values are stringified like `text(...)`. - `setTimeout(callback: () => void, delayMs?: number)`: schedules a callback to run later and returns a timeout id. Pending timeouts do not keep `exec` alive by themselves; await an explicit promise if you need to wait for one. - `clearTimeout(timeoutId?: number)`: cancels a timeout created by `setTimeout`. - `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries. @@ -253,9 +257,31 @@ pub fn build_exec_tool_description( namespace_descriptions: &BTreeMap, code_mode_only: bool, deferred_tools_available: bool, +) -> String { + build_exec_tool_description_with_config( + enabled_tools, + namespace_descriptions, + code_mode_only, + deferred_tools_available, + CodeModeSessionConfig::default(), + ) +} + +pub fn build_exec_tool_description_with_config( + enabled_tools: &[ToolDefinition], + namespace_descriptions: &BTreeMap, + code_mode_only: bool, + deferred_tools_available: bool, + session_config: CodeModeSessionConfig, ) -> String { let mut sections = Vec::new(); - sections.push(EXEC_DESCRIPTION_TEMPLATE.to_string()); + let store_load_helpers = match session_config.store_load { + StoreLoadMode::Enabled => STORE_LOAD_HELPERS_DESCRIPTION, + StoreLoadMode::Disabled => "", + }; + sections.push( + EXEC_DESCRIPTION_TEMPLATE.replace(STORE_LOAD_HELPERS_PLACEHOLDER, store_load_helpers), + ); if deferred_tools_available { sections.push(DEFERRED_NESTED_TOOLS_GUIDANCE.to_string()); } @@ -712,8 +738,11 @@ mod tests { use super::ToolNamespaceDescription; use super::augment_tool_definition; use super::build_exec_tool_description; + use super::build_exec_tool_description_with_config; use super::normalize_code_mode_identifier; use super::parse_exec_source; + use crate::CodeModeSessionConfig; + use crate::StoreLoadMode; use codex_protocol::ToolName; use pretty_assertions::assert_eq; use serde_json::Value as JsonValue; @@ -886,6 +915,22 @@ bar" assert!(description.contains("`clearTimeout(timeoutId?: number)`")); } + #[test] + fn exec_description_omits_disabled_store_load_helpers() { + let description = build_exec_tool_description_with_config( + &[], + &BTreeMap::new(), + /*code_mode_only*/ false, + /*deferred_tools_available*/ false, + CodeModeSessionConfig { + store_load: StoreLoadMode::Disabled, + }, + ); + + assert!(!description.contains("`store(key: string, value: any)`")); + assert!(!description.contains("`load(key: string)`")); + } + #[test] fn code_mode_only_description_groups_namespace_instructions_once() { let namespace_descriptions = BTreeMap::from([( diff --git a/codex-rs/code-mode/src/lib.rs b/codex-rs/code-mode/src/lib.rs index 37ba1d2a91..2267d69ab3 100644 --- a/codex-rs/code-mode/src/lib.rs +++ b/codex-rs/code-mode/src/lib.rs @@ -1,14 +1,19 @@ +mod config; mod description; mod response; mod runtime; mod service; +mod stored_values; +pub use config::CodeModeSessionConfig; +pub use config::StoreLoadMode; pub use description::CODE_MODE_PRAGMA_PREFIX; pub use description::CodeModeToolKind; pub use description::ToolDefinition; pub use description::ToolNamespaceDescription; pub use description::augment_tool_definition; pub use description::build_exec_tool_description; +pub use description::build_exec_tool_description_with_config; pub use description::build_wait_tool_description; pub use description::is_code_mode_nested_tool; pub use description::normalize_code_mode_identifier; diff --git a/codex-rs/code-mode/src/runtime/globals.rs b/codex-rs/code-mode/src/runtime/globals.rs index fe3df4c958..f5cd770d12 100644 --- a/codex-rs/code-mode/src/runtime/globals.rs +++ b/codex-rs/code-mode/src/runtime/globals.rs @@ -25,8 +25,6 @@ pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), St let text = helper_function(scope, "text", text_callback)?; let image = helper_function(scope, "image", image_callback)?; let generated_image = helper_function(scope, "generatedImage", generated_image_callback)?; - let store = helper_function(scope, "store", store_callback)?; - let load = helper_function(scope, "load", load_callback)?; let notify = helper_function(scope, "notify", notify_callback)?; let yield_control = helper_function(scope, "yield_control", yield_control_callback)?; let exit = helper_function(scope, "exit", exit_callback)?; @@ -38,8 +36,15 @@ pub(super) fn install_globals(scope: &mut v8::PinScope<'_, '_>) -> Result<(), St set_global(scope, global, "text", text.into())?; set_global(scope, global, "image", image.into())?; set_global(scope, global, "generatedImage", generated_image.into())?; - set_global(scope, global, "store", store.into())?; - set_global(scope, global, "load", load.into())?; + if scope + .get_slot::() + .is_some_and(|state| state.stored_values.is_enabled()) + { + let store = helper_function(scope, "store", store_callback)?; + let load = helper_function(scope, "load", load_callback)?; + set_global(scope, global, "store", store.into())?; + set_global(scope, global, "load", load.into())?; + } set_global(scope, global, "notify", notify.into())?; set_global(scope, global, "yield_control", yield_control.into())?; set_global(scope, global, "exit", exit.into())?; diff --git a/codex-rs/code-mode/src/runtime/mod.rs b/codex-rs/code-mode/src/runtime/mod.rs index 21757328ee..2da7adba9c 100644 --- a/codex-rs/code-mode/src/runtime/mod.rs +++ b/codex-rs/code-mode/src/runtime/mod.rs @@ -20,6 +20,7 @@ use crate::description::ToolDefinition; use crate::description::enabled_tool_metadata; use crate::response::FunctionCallOutputContentItem; use crate::service::CellId; +use crate::stored_values::StoredValues; pub const DEFAULT_EXEC_YIELD_TIME_MS: u64 = 10_000; pub const DEFAULT_WAIT_YIELD_TIME_MS: u64 = 10_000; @@ -165,13 +166,13 @@ pub(crate) enum RuntimeEvent { text: String, }, Result { - stored_value_writes: HashMap, + stored_value_writes: StoredValues, error_text: Option, }, } pub(crate) fn spawn_runtime( - stored_values: HashMap, + stored_values: StoredValues, request: ExecuteRequest, event_tx: mpsc::UnboundedSender, pending_mode: PendingRuntimeMode, @@ -224,15 +225,15 @@ struct RuntimeConfig { tool_call_id: String, enabled_tools: Vec, source: String, - stored_values: HashMap, + stored_values: StoredValues, } pub(super) struct RuntimeState { event_tx: mpsc::UnboundedSender, pending_tool_calls: HashMap>, pending_timeouts: HashMap, - stored_values: HashMap, - stored_value_writes: HashMap, + stored_values: StoredValues, + stored_value_writes: StoredValues, enabled_tools: Vec, next_tool_call_id: u64, next_timeout_id: u64, @@ -244,7 +245,7 @@ pub(super) struct RuntimeState { pub(super) enum CompletionState { Pending, Completed { - stored_value_writes: HashMap, + stored_value_writes: StoredValues, error_text: Option, }, } @@ -289,8 +290,8 @@ fn run_runtime( event_tx: event_tx.clone(), pending_tool_calls: HashMap::new(), pending_timeouts: HashMap::new(), + stored_value_writes: config.stored_values.empty_like(), stored_values: config.stored_values, - stored_value_writes: HashMap::new(), enabled_tools: config.enabled_tools, next_tool_call_id: 1, next_timeout_id: 1, @@ -300,7 +301,7 @@ fn run_runtime( }); if let Err(error_text) = globals::install_globals(scope) { - send_result(&event_tx, HashMap::new(), Some(error_text)); + capture_scope_send_error(scope, &event_tx, Some(error_text)); return; } @@ -418,7 +419,7 @@ fn capture_scope_send_error( fn send_result( event_tx: &mpsc::UnboundedSender, - stored_value_writes: HashMap, + stored_value_writes: StoredValues, error_text: Option, ) { let _ = event_tx.send(RuntimeEvent::Result { @@ -429,7 +430,6 @@ fn send_result( #[cfg(test)] mod tests { - use std::collections::HashMap; use std::time::Duration; use pretty_assertions::assert_eq; @@ -442,6 +442,8 @@ mod tests { use super::RuntimeEvent; use super::spawn_runtime; use crate::FunctionCallOutputContentItem; + use crate::StoreLoadMode; + use crate::stored_values::StoredValues; fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { @@ -457,7 +459,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(), + StoredValues::from_mode(StoreLoadMode::Enabled), execute_request("while (true) {}"), event_tx, PendingRuntimeMode::Continue, @@ -493,7 +495,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(), + StoredValues::from_mode(StoreLoadMode::Enabled), execute_request( r#" await new Promise((resolve) => setTimeout(resolve, 60_000)); diff --git a/codex-rs/code-mode/src/service.rs b/codex-rs/code-mode/src/service.rs index b348d1d35f..e0d93af05b 100644 --- a/codex-rs/code-mode/src/service.rs +++ b/codex-rs/code-mode/src/service.rs @@ -18,6 +18,7 @@ use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; use tracing::warn; +use crate::CodeModeSessionConfig; use crate::FunctionCallOutputContentItem; use crate::runtime::CodeModeNestedToolCall; use crate::runtime::DEFAULT_EXEC_YIELD_TIME_MS; @@ -33,6 +34,7 @@ use crate::runtime::WaitRequest; use crate::runtime::WaitToPendingOutcome; use crate::runtime::WaitToPendingRequest; use crate::runtime::spawn_runtime; +use crate::stored_values::StoredValues; pub type CodeModeSessionResultFuture<'a, T> = Pin> + Send + 'a>>; @@ -152,6 +154,7 @@ pub trait CodeModeSession: Send + Sync { pub trait CodeModeSessionProvider: Send + Sync { fn create_session<'a>( &'a self, + config: CodeModeSessionConfig, delegate: Arc, ) -> CodeModeSessionProviderFuture<'a>; } @@ -162,11 +165,12 @@ pub struct InProcessCodeModeSessionProvider; impl CodeModeSessionProvider for InProcessCodeModeSessionProvider { fn create_session<'a>( &'a self, + config: CodeModeSessionConfig, delegate: Arc, ) -> CodeModeSessionProviderFuture<'a> { Box::pin(async move { let session: Arc = - Arc::new(CodeModeService::with_delegate(delegate)); + Arc::new(CodeModeService::with_delegate_and_config(delegate, config)); Ok(session) }) } @@ -180,7 +184,7 @@ struct CellHandle { } struct Inner { - stored_values: Mutex>, + stored_values: Mutex, cells: Mutex>, delegate: Arc, shutting_down: AtomicBool, @@ -193,13 +197,24 @@ pub struct CodeModeService { impl CodeModeService { pub fn new() -> Self { - Self::with_delegate(Arc::new(NoopCodeModeSessionDelegate)) + Self::with_config(CodeModeSessionConfig::default()) + } + + pub fn with_config(config: CodeModeSessionConfig) -> Self { + Self::with_delegate_and_config(Arc::new(NoopCodeModeSessionDelegate), config) } pub fn with_delegate(delegate: Arc) -> Self { + Self::with_delegate_and_config(delegate, CodeModeSessionConfig::default()) + } + + pub fn with_delegate_and_config( + delegate: Arc, + config: CodeModeSessionConfig, + ) -> Self { Self { inner: Arc::new(Inner { - stored_values: Mutex::new(HashMap::new()), + stored_values: Mutex::new(StoredValues::from_mode(config.store_load)), cells: Mutex::new(HashMap::new()), delegate, shutting_down: AtomicBool::new(false), @@ -875,13 +890,16 @@ mod tests { use super::WaitToPendingOutcome; use super::WaitToPendingRequest; use super::run_cell_control; + use crate::CodeModeSessionConfig; use crate::CodeModeToolKind; use crate::FunctionCallOutputContentItem; + use crate::StoreLoadMode; use crate::ToolDefinition; use crate::runtime::ExecuteRequest; use crate::runtime::ExecuteToPendingOutcome; use crate::runtime::RuntimeEvent; use crate::runtime::spawn_runtime; + use crate::stored_values::StoredValues; fn execute_request(source: &str) -> ExecuteRequest { ExecuteRequest { @@ -909,7 +927,7 @@ mod tests { fn test_inner() -> Arc { Arc::new(Inner { - stored_values: Mutex::new(HashMap::new()), + stored_values: Mutex::new(StoredValues::from_mode(StoreLoadMode::Enabled)), cells: Mutex::new(HashMap::new()), delegate: Arc::new(NoopCodeModeSessionDelegate), shutting_down: std::sync::atomic::AtomicBool::new(false), @@ -1007,6 +1025,44 @@ mod tests { ); } + #[tokio::test] + async fn disabled_store_load_does_not_inject_globals_or_allocate_json_storage() { + let service = CodeModeService::with_config(CodeModeSessionConfig { + store_load: StoreLoadMode::Disabled, + }); + + let response = execute( + &service, + ExecuteRequest { + source: r#" +text(JSON.stringify([ + Object.hasOwn(globalThis, "store"), + Object.hasOwn(globalThis, "load"), +])); +"# + .to_string(), + yield_time_ms: None, + ..execute_request("") + }, + ) + .await; + + assert_eq!( + response, + RuntimeResponse::Result { + cell_id: cell_id("1"), + content_items: vec![FunctionCallOutputContentItem::InputText { + text: "[false,false]".to_string(), + }], + error_text: None, + } + ); + assert_eq!( + *service.inner.stored_values.lock().await, + StoredValues::Disabled + ); + } + #[tokio::test] async fn shutdown_interrupts_cpu_bound_cells() { let service = CodeModeService::new(); @@ -1809,7 +1865,7 @@ image({ let (initial_response_tx, initial_response_rx) = oneshot::channel(); let (runtime_event_tx, _runtime_event_rx) = mpsc::unbounded_channel(); let (runtime_tx, runtime_control_tx, runtime_terminate_handle) = spawn_runtime( - HashMap::new(), + StoredValues::from_mode(StoreLoadMode::Enabled), ExecuteRequest { source: "await new Promise(() => {})".to_string(), yield_time_ms: None, diff --git a/codex-rs/code-mode/src/stored_values.rs b/codex-rs/code-mode/src/stored_values.rs new file mode 100644 index 0000000000..9af04423c1 --- /dev/null +++ b/codex-rs/code-mode/src/stored_values.rs @@ -0,0 +1,55 @@ +use std::collections::HashMap; + +use serde_json::Value as JsonValue; + +use crate::StoreLoadMode; + +#[derive(Clone, Debug, Eq, PartialEq, Default)] +pub(crate) enum StoredValues { + Enabled(HashMap), + #[default] + Disabled, +} + +impl StoredValues { + pub(crate) fn from_mode(mode: StoreLoadMode) -> Self { + match mode { + StoreLoadMode::Enabled => Self::Enabled(HashMap::new()), + StoreLoadMode::Disabled => Self::Disabled, + } + } + + pub(crate) fn empty_like(&self) -> Self { + match self { + Self::Enabled(_) => Self::Enabled(HashMap::new()), + Self::Disabled => Self::Disabled, + } + } + + pub(crate) fn is_enabled(&self) -> bool { + matches!(self, Self::Enabled(_)) + } + + pub(crate) fn insert(&mut self, key: String, value: JsonValue) { + if let Self::Enabled(values) = self { + values.insert(key, value); + } + } + + pub(crate) fn get(&self, key: &str) -> Option<&JsonValue> { + match self { + Self::Enabled(values) => values.get(key), + Self::Disabled => None, + } + } + + pub(crate) fn extend(&mut self, writes: Self) { + match (self, writes) { + (Self::Enabled(values), Self::Enabled(writes)) => values.extend(writes), + (Self::Disabled, Self::Disabled) => {} + (Self::Enabled(_), Self::Disabled) | (Self::Disabled, Self::Enabled(_)) => { + debug_assert!(false, "stored value modes must match"); + } + } + } +}