Persist js_repl codex helpers across cells (#14503)

## Summary This changes `js_repl` so saved references to `codex.tool(...)` and `codex.emitImage(...)` keep working across cells. Previously, those helpers were recreated per exec and captured that exec's `message.id`. If a persisted object or saved closure reused an old helper in a later cell, the nested tool/image call could fail with `js_repl exec context not found`. This patch: - keeps stable `codex.tool` and `codex.emitImage` helper identities in the kernel - resolves the current exec dynamically at call time using `AsyncLocalStorage` - adds regression coverage for persisted helper references across cells - updates the js_repl docs and project-doc instructions to describe the new behavior and its limits ## Why We already support persistent top-level bindings across `js_repl` cells, so persisted objects should be able to reuse `codex` helpers in later active cells. The bug was that helper identity was exec-scoped, not kernel-scoped. Using `AsyncLocalStorage` fixes the cross-cell reuse case without falling back to a single global active exec that could accidentally attribute stale background callbacks to the wrong cell.
2026-07-01 00:31:56 +08:00 · 2026-03-12 15:41:54 -07:00
parent a314c7d3ae
commit b560494c9f
5 changed files with 276 additions and 66 deletions
@@ -3,6 +3,7 @@
 // Requires Node started with --experimental-vm-modules.

 const { Buffer } = require("node:buffer");
+const { AsyncLocalStorage } = require("node:async_hooks");
 const crypto = require("node:crypto");
 const fs = require("node:fs");
 const { builtinModules, createRequire } = require("node:module");
@@ -126,6 +127,7 @@ const pendingTool = new Map();
 const pendingEmitImage = new Map();
 let toolCounter = 0;
 let emitImageCounter = 0;
+const execContextStorage = new AsyncLocalStorage();
 const cwd = process.cwd();
 const tmpDir = process.env.CODEX_JS_TMP_DIR || cwd;
 const homeDir = process.env.HOME ?? null;
@@ -1122,6 +1124,14 @@ function sendFatalExecResultSync(kind, error) {
  }
 }

+function getCurrentExecState() {
+  const execState = execContextStorage.getStore();
+  if (!execState || typeof execState.id !== "string" || !execState.id) {
+    throw new Error("js_repl exec context not found");
+  }
+  return execState;
+}
+
 function scheduleFatalExit(kind, error) {
  if (fatalExitScheduled) {
    process.exitCode = 1;
@@ -1427,15 +1437,21 @@ function normalizeEmitImageValue(value) {
  throw new Error("codex.emitImage received an unsupported value");
 }

-async function handleExec(message) {
-  clearLocalFileModuleCaches();
-  activeExecId = message.id;
-  const pendingBackgroundTasks = new Set();
-  const tool = (toolName, args) => {
+const codex = {
+  cwd,
+  homeDir,
+  tmpDir,
+  tool(toolName, args) {
+    let execState;
+    try {
+      execState = getCurrentExecState();
+    } catch (error) {
+      return Promise.reject(error);
+    }
    if (typeof toolName !== "string" || !toolName) {
      return Promise.reject(new Error("codex.tool expects a tool name string"));
    }
-    const id = `${message.id}-tool-${toolCounter++}`;
+    const id = `${execState.id}-tool-${toolCounter++}`;
    let argumentsJson = "{}";
    if (typeof args === "string") {
      argumentsJson = args;
@@ -1447,7 +1463,7 @@ async function handleExec(message) {
      const payload = {
        type: "run_tool",
        id,
-        exec_id: message.id,
+        exec_id: execState.id,
        tool_name: toolName,
        arguments: argumentsJson,
      };
@@ -1460,15 +1476,31 @@ async function handleExec(message) {
        resolve(res.response);
      });
    });
-  };
-  const emitImage = (imageLike) => {
+  },
+  emitImage(imageLike) {
+    let execState;
+    try {
+      execState = getCurrentExecState();
+    } catch (error) {
+      return {
+        then(onFulfilled, onRejected) {
+          return Promise.reject(error).then(onFulfilled, onRejected);
+        },
+        catch(onRejected) {
+          return Promise.reject(error).catch(onRejected);
+        },
+        finally(onFinally) {
+          return Promise.reject(error).finally(onFinally);
+        },
+      };
+    }
    const operation = (async () => {
      const normalized = normalizeEmitImageValue(await imageLike);
-      const id = `${message.id}-emit-image-${emitImageCounter++}`;
+      const id = `${execState.id}-emit-image-${emitImageCounter++}`;
      const payload = {
        type: "emit_image",
        id,
-        exec_id: message.id,
+        exec_id: execState.id,
        image_url: normalized.image_url,
        detail: normalized.detail ?? null,
      };
@@ -1489,7 +1521,7 @@ async function handleExec(message) {
      () => ({ ok: true, error: null, observation }),
      (error) => ({ ok: false, error, observation }),
    );
-    pendingBackgroundTasks.add(trackedOperation);
+    execState.pendingBackgroundTasks.add(trackedOperation);
    return {
      then(onFulfilled, onRejected) {
        observation.observed = true;
@@ -1504,6 +1536,15 @@ async function handleExec(message) {
        return operation.finally(onFinally);
      },
    };
+  },
+};
+
+async function handleExec(message) {
+  clearLocalFileModuleCaches();
+  activeExecId = message.id;
+  const execState = {
+    id: message.id,
+    pendingBackgroundTasks: new Set(),
  };

  let module = null;
@@ -1534,63 +1575,67 @@ async function handleExec(message) {
    priorBindings = builtSource.priorBindings;
    let output = "";

-    context.codex = { cwd, homeDir, tmpDir, tool, emitImage };
+    context.codex = codex;
    context.tmpDir = tmpDir;

-    await withCapturedConsole(context, async (logs) => {
-      const cellIdentifier = path.join(
-        cwd,
-        `.codex_js_repl_cell_${cellCounter++}.mjs`,
-      );
-      module = new SourceTextModule(source, {
-        context,
-        identifier: cellIdentifier,
-        initializeImportMeta(meta, mod) {
-          setImportMeta(meta, mod, true);
-          meta.__codexInternalMarkCommittedBindings = markCommittedBindings;
-          meta.__codexInternalMarkPreludeCompleted = markPreludeCompleted;
-        },
-        importModuleDynamically(specifier, referrer) {
-          return importResolved(resolveSpecifier(specifier, referrer?.identifier));
-        },
-      });
+    await execContextStorage.run(execState, async () => {
+      await withCapturedConsole(context, async (logs) => {
+        const cellIdentifier = path.join(
+          cwd,
+          `.codex_js_repl_cell_${cellCounter++}.mjs`,
+        );
+        module = new SourceTextModule(source, {
+          context,
+          identifier: cellIdentifier,
+          initializeImportMeta(meta, mod) {
+            setImportMeta(meta, mod, true);
+            meta.__codexInternalMarkCommittedBindings = markCommittedBindings;
+            meta.__codexInternalMarkPreludeCompleted = markPreludeCompleted;
+          },
+          importModuleDynamically(specifier, referrer) {
+            return importResolved(resolveSpecifier(specifier, referrer?.identifier));
+          },
+        });

-      await module.link(async (specifier) => {
-        if (specifier === "@prev" && previousModule) {
-          const exportNames = previousBindings.map((b) => b.name);
-          // Build a synthetic module snapshot of the prior cell's exports.
-          // This is the bridge that carries values from cell N to cell N+1.
-          const synthetic = new SyntheticModule(
-            exportNames,
-            function initSynthetic() {
-              for (const binding of previousBindings) {
-                this.setExport(
-                  binding.name,
-                  previousModule.namespace[binding.name],
-                );
-              }
-            },
-            { context },
+        await module.link(async (specifier) => {
+          if (specifier === "@prev" && previousModule) {
+            const exportNames = previousBindings.map((b) => b.name);
+            // Build a synthetic module snapshot of the prior cell's exports.
+            // This is the bridge that carries values from cell N to cell N+1.
+            const synthetic = new SyntheticModule(
+              exportNames,
+              function initSynthetic() {
+                for (const binding of previousBindings) {
+                  this.setExport(
+                    binding.name,
+                    previousModule.namespace[binding.name],
+                  );
+                }
+              },
+              { context },
+            );
+            return synthetic;
+          }
+          throw new Error(
+            `Top-level static import "${specifier}" is not supported in js_repl. Use await import("${specifier}") instead.`,
          );
-          return synthetic;
-        }
-        throw new Error(
-          `Top-level static import "${specifier}" is not supported in js_repl. Use await import("${specifier}") instead.`,
-        );
-      });
-      moduleLinked = true;
+        });
+        moduleLinked = true;

-      await module.evaluate();
-      if (pendingBackgroundTasks.size > 0) {
-        const backgroundResults = await Promise.all([...pendingBackgroundTasks]);
-        const firstUnhandledBackgroundError = backgroundResults.find(
-          (result) => !result.ok && !result.observation.observed,
-        );
-        if (firstUnhandledBackgroundError) {
-          throw firstUnhandledBackgroundError.error;
+        await module.evaluate();
+        if (execState.pendingBackgroundTasks.size > 0) {
+          const backgroundResults = await Promise.all([
+            ...execState.pendingBackgroundTasks,
+          ]);
+          const firstUnhandledBackgroundError = backgroundResults.find(
+            (result) => !result.ok && !result.observation.observed,
+          );
+          if (firstUnhandledBackgroundError) {
+            throw firstUnhandledBackgroundError.error;
+          }
        }
-      }
-      output = logs.join("\n");
+        output = logs.join("\n");
+      });
    });

    previousModule = module;
@@ -818,6 +818,87 @@ console.log("cell-complete");
    Ok(())
 }

+#[tokio::test]
+async fn js_repl_persisted_tool_helpers_work_across_cells() -> anyhow::Result<()> {
+    if !can_run_js_repl_runtime_tests().await {
+        return Ok(());
+    }
+
+    let (session, mut turn) = make_session_and_context().await;
+    turn.approval_policy
+        .set(AskForApproval::Never)
+        .expect("test setup should allow updating approval policy");
+    set_danger_full_access(&mut turn);
+
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
+    let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::default()));
+    let manager = turn.js_repl.manager().await?;
+
+    let global_marker = turn
+        .cwd
+        .join(format!("js-repl-global-helper-{}.txt", Uuid::new_v4()));
+    let lexical_marker = turn
+        .cwd
+        .join(format!("js-repl-lexical-helper-{}.txt", Uuid::new_v4()));
+    let global_marker_json = serde_json::to_string(&global_marker.to_string_lossy().to_string())?;
+    let lexical_marker_json = serde_json::to_string(&lexical_marker.to_string_lossy().to_string())?;
+
+    manager
+        .execute(
+            Arc::clone(&session),
+            Arc::clone(&turn),
+            Arc::clone(&tracker),
+            JsReplArgs {
+                code: format!(
+                    r#"
+const globalMarker = {global_marker_json};
+const lexicalMarker = {lexical_marker_json};
+const savedTool = codex.tool;
+globalThis.globalToolHelper = {{
+  run: () => savedTool("shell_command", {{ command: `printf global_helper > "${{globalMarker}}"` }}),
+}};
+const lexicalToolHelper = {{
+  run: () => savedTool("shell_command", {{ command: `printf lexical_helper > "${{lexicalMarker}}"` }}),
+}};
+"#
+                ),
+                timeout_ms: Some(10_000),
+            },
+        )
+        .await?;
+
+    let next = manager
+        .execute(
+            Arc::clone(&session),
+            Arc::clone(&turn),
+            tracker,
+            JsReplArgs {
+                code: r#"
+await globalToolHelper.run();
+await lexicalToolHelper.run();
+console.log("helpers-ran");
+"#
+                .to_string(),
+                timeout_ms: Some(10_000),
+            },
+        )
+        .await?;
+
+    assert!(next.output.contains("helpers-ran"));
+    assert_eq!(
+        tokio::fs::read_to_string(&global_marker).await?,
+        "global_helper"
+    );
+    assert_eq!(
+        tokio::fs::read_to_string(&lexical_marker).await?,
+        "lexical_helper"
+    );
+    let _ = tokio::fs::remove_file(&global_marker).await;
+    let _ = tokio::fs::remove_file(&lexical_marker).await;
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn js_repl_does_not_auto_attach_image_via_view_image_tool() -> anyhow::Result<()> {
    if !can_run_js_repl_runtime_tests().await {
@@ -1114,6 +1195,88 @@ console.log("cell-complete");
    Ok(())
 }

+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn js_repl_persisted_emit_image_helpers_work_across_cells() -> anyhow::Result<()> {
+    if !can_run_js_repl_runtime_tests().await {
+        return Ok(());
+    }
+
+    let (session, turn) = make_session_and_context().await;
+    if !turn
+        .model_info
+        .input_modalities
+        .contains(&InputModality::Image)
+    {
+        return Ok(());
+    }
+
+    let session = Arc::new(session);
+    let turn = Arc::new(turn);
+    *session.active_turn.lock().await = Some(crate::state::ActiveTurn::default());
+
+    let tracker = Arc::new(tokio::sync::Mutex::new(TurnDiffTracker::default()));
+    let manager = turn.js_repl.manager().await?;
+    let data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR4nGP4z8DwHwAFAAH/iZk9HQAAAABJRU5ErkJggg==";
+
+    manager
+        .execute(
+            Arc::clone(&session),
+            Arc::clone(&turn),
+            Arc::clone(&tracker),
+            JsReplArgs {
+                code: format!(
+                    r#"
+const dataUrl = "{data_url}";
+const savedEmitImage = codex.emitImage;
+globalThis.globalEmitHelper = {{
+  run: () => savedEmitImage(dataUrl),
+}};
+const lexicalEmitHelper = {{
+  run: () => savedEmitImage(dataUrl),
+}};
+"#
+                ),
+                timeout_ms: Some(15_000),
+            },
+        )
+        .await?;
+
+    let next = manager
+        .execute(
+            Arc::clone(&session),
+            Arc::clone(&turn),
+            tracker,
+            JsReplArgs {
+                code: r#"
+await globalEmitHelper.run();
+await lexicalEmitHelper.run();
+console.log("helpers-ran");
+"#
+                .to_string(),
+                timeout_ms: Some(15_000),
+            },
+        )
+        .await?;
+
+    assert!(next.output.contains("helpers-ran"));
+    assert_eq!(
+        next.content_items,
+        vec![
+            FunctionCallOutputContentItem::InputImage {
+                image_url: data_url.to_string(),
+                detail: None,
+            },
+            FunctionCallOutputContentItem::InputImage {
+                image_url: data_url.to_string(),
+                detail: None,
+            },
+        ]
+    );
+    assert!(session.get_pending_input().await.is_empty());
+
+    Ok(())
+}
+
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
 async fn js_repl_unawaited_emit_image_errors_fail_cell() -> anyhow::Result<()> {
    if !can_run_js_repl_runtime_tests().await {