Code Mode

Code mode lets a model write a script that calls tools, transforms their results, and returns a compact answer. For example, it can fetch several reports in parallel, filter thousands of rows, and return three totals. The intermediate rows are available to the script without automatically entering the model's conversation.

@pai/code-mode uses the AI SDK's code-mode implementation and Vercel's run library. It requires Node.js 22.13 or later. No sandbox cluster or container is needed for these JavaScript calculations.

Add the tool

pnpm add @pai/code-mode

Pass ordinary PAI tool definitions in an explicit map. Its keys become the names available to generated code. There is no automatic access to other tools registered on the agent.

import { codeModeTool } from "@pai/code-mode";
import { defineAgent, defineTool } from "@pai/core";
import { z } from "zod";

const sales = defineTool({
  id: "example.sales",
  description: "Read fictional daily sales for a region, in cents.",
  inputSchema: z.object({ region: z.enum(["north", "south"]) }),
  outputSchema: z.array(z.object({ amountCents: z.number().int() })),
  execute: ({ input }) =>
    input.region === "north"
      ? [{ amountCents: 1200 }, { amountCents: 800 }]
      : [{ amountCents: 1500 }, { amountCents: 500 }],
});

const agent = defineAgent({
  name: "sales-analyst",
  identity: z.object({ id: z.string() }),
  model,
  instructions:
    "Use code for calculations and sales analysis. Return concise results.",
  tools: {
    code: codeModeTool({ tools: { sales } }),
  },
});

model is your application's AI SDK language model. The SDK generates the code tool's description from the allowed tools' descriptions and schemas, including TypeScript signatures. Provide output schemas so the model can see the shape it will receive. The reference is regenerated each model step using the code tool's currently enabled children and their resolved descriptions.

Prefer plain JSON output schemas. The upstream signature generator uses the schema's input representation, so a schema that transforms its type can describe the pre-transform shape rather than the returned value.

The model calls code with native input { js: string }. A script might be:

const reports = await Promise.all([
  tools.sales({ region: "north" }),
  tools.sales({ region: "south" }),
]);
return {
  totalCents: reports.flat().reduce((sum, row) => sum + row.amountCents, 0),
};

The result is { totalCents: 4000 }, without an extra result envelope. Generated programs may return different JSON shapes, so the output type is honestly unknown. Tool inputs, tool outputs, context requirements, and execution-policy options remain inferred from their definitions. TypeScript declarations guide the model; generated source is not statically type-checked.

For calculations alone, use codeModeTool(). This exposes no host tools.

What runs where

Host tool implementations run in your Node application. The guest does not receive their closures, credentials, or application context. It receives only serialized inputs and outputs through the allowed functions.

Context, capabilities, and policy

Existing tools retain their required context and capability declarations. The enclosing agent must provide compatible context, and the runtime must bind capabilities required by nested tools. There is no separate code-mode service or credential configuration.

The nested map can use the normal registration shape when a child needs mapContext, enabled, or onOutput. Those settings belong to that registration. Registering the same definition directly on an agent does not implicitly copy that separate registration's policy into code mode.

Children are enabled by default. Use a boolean or an asynchronous or synchronous callback to control availability:

codeModeTool({
  contextSchema: z.object({ canReadSales: z.boolean() }),
  tools: {
    sales: {
      tool: sales,
      enabled: ({ runtimeContext }) => runtimeContext.canReadSales,
    },
  },
});

The enclosing agent must supply this context or map it when registering code. Description preparation resolves the currently enabled children for each model step and omits disabled children from the generated reference. Before the code tool's implementation starts, PAI separately resolves all child registrations once. Every call in that script enforces the same resolved availability; the next code-tool invocation resolves it afresh.

Disabling a separate direct registration on the agent leaves an enabled code-mode registration available. Share a callback explicitly when both paths should follow the same policy. Keep availability callbacks free of side effects: description preparation and parent invocation each resolve them. A child's mapContext, capabilities, validation and lifecycle policy still run when that enabled child is actually called. Description preparation also maps an enabled child's context when that child's description is dynamic; static child descriptions skip mapping. Keep context mapping free of side effects. If availability resolution fails or is cancelled, the parent implementation does not start; a failure is not converted to enabled: false.

Use contextSchema when the code tool's context differs from a child's. It types the registration callbacks and keeps the enclosing agent's context requirement explicit, just as on defineTool:

codeModeTool({
  contextSchema: z.object({ organization: z.string() }),
  tools: {
    lookup: {
      tool: tenantLookup, // Requires { tenant: string } as its context.
      mapContext: ({ runtimeContext }) => ({
        tenant: runtimeContext.organization,
      }),
      onOutput: ({ output }) => {
        // output is inferred from tenantLookup's output schema.
        recordLookup(output);
      },
    },
  },
});

An agent lifecycle hook still receives the full agent runtime context. A child registration receives the enclosing tool's mapped context; neither can silently replace the other's authority.

Lifecycle policy applies to each child call. Its input and result are validated; policy rejection prevents execution. Raw implementation exceptions stay behind PAI's safe error projection. A script can catch an ordinary tool failure. Lifecycle-hook failures observed while its parent call is active remain fatal, even if the script catches the child's rejection.

Execution limits

Use the AI SDK's native executionPolicy options when the defaults do not fit:

codeModeTool({
  tools: { sales },
  executionPolicy: {
    timeoutMs: 10_000,
    maxBridgeRequests: 100,
    maxInFlightBridgeRequests: 8,
    maxResultBytes: 64 * 1024,
  },
});

The upstream defaults include a 30-second deadline, 64 MiB guest heap, 256 host calls, and 32 simultaneous calls. Other options bound source, stack, result, console output, and individual tool payloads. The worker pool also has a process-wide concurrency cap. See the AI SDK code-mode reference for the native options.

Limits reject excessive work; they do not automatically batch a large Promise.all. Batch within the script when needed. Host tools must honor their abort signals and bound their own allocations and external requests. Guest memory limits do not cap the Node process or a database query's memory.

Supported behavior and limits

• Scripts support top-level await and return, JavaScript, and type-stripped TypeScript. Await all tool work; detached calls fail execution.
• Only ordinary server tools can be nested. Client tools, suspension, child ToolData, and further nesting are rejected. Use a direct PAI tool when the workflow needs an approval or resumable action.
• Nested calls return validated data to the script. Their toModelOutput callbacks do not run; native media conversion belongs to model-facing results.
• The script has no ambient Node, filesystem, network, or package-import access. An allowed tool can explicitly provide those powers. Its authorization still matters even though the guest is isolated.
• Intermediate results are not automatically saved as conversation parts. Explicitly returning them exposes them to the model, and guest console output goes to application logs. Application lifecycle hooks and telemetry may also record data under the application's existing policies.
• Code mode does not checkpoint each child call. It has no transactional or exactly-once guarantee across failures or retries. Keep side-effecting tools idempotent where needed; cancellation cannot undo completed external work.
• JSON serialization applies. Use integer units for money when possible; arbitrary-precision libraries, Python, shell commands, and persistent files belong in a configured sandbox.

The web lab's Code Mode: Order Analysis demo shows the submitted script and the exact result. Its first prompt aggregates fictional orders; its second calculates a forecast without host calls.

Use the same API in ordinary tools

codeModeTool uses the general defineTool.tools property. An ordinary backend tool can declare the same children and call their typed execute methods:

const totalSales = defineTool({
  id: "example.total-sales",
  description: "Calculate total sales across the requested regions.",
  inputSchema: z.object({
    regions: z.array(z.enum(["north", "south"])),
  }),
  tools: { sales },
  async execute({ input, tools }) {
    const reports = await Promise.all(
      input.regions.map(region => tools.sales.execute({ region })),
    );
    return {
      totalCents: reports.flat().reduce((sum, row) => sum + row.amountCents, 0),
    };
  },
});

This reuses sales, defineTool, and z from above. PAI preserves each child's input and output types, validation, availability, context, and cancellation. The application writes this workflow; code mode instead supplies the same bound tools to the SDK to run a model-written workflow.

Each descriptor also exposes a readonly enabled boolean. An implementation can check ctx.tools.sales.enabled before choosing to call ctx.tools.sales.execute(input). The method always exists and enforces that same value for the enclosing invocation. Validation, lifecycle policy and the tool's execution still apply after an enabled check.

The Tools Calling Tools lab demonstrates this generic API directly. Its compareTrips tool lists destinations, fetches their quotes with Promise.all, and returns sorted totals for the requested number of nights. When getQuote.enabled is false, it returns the destinations without fetching prices. Two agent registrations supply the context for these branches. Read server/compare-trips.ts to see the complete composition without a code evaluator.

Build a custom code tool

codeModeTool is built from public APIs. A consumer can connect its own defineTool to the SDK evaluator with the same tool definitions:

import {
  experimental_createCodeModeTool as createNativeCodeModeTool,
  experimental_runCodeMode as runCodeMode,
} from "@ai-sdk/code-mode";

const customCode = defineTool({
  id: "example.custom-code",
  inputSchema: z.object({ js: z.string() }),
  tools: { sales },
  description: ({ tools }) => createNativeCodeModeTool(tools).description,
  execute: (ctx) => runCodeMode({
    js: ctx.input.js,
    tools: ctx.tools,
    toolExecutionOptions: {
      toolCallId: ctx.thread.toolCallId,
      abortSignal: ctx.signal,
    },
  }),
});

The sales, defineTool, and z definitions above are reused here. Description metadata contains only enabled children; passing the whole map requires no per-entry checks. Execution receives SDK-shaped tools with their existing schemas and bound execute functions. PAI retains validation and invocation policy, so consumers need no schema conversion or per-child execution adapter. Ordinary application code can call ctx.tools.sales.execute({ region: "north" }); generated scripts continue to use tools.sales({ region: "north" }) through the SDK's guest interface. The enabled property belongs to the host descriptor; the guest API remains a map of callable functions.

The optional @pai/code-mode package supplies the reusable factory, execution-policy configuration, and evaluator-specific safe error conversion around this public API. Consumers can use it when they want the standard code-mode behavior.

Related APIs

PAI's reusable nested tool calls power this integration. Consumers normally use codeModeTool directly. The implementation uses @ai-sdk/code-mode, built on run, rather than maintaining an evaluator or worker protocol in PAI.