Add store/load support for code mode (#14259)

adds support for transferring state across code mode invocations.

codex-rs/core/src/codex.rs

@@ -1644,6 +1644,7 @@ impl Session {
                 config.features.enabled(Feature::RuntimeMetrics),
                 Self::build_model_client_beta_features_header(config.as_ref()),
             ),
+            code_mode_store: Default::default(),
         };
         let js_repl = Arc::new(JsReplHandle::with_node_path(
             config.js_repl_node_path.clone(),

codex-rs/core/src/codex_tests.rs

@@ -2235,6 +2235,7 @@ pub(crate) async fn make_session_and_context() -> (Session, TurnContext) {
             config.features.enabled(Feature::RuntimeMetrics),
             Session::build_model_client_beta_features_header(config.as_ref()),
         ),
+        code_mode_store: Default::default(),
     };
     let js_repl = Arc::new(JsReplHandle::with_node_path(
         config.js_repl_node_path.clone(),
@@ -2792,6 +2793,7 @@ pub(crate) async fn make_session_and_context_with_dynamic_tools_and_rx(
             config.features.enabled(Feature::RuntimeMetrics),
             Session::build_model_client_beta_features_header(config.as_ref()),
         ),
+        code_mode_store: Default::default(),
     };
     let js_repl = Arc::new(JsReplHandle::with_node_path(
         config.js_repl_node_path.clone(),

codex-rs/core/src/state/service.rs

@@ -22,12 +22,35 @@ use crate::unified_exec::UnifiedExecProcessManager;
 use codex_hooks::Hooks;
 use codex_otel::SessionTelemetry;
 use codex_utils_absolute_path::AbsolutePathBuf;
+use serde_json::Value as JsonValue;
 use std::path::PathBuf;
 use tokio::sync::Mutex;
 use tokio::sync::RwLock;
 use tokio::sync::watch;
 use tokio_util::sync::CancellationToken;
 
+pub(crate) struct CodeModeStoreService {
+    stored_values: Mutex<HashMap<String, JsonValue>>,
+}
+
+impl Default for CodeModeStoreService {
+    fn default() -> Self {
+        Self {
+            stored_values: Mutex::new(HashMap::new()),
+        }
+    }
+}
+
+impl CodeModeStoreService {
+    pub(crate) async fn stored_values(&self) -> HashMap<String, JsonValue> {
+        self.stored_values.lock().await.clone()
+    }
+
+    pub(crate) async fn replace_stored_values(&self, values: HashMap<String, JsonValue>) {
+        *self.stored_values.lock().await = values;
+    }
+}
+
 pub(crate) struct SessionServices {
     pub(crate) mcp_connection_manager: Arc<RwLock<McpConnectionManager>>,
     pub(crate) mcp_startup_cancellation_token: Mutex<CancellationToken>,
@@ -59,4 +82,5 @@ pub(crate) struct SessionServices {
     pub(crate) state_db: Option<StateDbHandle>,
     /// Session-scoped model client shared across turns.
     pub(crate) model_client: ModelClient,
+    pub(crate) code_mode_store: CodeModeStoreService,
 }

codex-rs/core/src/tools/code_mode.rs

@@ -1,3 +1,4 @@
+use std::collections::HashMap;
 use std::process::ExitStatus;
 use std::sync::Arc;
 
@@ -57,6 +58,7 @@ struct EnabledTool {
 enum HostToNodeMessage {
     Init {
         enabled_tools: Vec<EnabledTool>,
+        stored_values: HashMap<String, JsonValue>,
         source: String,
     },
     Response {
@@ -76,6 +78,7 @@ enum NodeToHostMessage {
     },
     Result {
         content_items: Vec<JsonValue>,
+        stored_values: HashMap<String, JsonValue>,
         #[serde(default)]
         max_output_tokens_per_exec_call: Option<usize>,
     },
@@ -94,7 +97,7 @@ pub(crate) fn instructions(config: &Config) -> Option<String> {
     section.push_str("- Direct tool calls remain available while `code_mode` is enabled.\n");
     section.push_str("- `code_mode` uses the same Node runtime resolution as `js_repl`. If needed, point `js_repl_node_path` at the Node binary you want Codex to use.\n");
     section.push_str("- Import nested tools from `tools.js`, for example `import { exec_command } from \"tools.js\"` or `import { tools } from \"tools.js\"`. Namespaced tools are also available from `tools/<namespace...>.js`; MCP tools use `tools/mcp/<server>.js`, for example `import { append_notebook_logs_chart } from \"tools/mcp/ologs.js\"`. `tools[name]` and identifier wrappers like `await exec_command(args)` remain available for compatibility. Nested tool calls resolve to their code-mode result values.\n");
-    section.push_str("- Import `{ output_text, output_image, set_max_output_tokens_per_exec_call }` from `@openai/code_mode`. `output_text(value)` surfaces text back to the model and stringifies non-string objects with `JSON.stringify(...)` when possible. `output_image(imageUrl)` appends an `input_image` content item for `http(s)` or `data:` URLs. `set_max_output_tokens_per_exec_call(value)` sets the token budget used to truncate the final Rust-side result of the current `code_mode` execution; the default is `10000`. This guards the overall `code_mode` output, not individual nested tool invocations. When truncation happens, the final text uses the unified-exec style `Original token count:` / `Output:` wrapper and the usual `…N tokens truncated…` marker.\n");
+    section.push_str("- Import `{ output_text, output_image, set_max_output_tokens_per_exec_call, store, load }` from `@openai/code_mode` (or `\"openai/code_mode\"`). `output_text(value)` surfaces text back to the model and stringifies non-string objects with `JSON.stringify(...)` when possible. `output_image(imageUrl)` appends an `input_image` content item for `http(s)` or `data:` URLs. `store(key, value)` persists JSON-serializable values across `code_mode` calls in the current session, and `load(key)` returns a cloned stored value or `undefined`. `set_max_output_tokens_per_exec_call(value)` sets the token budget used to truncate the final Rust-side result of the current `code_mode` execution; the default is `10000`. This guards the overall `code_mode` output, not individual nested tool invocations. When truncation happens, the final text uses the unified-exec style `Original token count:` / `Output:` wrapper and the usual `…N tokens truncated…` marker.\n");
     section.push_str(
         "- Function tools require JSON object arguments. Freeform tools require raw strings.\n",
     );
@@ -116,8 +119,9 @@ pub(crate) async fn execute(
         tracker,
     };
     let enabled_tools = build_enabled_tools(&exec).await;
+    let stored_values = exec.session.services.code_mode_store.stored_values().await;
     let source = build_source(&code, &enabled_tools).map_err(FunctionCallError::RespondToModel)?;
-    execute_node(exec, source, enabled_tools)
+    execute_node(exec, source, enabled_tools, stored_values)
         .await
         .map_err(FunctionCallError::RespondToModel)
 }
@@ -126,6 +130,7 @@ async fn execute_node(
     exec: ExecContext,
     source: String,
     enabled_tools: Vec<EnabledTool>,
+    stored_values: HashMap<String, JsonValue>,
 ) -> Result<Vec<FunctionCallOutputContentItem>, String> {
     let node_path = resolve_compatible_node(exec.turn.config.js_repl_node_path.as_deref()).await?;
 
@@ -169,6 +174,7 @@ async fn execute_node(
         &mut stdin,
         &HostToNodeMessage::Init {
             enabled_tools: enabled_tools.clone(),
+            stored_values,
             source,
         },
     )
@@ -196,8 +202,14 @@ async fn execute_node(
             }
             NodeToHostMessage::Result {
                 content_items,
+                stored_values,
                 max_output_tokens_per_exec_call,
             } => {
+                exec.session
+                    .services
+                    .code_mode_store
+                    .replace_stored_values(stored_values)
+                    .await;
                 final_content_items = Some(truncate_code_mode_result(
                     output_content_items_from_json_values(content_items)?,
                     max_output_tokens_per_exec_call,

codex-rs/core/src/tools/code_mode_runner.cjs

@@ -108,6 +108,10 @@ function isValidIdentifier(name) {
   return /^[A-Za-z_$][0-9A-Za-z_$]*$/.test(name);
 }
 
+function cloneJsonValue(value) {
+  return JSON.parse(JSON.stringify(value));
+}
+
 function createToolCaller(protocol) {
   return (name, input) =>
     protocol.request('tool_call', {
@@ -197,6 +201,21 @@ function normalizeOutputImageUrl(value) {
 }
 
 function createCodeModeModule(context, state) {
+  const load = (key) => {
+    if (typeof key !== 'string') {
+      throw new TypeError('load key must be a string');
+    }
+    if (!Object.prototype.hasOwnProperty.call(state.storedValues, key)) {
+      return undefined;
+    }
+    return cloneJsonValue(state.storedValues[key]);
+  };
+  const store = (key, value) => {
+    if (typeof key !== 'string') {
+      throw new TypeError('store key must be a string');
+    }
+    state.storedValues[key] = cloneJsonValue(value);
+  };
   const outputText = (value) => {
     const item = {
       type: 'input_text',
@@ -215,8 +234,9 @@ function createCodeModeModule(context, state) {
   };
 
   return new SyntheticModule(
-    ['output_text', 'output_image', 'set_max_output_tokens_per_exec_call'],
+    ['load', 'output_text', 'output_image', 'set_max_output_tokens_per_exec_call', 'store'],
     function initCodeModeModule() {
+      this.setExport('load', load);
       this.setExport('output_text', outputText);
       this.setExport('output_image', outputImage);
       this.setExport('set_max_output_tokens_per_exec_call', (value) => {
@@ -224,6 +244,7 @@ function createCodeModeModule(context, state) {
         state.maxOutputTokensPerExecCall = normalized;
         return normalized;
       });
+      this.setExport('store', store);
     },
     { context }
   );
@@ -291,10 +312,9 @@ function createModuleResolver(context, callTool, enabledTools, state) {
     if (specifier === 'tools.js') {
       return toolsModule;
     }
-    if (specifier === '@openai/code_mode') {
+    if (specifier === '@openai/code_mode' || specifier === 'openai/code_mode') {
       return codeModeModule;
     }
-
     const namespacedMatch = /^tools\/(.+)\.js$/.exec(specifier);
     if (!namespacedMatch) {
       throw new Error(`Unsupported import in code_mode: ${specifier}`);
@@ -318,7 +338,7 @@ function createModuleResolver(context, callTool, enabledTools, state) {
   };
 }
 
-async function runModule(context, protocol, request, state, callTool) {
+async function runModule(context, request, state, callTool) {
   const resolveModule = createModuleResolver(
     context,
     callTool,
@@ -340,6 +360,7 @@ async function main() {
   const request = await protocol.init;
   const state = {
     maxOutputTokensPerExecCall: DEFAULT_MAX_OUTPUT_TOKENS_PER_EXEC_CALL,
+    storedValues: cloneJsonValue(request.stored_values ?? {}),
   };
   const callTool = createToolCaller(protocol);
   const context = vm.createContext({
@@ -348,10 +369,11 @@ async function main() {
   });
 
   try {
-    await runModule(context, protocol, request, state, callTool);
+    await runModule(context, request, state, callTool);
     await protocol.send({
       type: 'result',
       content_items: readContentItems(context),
+      stored_values: state.storedValues,
       max_output_tokens_per_exec_call: state.maxOutputTokensPerExecCall,
     });
     process.exit(0);
@@ -360,6 +382,7 @@ async function main() {
     await protocol.send({
       type: 'result',
       content_items: readContentItems(context),
+      stored_values: state.storedValues,
       max_output_tokens_per_exec_call: state.maxOutputTokensPerExecCall,
     });
     process.exit(1);

codex-rs/core/src/tools/spec.rs

@@ -1620,7 +1620,7 @@ source: /[\s\S]+/
         enabled_tool_names.join(", ")
     };
     let description = format!(
-        "Runs JavaScript in a Node-backed `node:vm` context. This is a freeform tool: send raw JavaScript source text (no JSON/quotes/markdown fences). Direct tool calls remain available while `code_mode` is enabled. Inside JavaScript, import nested tools from `tools.js`, for example `import {{ exec_command }} from \"tools.js\"` or `import {{ tools }} from \"tools.js\"`. Namespaced tools are also available from `tools/<namespace...>.js`; MCP tools use `tools/mcp/<server>.js`, for example `import {{ append_notebook_logs_chart }} from \"tools/mcp/ologs.js\"`. `tools[name]` and identifier wrappers like `await shell(args)` remain available for compatibility when the tool name is a valid JS identifier. Nested tool calls resolve to their code-mode result values. Import `{{ output_text, output_image, set_max_output_tokens_per_exec_call }}` from `\"@openai/code_mode\"`; `output_text(value)` surfaces text back to the model and stringifies non-string objects when possible, `output_image(imageUrl)` appends an `input_image` content item for `http(s)` or `data:` URLs, and `set_max_output_tokens_per_exec_call(value)` sets the token budget used to truncate the final Rust-side result of the current `code_mode` execution. The default is `10000`. This guards the overall `code_mode` output, not individual nested tool invocations. When truncation happens, the final text uses the unified-exec style `Original token count:` / `Output:` wrapper and the usual `…N tokens truncated…` marker. Function tools require JSON object arguments. Freeform tools require raw strings. `add_content(value)` remains available for compatibility with a content item, content-item array, or string. Structured nested-tool results should be converted to text first, for example with `JSON.stringify(...)`. Only content passed to `output_text(...)`, `output_image(...)`, or `add_content(value)` is surfaced back to the model. Enabled nested tools: {enabled_list}."
+        "Runs JavaScript in a Node-backed `node:vm` context. This is a freeform tool: send raw JavaScript source text (no JSON/quotes/markdown fences). Direct tool calls remain available while `code_mode` is enabled. Inside JavaScript, import nested tools from `tools.js`, for example `import {{ exec_command }} from \"tools.js\"` or `import {{ tools }} from \"tools.js\"`. Namespaced tools are also available from `tools/<namespace...>.js`; MCP tools use `tools/mcp/<server>.js`, for example `import {{ append_notebook_logs_chart }} from \"tools/mcp/ologs.js\"`. `tools[name]` and identifier wrappers like `await shell(args)` remain available for compatibility when the tool name is a valid JS identifier. Nested tool calls resolve to their code-mode result values. Import `{{ output_text, output_image, set_max_output_tokens_per_exec_call, store, load }}` from `\"@openai/code_mode\"` (or `\"openai/code_mode\"`); `output_text(value)` surfaces text back to the model and stringifies non-string objects when possible, `output_image(imageUrl)` appends an `input_image` content item for `http(s)` or `data:` URLs, `store(key, value)` persists JSON-serializable values across `code_mode` calls in the current session, `load(key)` returns a cloned stored value or `undefined`, and `set_max_output_tokens_per_exec_call(value)` sets the token budget used to truncate the final Rust-side result of the current `code_mode` execution. The default is `10000`. This guards the overall `code_mode` output, not individual nested tool invocations. When truncation happens, the final text uses the unified-exec style `Original token count:` / `Output:` wrapper and the usual `…N tokens truncated…` marker. Function tools require JSON object arguments. Freeform tools require raw strings. `add_content(value)` remains available for compatibility with a content item, content-item array, or string. Structured nested-tool results should be converted to text first, for example with `JSON.stringify(...)`. Only content passed to `output_text(...)`, `output_image(...)`, or `add_content(value)` is surfaced back to the model. Enabled nested tools: {enabled_list}."
     );
 
     ToolSpec::Freeform(FreeformTool {

codex-rs/core/tests/suite/code_mode.rs

@@ -524,3 +524,96 @@ structuredContent=null"
 
     Ok(())
 }
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn code_mode_can_store_and_load_values_across_turns() -> Result<()> {
+    skip_if_no_network!(Ok(()));
+
+    let server = responses::start_mock_server().await;
+    let mut builder = test_codex().with_config(move |config| {
+        let _ = config.features.enable(Feature::CodeMode);
+    });
+    let test = builder.build(&server).await?;
+
+    responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_response_created("resp-1"),
+            ev_custom_tool_call(
+                "call-1",
+                "code_mode",
+                r#"
+import { store } from "@openai/code_mode";
+
+store("nb", { title: "Notebook", items: [1, true, null] });
+add_content("stored");
+"#,
+            ),
+            ev_completed("resp-1"),
+        ]),
+    )
+    .await;
+    let first_follow_up = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_assistant_message("msg-1", "stored"),
+            ev_completed("resp-2"),
+        ]),
+    )
+    .await;
+
+    test.submit_turn("store value for later").await?;
+
+    let first_request = first_follow_up.single_request();
+    let (first_output, first_success) =
+        custom_tool_output_text_and_success(&first_request, "call-1");
+    assert_ne!(
+        first_success,
+        Some(false),
+        "code_mode store call failed unexpectedly: {first_output}"
+    );
+    assert_eq!(first_output, "stored");
+
+    responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_response_created("resp-3"),
+            ev_custom_tool_call(
+                "call-2",
+                "code_mode",
+                r#"
+import { load } from "openai/code_mode";
+
+add_content(JSON.stringify(load("nb")));
+"#,
+            ),
+            ev_completed("resp-3"),
+        ]),
+    )
+    .await;
+    let second_follow_up = responses::mount_sse_once(
+        &server,
+        sse(vec![
+            ev_assistant_message("msg-2", "loaded"),
+            ev_completed("resp-4"),
+        ]),
+    )
+    .await;
+
+    test.submit_turn("load the stored value").await?;
+
+    let second_request = second_follow_up.single_request();
+    let (second_output, second_success) =
+        custom_tool_output_text_and_success(&second_request, "call-2");
+    assert_ne!(
+        second_success,
+        Some(false),
+        "code_mode load call failed unexpectedly: {second_output}"
+    );
+    let loaded: Value = serde_json::from_str(&second_output)?;
+    assert_eq!(
+        loaded,
+        serde_json::json!({ "title": "Notebook", "items": [1, true, null] })
+    );
+
+    Ok(())
+}