diff --git a/codex-rs/core/src/codex.rs b/codex-rs/core/src/codex.rs index 96efcbdfe..3ecb46963 100644 --- a/codex-rs/core/src/codex.rs +++ b/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(), diff --git a/codex-rs/core/src/codex_tests.rs b/codex-rs/core/src/codex_tests.rs index 7a17bdd98..b94f0d92a 100644 --- a/codex-rs/core/src/codex_tests.rs +++ b/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(), diff --git a/codex-rs/core/src/state/service.rs b/codex-rs/core/src/state/service.rs index 012e17bbf..5c0a741a1 100644 --- a/codex-rs/core/src/state/service.rs +++ b/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>, +} + +impl Default for CodeModeStoreService { + fn default() -> Self { + Self { + stored_values: Mutex::new(HashMap::new()), + } + } +} + +impl CodeModeStoreService { + pub(crate) async fn stored_values(&self) -> HashMap { + self.stored_values.lock().await.clone() + } + + pub(crate) async fn replace_stored_values(&self, values: HashMap) { + *self.stored_values.lock().await = values; + } +} + pub(crate) struct SessionServices { pub(crate) mcp_connection_manager: Arc>, pub(crate) mcp_startup_cancellation_token: Mutex, @@ -59,4 +82,5 @@ pub(crate) struct SessionServices { pub(crate) state_db: Option, /// Session-scoped model client shared across turns. pub(crate) model_client: ModelClient, + pub(crate) code_mode_store: CodeModeStoreService, } diff --git a/codex-rs/core/src/tools/code_mode.rs b/codex-rs/core/src/tools/code_mode.rs index abe11b248..cd75bc61a 100644 --- a/codex-rs/core/src/tools/code_mode.rs +++ b/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, + stored_values: HashMap, source: String, }, Response { @@ -76,6 +78,7 @@ enum NodeToHostMessage { }, Result { content_items: Vec, + stored_values: HashMap, #[serde(default)] max_output_tokens_per_exec_call: Option, }, @@ -94,7 +97,7 @@ pub(crate) fn instructions(config: &Config) -> Option { 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/.js`; MCP tools use `tools/mcp/.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, + stored_values: HashMap, ) -> Result, 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, diff --git a/codex-rs/core/src/tools/code_mode_runner.cjs b/codex-rs/core/src/tools/code_mode_runner.cjs index e66f9bdb7..7dfaf4480 100644 --- a/codex-rs/core/src/tools/code_mode_runner.cjs +++ b/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); diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index e303a22df..c61a1e46b 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/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/.js`; MCP tools use `tools/mcp/.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/.js`; MCP tools use `tools/mcp/.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 { diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 4aca988ed..5a60ed85f 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/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(()) +}