From 710682598d20c8a51d41b99a4d709b3a7b827115 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 9 Mar 2026 22:59:09 -0600 Subject: [PATCH] Export tools module into code mode runner (#14167) **Summary** - allow `code_mode` to pass enabled tools metadata to the runner and expose them via `tools.js` - import tools inside JavaScript rather than relying only on globals or proxies for nested tool calls - update specs, docs, and tests to exercise the new bridge and explain the tooling changes **Testing** - Not run (not requested) --- codex-rs/core/src/tools/code_mode.rs | 67 +++++++++++----- codex-rs/core/src/tools/code_mode_bridge.js | 27 +------ codex-rs/core/src/tools/code_mode_runner.cjs | 84 ++++++++++++++++++-- codex-rs/core/src/tools/spec.rs | 2 +- codex-rs/core/tests/suite/code_mode.rs | 6 +- 5 files changed, 131 insertions(+), 55 deletions(-) diff --git a/codex-rs/core/src/tools/code_mode.rs b/codex-rs/core/src/tools/code_mode.rs index 6203763e9..9c1df684c 100644 --- a/codex-rs/core/src/tools/code_mode.rs +++ b/codex-rs/core/src/tools/code_mode.rs @@ -55,6 +55,7 @@ struct EnabledTool { #[serde(tag = "type", rename_all = "snake_case")] enum HostToNodeMessage { Init { + enabled_tools: Vec, source: String, }, Response { @@ -69,7 +70,8 @@ enum NodeToHostMessage { ToolCall { id: String, name: String, - input: String, + #[serde(default)] + input: Option, }, Result { content_items: Vec, @@ -88,7 +90,7 @@ pub(crate) fn instructions(config: &Config) -> Option { section.push_str("- `code_mode` is a freeform/custom tool. Direct `code_mode` calls must send raw JavaScript tool input. Do not wrap code in JSON, quotes, or markdown code fences.\n"); 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("- Call nested tools with `await tools[name](args)` or identifier wrappers like `await exec_command(args)` when the tool name is a valid JavaScript identifier. Nested tool calls resolve to arrays of content items.\n"); + section.push_str("- Import nested tools from `tools.js`, for example `import { exec_command } from \"tools.js\"` or `import { tools } from \"tools.js\"`. `tools[name]` and identifier wrappers like `await exec_command(args)` remain available for compatibility. Nested tool calls resolve to arrays of content items.\n"); section.push_str( "- Function tools require JSON object arguments. Freeform tools require raw strings.\n", ); @@ -111,7 +113,7 @@ pub(crate) async fn execute( }; let enabled_tools = build_enabled_tools(&exec); let source = build_source(&code, &enabled_tools).map_err(FunctionCallError::RespondToModel)?; - execute_node(exec, source) + execute_node(exec, source, enabled_tools) .await .map_err(FunctionCallError::RespondToModel) } @@ -119,11 +121,13 @@ pub(crate) async fn execute( async fn execute_node( exec: ExecContext, source: String, + enabled_tools: Vec, ) -> Result, String> { let node_path = resolve_compatible_node(exec.turn.config.js_repl_node_path.as_deref()).await?; let env = create_env(&exec.turn.shell_environment_policy, None); let mut cmd = tokio::process::Command::new(&node_path); + cmd.arg("--experimental-vm-modules"); cmd.arg("--eval"); cmd.arg(CODE_MODE_RUNNER_SOURCE); cmd.current_dir(&exec.turn.cwd); @@ -157,7 +161,14 @@ async fn execute_node( String::from_utf8_lossy(&buf).trim().to_string() }); - write_message(&mut stdin, &HostToNodeMessage::Init { source }).await?; + write_message( + &mut stdin, + &HostToNodeMessage::Init { + enabled_tools: enabled_tools.clone(), + source, + }, + ) + .await?; let mut stdout_lines = BufReader::new(stdout).lines(); let mut final_content_items = None; @@ -275,7 +286,11 @@ fn build_enabled_tools(exec: &ExecContext) -> Vec { out } -async fn call_nested_tool(exec: ExecContext, tool_name: String, input: String) -> Vec { +async fn call_nested_tool( + exec: ExecContext, + tool_name: String, + input: Option, +) -> Vec { if tool_name == "code_mode" { return error_content_items_json("code_mode cannot invoke itself".to_string()); } @@ -336,27 +351,39 @@ fn tool_kind_for_name(specs: &[ToolSpec], tool_name: &str) -> Result, ) -> Result { let actual_kind = tool_kind_for_name(specs, tool_name)?; match actual_kind { - CodeModeToolKind::Function => { - validate_function_arguments(tool_name, &input)?; - Ok(ToolPayload::Function { arguments: input }) - } - CodeModeToolKind::Freeform => Ok(ToolPayload::Custom { input }), + CodeModeToolKind::Function => build_function_tool_payload(tool_name, input), + CodeModeToolKind::Freeform => build_freeform_tool_payload(tool_name, input), } } -fn validate_function_arguments(tool_name: &str, input: &str) -> Result<(), String> { - let value: JsonValue = serde_json::from_str(input) - .map_err(|err| format!("tool `{tool_name}` expects a JSON object for arguments: {err}"))?; - if value.is_object() { - Ok(()) - } else { - Err(format!( - "tool `{tool_name}` expects a JSON object for arguments" - )) +fn build_function_tool_payload( + tool_name: &str, + input: Option, +) -> Result { + let arguments = match input { + None => "{}".to_string(), + Some(JsonValue::Object(map)) => serde_json::to_string(&JsonValue::Object(map)) + .map_err(|err| format!("failed to serialize tool `{tool_name}` arguments: {err}"))?, + Some(_) => { + return Err(format!( + "tool `{tool_name}` expects a JSON object for arguments" + )); + } + }; + Ok(ToolPayload::Function { arguments }) +} + +fn build_freeform_tool_payload( + tool_name: &str, + input: Option, +) -> Result { + match input { + Some(JsonValue::String(input)) => Ok(ToolPayload::Custom { input }), + _ => Err(format!("tool `{tool_name}` expects a string input")), } } diff --git a/codex-rs/core/src/tools/code_mode_bridge.js b/codex-rs/core/src/tools/code_mode_bridge.js index 406d2f35b..eba69c9f3 100644 --- a/codex-rs/core/src/tools/code_mode_bridge.js +++ b/codex-rs/core/src/tools/code_mode_bridge.js @@ -1,7 +1,5 @@ -(async () => { const __codexEnabledTools = __CODE_MODE_ENABLED_TOOLS_PLACEHOLDER__; const __codexEnabledToolNames = __codexEnabledTools.map((tool) => tool.name); -const __codexToolKinds = new Map(__codexEnabledTools.map((tool) => [tool.name, tool.kind])); const __codexContentItems = []; function __codexCloneContentItem(item) { @@ -31,26 +29,6 @@ function __codexNormalizeContentItems(value) { return [__codexCloneContentItem(value)]; } -async function __codexCallTool(name, args) { - const toolKind = __codexToolKinds.get(name); - if (toolKind === undefined) { - throw new Error(`Tool "${name}" is not enabled in code_mode`); - } - if (toolKind === 'freeform') { - if (typeof args !== 'string') { - throw new TypeError(`Tool "${name}" expects a string input`); - } - return await __codex_tool_call(name, args); - } - if (args === undefined) { - return await __codex_tool_call(name, '{}'); - } - if (!args || typeof args !== 'object' || Array.isArray(args)) { - throw new TypeError(`Tool "${name}" expects a JSON object for arguments`); - } - return await __codex_tool_call(name, JSON.stringify(args)); -} - Object.defineProperty(globalThis, '__codexContentItems', { value: __codexContentItems, configurable: true, @@ -71,7 +49,7 @@ globalThis.add_content = (value) => { globalThis.tools = new Proxy(Object.create(null), { get(_target, prop) { const name = String(prop); - return async (args) => __codexCallTool(name, args); + return async (args) => __codex_tool_call(name, args); }, }); @@ -86,7 +64,7 @@ globalThis.console = Object.freeze({ for (const name of __codexEnabledToolNames) { if (/^[A-Za-z_$][0-9A-Za-z_$]*$/.test(name) && !(name in globalThis)) { Object.defineProperty(globalThis, name, { - value: async (args) => __codexCallTool(name, args), + value: async (args) => __codex_tool_call(name, args), configurable: true, enumerable: false, writable: false, @@ -95,4 +73,3 @@ for (const name of __codexEnabledToolNames) { } __CODE_MODE_USER_CODE_PLACEHOLDER__ -})(); diff --git a/codex-rs/core/src/tools/code_mode_runner.cjs b/codex-rs/core/src/tools/code_mode_runner.cjs index ca9fcac67..09fe9e8af 100644 --- a/codex-rs/core/src/tools/code_mode_runner.cjs +++ b/codex-rs/core/src/tools/code_mode_runner.cjs @@ -3,6 +3,8 @@ const readline = require('node:readline'); const vm = require('node:vm'); +const { SourceTextModule, SyntheticModule } = vm; + function createProtocol() { const rl = readline.createInterface({ input: process.stdin, @@ -86,10 +88,7 @@ function createProtocol() { function readContentItems(context) { try { - const serialized = vm.runInContext( - 'JSON.stringify(globalThis.__codexContentItems ?? [])', - context - ); + const serialized = vm.runInContext('JSON.stringify(globalThis.__codexContentItems ?? [])', context); const contentItems = JSON.parse(serialized); return Array.isArray(contentItems) ? contentItems : []; } catch { @@ -97,6 +96,78 @@ function readContentItems(context) { } } +function isValidIdentifier(name) { + return /^[A-Za-z_$][0-9A-Za-z_$]*$/.test(name); +} + +function createToolsNamespace(protocol, enabledTools) { + const tools = Object.create(null); + + for (const { name } of enabledTools) { + const callTool = async (args) => + protocol.request('tool_call', { + name: String(name), + input: args, + }); + Object.defineProperty(tools, name, { + value: callTool, + configurable: false, + enumerable: true, + writable: false, + }); + } + + return Object.freeze(tools); +} + +function createToolsModule(context, protocol, enabledTools) { + const tools = createToolsNamespace(protocol, enabledTools); + const exportNames = ['tools']; + + for (const { name } of enabledTools) { + if (name !== 'tools' && isValidIdentifier(name)) { + exportNames.push(name); + } + } + + const uniqueExportNames = [...new Set(exportNames)]; + + return new SyntheticModule( + uniqueExportNames, + function initToolsModule() { + this.setExport('tools', tools); + for (const exportName of uniqueExportNames) { + if (exportName !== 'tools') { + this.setExport(exportName, tools[exportName]); + } + } + }, + { context } + ); +} + +async function runModule(context, protocol, request) { + const toolsModule = createToolsModule(context, protocol, request.enabled_tools ?? []); + const mainModule = new SourceTextModule(request.source, { + context, + identifier: 'code_mode_main.mjs', + importModuleDynamically(specifier) { + if (specifier === 'tools.js') { + return toolsModule; + } + throw new Error(`Unsupported import in code_mode: ${specifier}`); + }, + }); + + await mainModule.link(async (specifier) => { + if (specifier === 'tools.js') { + return toolsModule; + } + throw new Error(`Unsupported import in code_mode: ${specifier}`); + }); + await mainModule.evaluate(); +} + async function main() { const protocol = createProtocol(); const request = await protocol.init; @@ -109,10 +180,7 @@ async function main() { }); try { - await vm.runInContext(request.source, context, { - displayErrors: true, - microtaskMode: 'afterEvaluate', - }); + await runModule(context, protocol, request); await protocol.send({ type: 'result', content_items: readContentItems(context), diff --git a/codex-rs/core/src/tools/spec.rs b/codex-rs/core/src/tools/spec.rs index 42a7558c6..1553631d1 100644 --- a/codex-rs/core/src/tools/spec.rs +++ b/codex-rs/core/src/tools/spec.rs @@ -1544,7 +1544,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, call nested tools with `await tools[name](args)` or identifier wrappers like `await shell(args)` when the tool name is a valid JS identifier. Nested tool calls resolve to arrays of content items. Function tools require JSON object arguments. Freeform tools require raw strings. Use synchronous `add_content(value)` with a content item or content-item array, including `add_content(await exec_command(...))`, to return the same content items a direct tool call would expose to the model. Only content passed to `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\"`. `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 arrays of content items. Function tools require JSON object arguments. Freeform tools require raw strings. Use synchronous `add_content(value)` with a content item or content-item array, including `add_content(await exec_command(...))`, to return the same content items a direct tool call would expose to the model. Only content passed to `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 778ba137d..55e23ce1c 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -73,6 +73,8 @@ async fn code_mode_can_return_exec_command_output() -> Result<()> { &server, "use code_mode to run exec_command", r#" +import { exec_command } from "tools.js"; + add_content(await exec_command({ cmd: "printf code_mode_exec_marker" })); "#, false, @@ -112,7 +114,9 @@ async fn code_mode_can_apply_patch_via_nested_tool() -> Result<()> { let patch = format!( "*** Begin Patch\n*** Add File: {file_name}\n+hello from code_mode\n*** End Patch\n" ); - let code = format!("const items = await apply_patch({patch:?});\nadd_content(items);\n"); + let code = format!( + "import {{ apply_patch }} from \"tools.js\";\nconst items = await apply_patch({patch:?});\nadd_content(items);\n" + ); let (test, second_mock) = run_code_mode_turn(&server, "use code_mode to run apply_patch", &code, true).await?;