Code Mode

Code mode gives the model one tool that runs JavaScript or type-stripped TypeScript in Node.js inside the conversation's sandbox. The script can use Node APIs and installed packages, call an allowlist of your application tools, and return a small result. Intermediate data stays inside the script. Only the script, its return value and what it printed are saved on the thread.

Use it for calculations, file work, or combining several tool results: totals across months, a ranking, a join of two lists. For a single lookup, register the tool directly. Use a sandbox shell tool for shell commands, other languages, or streamed command output.

Use one script when the dependencies are clear but input values are not known until fetched. Await a tool's result and compute the next call's inputs from it, avoiding another model inference step just to pass data between tools. When enabled child tools are available, Code Mode's description gives this guidance to the model. The top-level sequential-call rules do not constrain calls inside a script; the script controls their order with await.

pnpm add @pai/code-mode @pai/sandbox @pai/sandbox-gke

Each script runs in a fresh Node.js process in the configured sandbox, whose image must include Node.js 22.13 or later. GKE provides process isolation and blocks outbound network access. Dependencies must already be in the image (installed outside /workspace) or workspace. Allowed application tools execute in the app and return results over the sandbox process channel. Generated code never executes in the app, including when the sandbox is unavailable.

Add a code tool

Start from ordinary tools. Give each an outputSchema: the model writes code against it.

import { defineTool } from "@pai/core";
import { z } from "zod";
import type { InvoiceStore } from "../services/invoices";

export const listInvoices = defineTool({
  id: "listInvoices",
  description: "List invoices for a month (YYYY-MM). Amounts are integer cents.",
  inputSchema: z.object({ month: z.string().regex(/^\d{4}-\d{2}$/) }),
  outputSchema: z.array(
    z.object({
      firmId: z.string(),
      amountCents: z.number().int(),
      status: z.enum(["paid", "void"]),
    }),
  ),
  contextSchema: z.object({
    workspaceId: z.string(),
    invoices: z.custom<InvoiceStore>(),
  }),
  execute: (ctx) =>
    ctx.context.invoices.list({
      workspaceId: ctx.context.workspaceId,
      month: ctx.input.month,
      signal: ctx.signal,
    }),
});

listFirms has the same shape. It takes {} and returns the workspace's firms as { id, name }[] from a firms store in its context.

Wrap the tools in codeModeTool. The keys are the names a script calls:

import { codeModeTool } from "@pai/code-mode";
import { listFirms, listInvoices } from "./invoices";

export const code = codeModeTool({
  tools: { listInvoices, listFirms },
});

Register code like any tool, and supply the context the children need:

export const assistantAgent = defineAgent({
  // …name, model, identity
  runtimeContext: async ({ identity }) => ({
    workspaceId: identity.workspaceId,
    invoices: createInvoiceStore(),
    firms: createFirmStore(),
  }),
  instructions:
    "For questions that need several reads or a calculation, write one script with the code tool. Return only the answer.",
  tools: { code },
});

Children inherit the agent's context, as if they were registered directly. If the agent lacks a field they need, registering code is a type error. TypeScript reports it as Property 'tool' is missing. Add the field to runtimeContext, or register { tool: code, mapContext }.

Bind the existing sandbox capability when creating PAI. Code mode and ordinary sandbox tools share the conversation's sandbox:

import { createPai } from "@pai/core";
import { sandboxCapability } from "@pai/sandbox";
import { createGkeSandboxProvider } from "@pai/sandbox-gke";

export const pai = createPai({
  agents: {
    main: {
      agent: assistantAgent,
      capabilities: {
        sandbox: sandboxCapability({
          provider: createGkeSandboxProvider({ namespace: "pai-sandbox-chat" }),
        }),
      },
    },
  },
  scopeKey: (identity) => identity.workspaceId,
});

PAI supplies the runner; the sandbox needs no npm installation or callback endpoint. See GKE setup for credentials and local kind transport.

What the model sees

One tool, code, with input { reason: string; js: string }. reason is a short user-facing label for the step; render it in your UI instead of the code. On every model step, its description declares the enabled children as TypeScript, the way the AI SDK's own code mode does:

declare const tools: {
  /** List invoices for a month (YYYY-MM). Amounts are integer cents. */
  listInvoices: (input: { month: string; }) => Promise<{ firmId: string; amountCents: number; status: "paid" | "void"; }[]>;
  // …listFirms
};

Each object property's description becomes a doc comment too. Validation rules such as a regex have no TypeScript form, and nor does a description on an array item or a union member, so state those in the description, as listInvoices does. If the agent also registers a child directly, under the same name, the step offers it to the model as its own tool, and code declares it by its result type alone: listInvoices: (input) => Promise<…>. The model already has its description and input schema.

The script calls them as await tools.listInvoices({ month: "2026-08" }) or await tools.listFirms({}).

The model then writes a script like this:

const [invoices, firms] = await Promise.all([
  tools.listInvoices({ month: "2026-08" }),
  tools.listFirms({}),
]);
const names = new Map(firms.map((firm) => [firm.id, firm.name]));
const totals = new Map();
for (const invoice of invoices) {
  if (invoice.status !== "paid") continue;
  totals.set(
    invoice.firmId,
    (totals.get(invoice.firmId) ?? 0) + invoice.amountCents,
  );
}
return [...totals]
  .map(([firmId, cents]) => ({ firm: names.get(firmId), cents }))
  .sort((a, b) => b.cents - a.cents)
  .slice(0, 3);

The three returned rows are the tool output's result. The invoice list never enters the conversation.

• The script is an async function body: top-level await and return work. Node strips TypeScript annotations without type-checking; syntax that needs transformation, such as enum, fails.
• Use require(...) or await import(...) for Node built-ins and installed packages. Static import declarations are not valid inside the function body.
• Node APIs can read and write sandbox files directly. Application tools remain limited to the explicit tools allowlist.
• The return value is JSON-serialized into the output's result. Return plain objects and arrays: a returned Map or Set becomes {}. A script with no return has no result.
• What the script prints through console, process.stdout or process.stderr comes back in the output's logs, in order, keeping the last maxResultBytes; logsTruncated: true means earlier output was dropped. When the script ends with an error, such as one it throws or a result over the limit, that error ends with its logs; a timeout, an oversized source, a process that ends before the script returns or a failure outside the script carries none. Output from a subprocess that inherits the script's stdio is not captured: print what the subprocess returns instead, as execSync does.

For example, a script can save data without a separate file-writing tool:

const { writeFile } = await import("node:fs/promises");
const invoices = await tools.listInvoices({ month: "2026-08" });
await writeFile("invoices.json", JSON.stringify(invoices, null, 2));
return { path: "invoices.json", count: invoices.length };

Relative paths use the sandbox workspace. Files survive subsequent code calls for as long as that sandbox exists; JavaScript variables do not.

How a call runs

1. The model calls code with a script.
2. PAI resolves which children are enabled, once. That answer holds until the script finishes.
3. PAI acquires the conversation's sandbox and starts a fresh Node process in its workspace. Script variables do not carry over between calls.
4. Each await tools.name(input) sends a request to the application, which runs that child through PAI, as a direct call would. PAI validates input and output, and applies context, capabilities, registration callbacks, and safe errors. The child's toModelOutput does not run: the script receives the JSON over stdin.
5. The script's return value and its logs become the output of the code call. PAI stops the runner and its ordinary child processes when the invocation ends. Await subprocess work before returning.

What is saved and observed

Child calls are not stored as parts and do not stream to the UI. Code mode keeps intermediate rows out of context, but it does not hide data from the model. Whatever the script returns, the model reads. Ask for compact results in instructions.

Control each child

An allowlist entry is either a tool or a registration with enabled, mapContext, and onOutput, the same shape as defineAgent({ tools }):

export const code = codeModeTool({
  contextSchema: z.object({
    workspaceId: z.string(),
    stores: z.custom<{ invoices: InvoiceStore; firms: FirmStore }>(),
    permissions: z.object({ canReadInvoices: z.boolean() }),
  }),
  tools: {
    listFirms: {
      tool: listFirms,
      mapContext: ({ runtimeContext }) => ({
        workspaceId: runtimeContext.workspaceId,
        firms: runtimeContext.stores.firms,
      }),
    },
    listInvoices: {
      tool: listInvoices,
      enabled: ({ runtimeContext }) =>
        runtimeContext.permissions.canReadInvoices,
      mapContext: ({ runtimeContext }) => ({
        workspaceId: runtimeContext.workspaceId,
        invoices: runtimeContext.stores.invoices,
      }),
    },
  },
});

An agent that registers this code tool supplies that shape:

runtimeContext: async ({ identity }) => ({
  workspaceId: identity.workspaceId,
  stores: await loadStores(identity.workspaceId),
  permissions: await loadPermissions(identity),
}),

• Without contextSchema, the code tool's context is the combined context of the entries registered without mapContext. Mapped entries add nothing the agent must supply. enabled, mapContext, and onOutput receive that context as runtimeContext.
• Declare contextSchema when a callback reads anything else. The example above needs it: both entries map, so without it the code tool's context would be empty, and the callbacks read workspaceId, stores, and permissions.
• Children are enabled by default. A disabled child is left out of the description. If a script calls it anyway, the call throws Tool is disabled: listInvoices, which the script can catch.

Limits

Use executionPolicy to bound each script and its application-tool requests:

codeModeTool({
  tools: { listInvoices, listFirms },
  executionPolicy: { timeoutMs: 10_000, maxBridgeRequests: 50 },
});

The defaults are a 30-second timeout, 256 child requests per script, 32 child calls executing at once, and a 32 KiB return value. Every option and default is on codeModeTool(). The sandbox deployment sets pod CPU and memory limits; there is no separate JavaScript heap or stack option.

The tool description states the result budget. An oversized result or error fails with its size, the limit and advice to return a summary or save full output to a sandbox file; successful JSON is never truncated.

• Source, result and time limits fail the whole script.
• Child calls beyond maxInFlightBridgeRequests (32) queue automatically until capacity is available. A large Promise.all needs no manual batching. Queued calls count toward the total request budget and share the script deadline.
• Total request count and child input/output limits deny the affected call; the script can catch that error.
• Guest limits do not bound your services. A child still owns its query size and must honor ctx.signal, which aborts on timeout as well as on Stop.

Errors

Rows marked script are catchable inside the script. An uncaught one becomes the code call's error. The other rows end the call, or the whole run, as described.

When the runner ends early, the error onRuntimeError receives has a CodeModeProcessExitError as its cause, and content-free logs keep that error's code, the exit code.

Failed calls are ordinary tool errors. See Failures.

Child tools

Reuse existing server tools, including tools that declare dataSchemas. During a child call, ctx.data.write(...) resolves without validating, retaining or emitting the value, and ctx.data.list(...) returns []. The script receives the tool's returned data; intermediate progress does not enter the chat. Direct agent-tool calls retain their normal ToolData behavior.

A child needs execute, and cannot declare suspend or its own tools. Listing either is a type error, and codeModeTool() throws when it is called: Nested tool "publish" cannot declare suspension or nested tools. Register approval tools and client tools on the agent directly.

Before production

• Sandbox capacity. Each invocation starts a fresh Node process in the conversation's sandbox. Pod CPU/memory limits apply to all its processes; warm-pool and namespace quotas remain deployment policy.
• Startup and execution. Acquiring a sandbox can include a cold start. Runner startup has a separate bounded allowance before the script's execution deadline begins; the script deadline includes time waiting for child tools.
• Retries. Started scripts are never automatically replayed on sandbox loss. Child calls are not checkpointed. When the model requests a new script, child writes that already happened may repeat; prefer read-only children and make writes idempotent.

Render the script

The default tool card works. For a custom card, input.js is a string and output is a CodeModeToolOutput: result (unknown), logs and logsTruncated:

function CodeCard({
  tool,
}: {
  tool: ToolRenderProps<AssistantPai["agents"]["main"], "code">;
}) {
  return (
    <figure>
      <pre>{tool.part.input?.js}</pre>
      {tool.part.state === "output-available" ? (
        <>
          {/* Calls saved before logs existed hold the bare value. */}
          <pre>{JSON.stringify(tool.part.output?.result, null, 2)}</pre>
          {tool.part.output?.logs && <pre>{tool.part.output.logs}</pre>}
        </>
      ) : tool.part.state === "output-error" ? (
        <p>{tool.part.errorText}</p>
      ) : (
        <p>Running…</p>
      )}
    </figure>
  );
}

Register it under the key you gave the tool (code). See Tool UI.

Test it

Use a runtime test, not runTool. Queue a code call with the script you expect. It runs for real against your tools and the fake services you pass. This test uses the agent from Add a code tool:

import { mkdtemp } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { InferAgentContract } from "@pai/core";
import { sandboxCapability } from "@pai/sandbox";
import {
  createStubSandboxProvider,
  startTestProcess,
} from "@pai/sandbox/test/fixtures";
import { createTestRuntime, createTestUtils } from "@pai/test-utils";
import { assistantAgent } from "./assistant-agent";

const workspace = await mkdtemp(join(tmpdir(), "code-mode-test-"));
const testSandboxProvider = createStubSandboxProvider({
  startProcess: (input) => startTestProcess({ ...input, cwd: workspace }),
}).provider;

const t = createTestUtils<InferAgentContract<typeof assistantAgent>>();
const model = t.model.queue((m) => [
  m.response([
    m.toolCall("code", {
      reason: "Counting August invoices",
      js: 'const rows = await tools.listInvoices({ month: "2026-08" }); return rows.length;',
    }),
  ]),
  m.response([m.text("There were 4 invoices.")]),
]);

await using runtime = createTestRuntime({
  agent: assistantAgent,
  model,
  capabilities: { sandbox: sandboxCapability({ provider: testSandboxProvider }) },
  overrideRuntimeContext: {
    workspaceId: "test-workspace",
    invoices: fakeInvoices,
    firms: fakeFirms,
  },
});
const client = runtime.client({
  identity: { userId: "test-user", workspaceId: "test-workspace" },
});
const thread = client.thread(client.newThreadId());
const state = await (
  await thread.send("How many invoices were there in August?")
).waitUntilIdle();

expect(t.getTool(state, "code").output?.result).toBe(4);

fakeInvoices returns four invoices here. startTestProcess runs the script in a local Node child process with no isolation, so use it only for fixed test scripts like this one. For integration coverage, use createGkeSandboxProvider with a local kind deployment instead. A scripted model makes the program deterministic; Node still runs the program for real.

Next

• codeModeTool() — every option and default
• Nested tools — call tools from your own execute
• Tool context
• Sandbox — Python, shell, and files