From 37bcef99e35cf1c4fa3b36fb32213cb01d3ab3f8 Mon Sep 17 00:00:00 2001 From: pakrym-oai Date: Mon, 16 Mar 2026 14:25:06 -0700 Subject: [PATCH] Add exit helper for code mode --- codex-rs/core/src/tools/code_mode/bridge.js | 1 + .../core/src/tools/code_mode/description.md | 1 + codex-rs/core/src/tools/code_mode/runner.cjs | 35 +++++++++++++++- codex-rs/core/tests/suite/code_mode.rs | 42 +++++++++++++++++++ 4 files changed, 78 insertions(+), 1 deletion(-) diff --git a/codex-rs/core/src/tools/code_mode/bridge.js b/codex-rs/core/src/tools/code_mode/bridge.js index ab353feedeb..3bce192902d 100644 --- a/codex-rs/core/src/tools/code_mode/bridge.js +++ b/codex-rs/core/src/tools/code_mode/bridge.js @@ -27,6 +27,7 @@ Object.defineProperty(globalThis, '__codexContentItems', { } defineGlobal('ALL_TOOLS', __codexRuntime.ALL_TOOLS); + defineGlobal('exit', __codexRuntime.exit); defineGlobal('image', __codexRuntime.image); defineGlobal('load', __codexRuntime.load); defineGlobal('store', __codexRuntime.store); diff --git a/codex-rs/core/src/tools/code_mode/description.md b/codex-rs/core/src/tools/code_mode/description.md index 6aa21c5dc18..79e51ebf6ee 100644 --- a/codex-rs/core/src/tools/code_mode/description.md +++ b/codex-rs/core/src/tools/code_mode/description.md @@ -9,6 +9,7 @@ - They return either a structured value or a string based on the description above. - Global helpers: +- `exit()`: Immediately ends the current script successfully (like an early return from the top level). - `text(value: string | number | boolean | undefined | null)`: Appends a text item and returns it. Non-string values are stringified with `JSON.stringify(...)` when possible. - `image(imageUrl: string)`: Appends an image item and returns it. `image_url` can be an HTTPS URL or a base64-encoded `data:` URL. - `store(key: string, value: any)`: stores a serializable value under a string key for later `exec` calls in the same session. diff --git a/codex-rs/core/src/tools/code_mode/runner.cjs b/codex-rs/core/src/tools/code_mode/runner.cjs index a5d457e5bc4..9ebb9c98820 100644 --- a/codex-rs/core/src/tools/code_mode/runner.cjs +++ b/codex-rs/core/src/tools/code_mode/runner.cjs @@ -55,6 +55,17 @@ function codeModeWorkerMain() { return JSON.parse(JSON.stringify(value)); } + class CodeModeExitSignal extends Error { + constructor() { + super('code mode exit'); + this.name = 'CodeModeExitSignal'; + } + } + + function isCodeModeExitSignal(error) { + return error instanceof CodeModeExitSignal; + } + function createToolCaller() { let nextId = 0; const pending = new Map(); @@ -257,8 +268,12 @@ function codeModeWorkerMain() { const yieldControl = () => { parentPort.postMessage({ type: 'yield' }); }; + const exit = () => { + throw new CodeModeExitSignal(); + }; return Object.freeze({ + exit, image, load, output_image: image, @@ -271,8 +286,18 @@ function codeModeWorkerMain() { function createCodeModeModule(context, helpers) { return new SyntheticModule( - ['image', 'load', 'output_text', 'output_image', 'store', 'text', 'yield_control'], + [ + 'exit', + 'image', + 'load', + 'output_text', + 'output_image', + 'store', + 'text', + 'yield_control', + ], function initCodeModeModule() { + this.setExport('exit', helpers.exit); this.setExport('image', helpers.image); this.setExport('load', helpers.load); this.setExport('output_text', helpers.output_text); @@ -288,6 +313,7 @@ function codeModeWorkerMain() { function createBridgeRuntime(callTool, enabledTools, helpers) { return Object.freeze({ ALL_TOOLS: createAllToolsMetadata(enabledTools), + exit: helpers.exit, image: helpers.image, load: helpers.load, store: helpers.store, @@ -446,6 +472,13 @@ function codeModeWorkerMain() { stored_values: state.storedValues, }); } catch (error) { + if (isCodeModeExitSignal(error)) { + parentPort.postMessage({ + type: 'result', + stored_values: state.storedValues, + }); + return; + } parentPort.postMessage({ type: 'result', stored_values: state.storedValues, diff --git a/codex-rs/core/tests/suite/code_mode.rs b/codex-rs/core/tests/suite/code_mode.rs index 1ead4cb8929..e4f1f1a4975 100644 --- a/codex-rs/core/tests/suite/code_mode.rs +++ b/codex-rs/core/tests/suite/code_mode.rs @@ -1551,6 +1551,47 @@ text({ json: true }); Ok(()) } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn code_mode_exit_stops_script_immediately() -> Result<()> { + skip_if_no_network!(Ok(())); + + let server = responses::start_mock_server().await; + let (_test, second_mock) = run_code_mode_turn( + &server, + "use exec to stop script early with exit helper", + r#" +import { exit, text } from "@openai/code_mode"; + +text("before"); +exit(); +text("after"); +"#, + false, + ) + .await?; + + let req = second_mock.single_request(); + let items = custom_tool_output_items(&req, "call-1"); + let (output, success) = custom_tool_output_body_and_success(&req, "call-1"); + assert_ne!( + success, + Some(false), + "exec exit helper call failed unexpectedly: {output}" + ); + assert_eq!(items.len(), 2); + assert_regex_match( + concat!( + r"(?s)\A", + r"Script completed\nWall time \d+\.\d seconds\nOutput:\n\z" + ), + text_item(&items, 0), + ); + assert_eq!(text_item(&items, 1), "before"); + assert_eq!(output, "before"); + + Ok(()) +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn code_mode_surfaces_text_stringify_errors() -> Result<()> { skip_if_no_network!(Ok(())); @@ -1909,6 +1950,7 @@ text(JSON.stringify(Object.getOwnPropertyNames(globalThis).sort())); "encodeURI", "encodeURIComponent", "escape", + "exit", "eval", "globalThis", "image",