code_mode: Move exec params from runtime declarations to @pragma

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

@@ -29,11 +29,6 @@ Object.defineProperty(globalThis, '__codexContentItems', {
   defineGlobal('ALL_TOOLS', __codexRuntime.ALL_TOOLS);
   defineGlobal('image', __codexRuntime.image);
   defineGlobal('load', __codexRuntime.load);
-  defineGlobal(
-    'set_max_output_tokens_per_exec_call',
-    __codexRuntime.set_max_output_tokens_per_exec_call
-  );
-  defineGlobal('set_yield_time', __codexRuntime.set_yield_time);
   defineGlobal('store', __codexRuntime.store);
   defineGlobal('text', __codexRuntime.text);
   defineGlobal('tools', __codexRuntime.tools);

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

@@ -1,16 +1,17 @@
 ## exec
 - Runs raw JavaScript in an isolated context (no Node, no file system, or network access, no console).
 - Send raw JavaScript source text, not JSON, quoted strings, or markdown code fences.
+- You may optionally start the tool input with a first-line pragma like `// @exec: {"yield_time_ms": 10000, "max_output_tokens": 1000}`.
+- `yield_time_ms` asks `exec` to yield early after that many milliseconds if the script is still running.
+- `max_output_tokens` sets the token budget for direct `exec` results. By default the result is truncated to 10000 tokens.
 - 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(...)`.
 - Tool methods take either string or object as parameter.
 - They return either a structured value or a string based on the description above.
 
 - Global helpers:
 - `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 serializeable value under a string key for later `exec` calls in the same session.
+- `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.
 - `ALL_TOOLS`: metadata for the enabled nested tools as `{ name, description }` entries.
-- `set_max_output_tokens_per_exec_call(value)`: sets the token budget for direct `exec` results. By default the result is truncated to 10000 tokens.
-- `set_yield_time(value)`: asks `exec` to yield early after that many milliseconds if the script is still running.
 - `yield_control()`: yields the accumulated output to the model immediately while the script keeps running.

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

@@ -1,4 +1,5 @@
 use async_trait::async_trait;
+use serde::Deserialize;
 
 use crate::codex::Session;
 use crate::codex::TurnContext;
@@ -9,6 +10,7 @@ use crate::tools::context::ToolPayload;
 use crate::tools::registry::ToolHandler;
 use crate::tools::registry::ToolKind;
 
+use super::CODE_MODE_PRAGMA_PREFIX;
 use super::CodeModeSessionProgress;
 use super::ExecContext;
 use super::PUBLIC_TOOL_NAME;
@@ -18,6 +20,23 @@ use super::protocol::HostToNodeMessage;
 use super::protocol::build_source;
 
 pub struct CodeModeExecuteHandler;
+const MAX_JS_SAFE_INTEGER: u64 = (1_u64 << 53) - 1;
+
+#[derive(Debug, Default, Deserialize, PartialEq, Eq)]
+#[serde(deny_unknown_fields)]
+struct CodeModeExecPragma {
+    #[serde(default)]
+    yield_time_ms: Option<u64>,
+    #[serde(default)]
+    max_output_tokens: Option<usize>,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+struct CodeModeExecArgs {
+    code: String,
+    yield_time_ms: Option<u64>,
+    max_output_tokens: Option<usize>,
+}
 
 impl CodeModeExecuteHandler {
     async fn execute(
@@ -26,12 +45,13 @@ impl CodeModeExecuteHandler {
         turn: std::sync::Arc<TurnContext>,
         code: String,
     ) -> Result<FunctionToolOutput, FunctionCallError> {
+        let args = parse_freeform_args(&code)?;
         let exec = ExecContext { session, turn };
         let enabled_tools = build_enabled_tools(&exec).await;
         let service = &exec.session.services.code_mode_service;
         let stored_values = service.stored_values().await;
         let source =
-            build_source(&code, &enabled_tools).map_err(FunctionCallError::RespondToModel)?;
+            build_source(&args.code, &enabled_tools).map_err(FunctionCallError::RespondToModel)?;
         let cell_id = service.allocate_cell_id().await;
         let request_id = service.allocate_request_id().await;
         let process_slot = service
@@ -46,6 +66,8 @@ impl CodeModeExecuteHandler {
             enabled_tools,
             stored_values,
             source,
+            yield_time_ms: args.yield_time_ms,
+            max_output_tokens: args.max_output_tokens,
         };
         let result = {
             let mut process_slot = process_slot;
@@ -72,6 +94,91 @@ impl CodeModeExecuteHandler {
     }
 }
 
+fn parse_freeform_args(input: &str) -> Result<CodeModeExecArgs, FunctionCallError> {
+    if input.trim().is_empty() {
+        return Err(FunctionCallError::RespondToModel(
+            "exec expects raw JavaScript source text (non-empty). Provide JS only, optionally with first-line `// @exec: {\"yield_time_ms\": 10000, \"max_output_tokens\": 1000}`.".to_string(),
+        ));
+    }
+
+    let mut args = CodeModeExecArgs {
+        code: input.to_string(),
+        yield_time_ms: None,
+        max_output_tokens: None,
+    };
+
+    let mut lines = input.splitn(2, '\n');
+    let first_line = lines.next().unwrap_or_default();
+    let rest = lines.next().unwrap_or_default();
+    let trimmed = first_line.trim_start();
+    let Some(pragma) = trimmed.strip_prefix(CODE_MODE_PRAGMA_PREFIX) else {
+        return Ok(args);
+    };
+
+    if rest.trim().is_empty() {
+        return Err(FunctionCallError::RespondToModel(
+            "exec pragma must be followed by JavaScript source on subsequent lines".to_string(),
+        ));
+    }
+
+    let directive = pragma.trim();
+    if directive.is_empty() {
+        return Err(FunctionCallError::RespondToModel(
+            "exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`"
+                .to_string(),
+        ));
+    }
+
+    let value: serde_json::Value = serde_json::from_str(directive).map_err(|err| {
+        FunctionCallError::RespondToModel(format!(
+            "exec pragma must be valid JSON with supported fields `yield_time_ms` and `max_output_tokens`: {err}"
+        ))
+    })?;
+    let object = value.as_object().ok_or_else(|| {
+        FunctionCallError::RespondToModel(
+            "exec pragma must be a JSON object with supported fields `yield_time_ms` and `max_output_tokens`"
+                .to_string(),
+        )
+    })?;
+    for key in object.keys() {
+        match key.as_str() {
+            "yield_time_ms" | "max_output_tokens" => {}
+            _ => {
+                return Err(FunctionCallError::RespondToModel(format!(
+                    "exec pragma only supports `yield_time_ms` and `max_output_tokens`; got `{key}`"
+                )));
+            }
+        }
+    }
+
+    let pragma: CodeModeExecPragma = serde_json::from_value(value).map_err(|err| {
+        FunctionCallError::RespondToModel(format!(
+            "exec pragma fields `yield_time_ms` and `max_output_tokens` must be non-negative safe integers: {err}"
+        ))
+    })?;
+    if pragma
+        .yield_time_ms
+        .is_some_and(|yield_time_ms| yield_time_ms > MAX_JS_SAFE_INTEGER)
+    {
+        return Err(FunctionCallError::RespondToModel(
+            "exec pragma field `yield_time_ms` must be a non-negative safe integer".to_string(),
+        ));
+    }
+    if pragma.max_output_tokens.is_some_and(|max_output_tokens| {
+        u64::try_from(max_output_tokens)
+            .map(|max_output_tokens| max_output_tokens > MAX_JS_SAFE_INTEGER)
+            .unwrap_or(true)
+    }) {
+        return Err(FunctionCallError::RespondToModel(
+            "exec pragma field `max_output_tokens` must be a non-negative safe integer".to_string(),
+        ));
+    }
+    args.code = rest.to_string();
+    args.yield_time_ms = pragma.yield_time_ms;
+    args.max_output_tokens = pragma.max_output_tokens;
+    Ok(args)
+}
+
 #[async_trait]
 impl ToolHandler for CodeModeExecuteHandler {
     type Output = FunctionToolOutput;
@@ -103,3 +210,7 @@ impl ToolHandler for CodeModeExecuteHandler {
         }
     }
 }
+
+#[cfg(test)]
+#[path = "execute_handler_tests.rs"]
+mod execute_handler_tests;

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

@@ -0,0 +1,41 @@
+use super::parse_freeform_args;
+use pretty_assertions::assert_eq;
+
+#[test]
+fn parse_freeform_args_without_pragma() {
+    let args = parse_freeform_args("output_text('ok');").expect("parse args");
+    assert_eq!(args.code, "output_text('ok');");
+    assert_eq!(args.yield_time_ms, None);
+    assert_eq!(args.max_output_tokens, None);
+}
+
+#[test]
+fn parse_freeform_args_with_pragma() {
+    let input = concat!(
+        "// @exec: {\"yield_time_ms\": 15000, \"max_output_tokens\": 2000}\n",
+        "output_text('ok');",
+    );
+    let args = parse_freeform_args(input).expect("parse args");
+    assert_eq!(args.code, "output_text('ok');");
+    assert_eq!(args.yield_time_ms, Some(15_000));
+    assert_eq!(args.max_output_tokens, Some(2_000));
+}
+
+#[test]
+fn parse_freeform_args_rejects_unknown_key() {
+    let err = parse_freeform_args("// @exec: {\"nope\": 1}\noutput_text('ok');")
+        .expect_err("expected error");
+    assert_eq!(
+        err.to_string(),
+        "exec pragma only supports `yield_time_ms` and `max_output_tokens`; got `nope`"
+    );
+}
+
+#[test]
+fn parse_freeform_args_rejects_missing_source() {
+    let err = parse_freeform_args("// @exec: {\"yield_time_ms\": 10}").expect_err("expected error");
+    assert_eq!(
+        err.to_string(),
+        "exec pragma must be followed by JavaScript source on subsequent lines"
+    );
+}

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

@@ -33,6 +33,7 @@ const CODE_MODE_RUNNER_SOURCE: &str = include_str!("runner.cjs");
 const CODE_MODE_BRIDGE_SOURCE: &str = include_str!("bridge.js");
 const CODE_MODE_DESCRIPTION_TEMPLATE: &str = include_str!("description.md");
 const CODE_MODE_WAIT_DESCRIPTION_TEMPLATE: &str = include_str!("wait_description.md");
+const CODE_MODE_PRAGMA_PREFIX: &str = "// @exec:";
 
 pub(crate) const PUBLIC_TOOL_NAME: &str = "exec";
 pub(crate) const WAIT_TOOL_NAME: &str = "exec_wait";
@@ -222,6 +223,7 @@ fn enabled_tool_from_spec(spec: ToolSpec) -> Option<protocol::EnabledTool> {
     }
 
     let reference = code_mode_tool_reference(&tool_name);
+    let global_name = normalize_code_mode_identifier(&tool_name);
     let (description, kind) = match spec {
         ToolSpec::Function(tool) => (tool.description, protocol::CodeModeToolKind::Function),
         ToolSpec::Freeform(tool) => (tool.description, protocol::CodeModeToolKind::Freeform),
@@ -234,8 +236,8 @@ fn enabled_tool_from_spec(spec: ToolSpec) -> Option<protocol::EnabledTool> {
     };
 
     Some(protocol::EnabledTool {
-        global_name: normalize_code_mode_identifier(&tool_name),
         tool_name,
+        global_name,
         module_path: reference.module_path,
         namespace: reference.namespace,
         name: normalize_code_mode_identifier(&reference.tool_key),

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

@@ -46,6 +46,8 @@ pub(super) enum HostToNodeMessage {
         enabled_tools: Vec<EnabledTool>,
         stored_values: HashMap<String, JsonValue>,
         source: String,
+        yield_time_ms: Option<u64>,
+        max_output_tokens: Option<usize>,
     },
     Poll {
         request_id: String,