diff --git a/.changeset/code-mode-harness.md b/.changeset/code-mode-harness.md new file mode 100644 index 0000000000..6b2b43211d --- /dev/null +++ b/.changeset/code-mode-harness.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-code-mode': minor +--- + +Add `@tanstack/ai-code-mode/harness` with `codeMode({ driver })`: a harness plugin that gives the model one `execute_typescript` tool. The model writes a program that calls several tools, and the program runs in the isolate of any `@tanstack/ai-isolate-*` driver. By default, only tools that are safe to run without a question move into code mode: server tools with no `needsApproval` that the permission rules allow in plan mode, and never a tool that declares permission resources. `include` picks the tools yourself, a safe tool with `metadata.codeMode: true` always moves, and `lazy: true` lists only the tool names and adds `discover_tools`. Tool names that are not JavaScript identifiers get a safe name inside the program. `@tanstack/ai-harness` is an optional peer dependency. diff --git a/.changeset/code-mode-result-shape.md b/.changeset/code-mode-result-shape.md new file mode 100644 index 0000000000..98673986c6 --- /dev/null +++ b/.changeset/code-mode-result-shape.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-code-mode': patch +--- + +The `execute_typescript` tool description and the code-mode system prompt now tell the model that the tool returns `{ success, result, logs, error }`, where `result` is the value its code returned. diff --git a/docs/code-mode/code-mode.md b/docs/code-mode/code-mode.md index db693b6279..688e444f1a 100644 --- a/docs/code-mode/code-mode.md +++ b/docs/code-mode/code-mode.md @@ -418,3 +418,4 @@ pnpm eval -- --no-judge # skip Anthropic-based judging - [Showing Code Mode in the UI](./client-integration) — Display execution progress in your React app - [Code Mode with Snippets](./code-mode-with-snippets) — Add persistent, reusable snippet libraries - [Isolate Drivers](./code-mode-isolates): compare Node, QuickJS, QuickJS Bun, Cloudflare, Daytona, and E2B. +- [Code mode in a harness](../harness/code-mode): add code mode to a harness session as a plugin. diff --git a/docs/config.json b/docs/config.json index 6625d8b157..9b86e0f327 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1036,6 +1036,11 @@ "addedAt": "2026-10-13", "updatedAt": "2026-10-14" }, + { + "label": "Code mode in a harness", + "to": "harness/code-mode", + "addedAt": "2026-10-13" + }, { "label": "Run agents from a harness", "to": "harness/subagents", diff --git a/docs/harness/code-mode.md b/docs/harness/code-mode.md new file mode 100644 index 0000000000..266c022390 --- /dev/null +++ b/docs/harness/code-mode.md @@ -0,0 +1,137 @@ +--- +title: Code mode in a harness +id: harness-code-mode +order: 13 +description: "Let the harness model write one TypeScript program that calls many tools, and run it in an isolate. Any TanStack AI isolate driver plugs in." +keywords: + - tanstack ai + - harness + - code mode + - isolate + - quickjs +--- + +Your agent has many read-only tools, and it calls them one at a time. Each call is a model round trip, and every tool schema goes into every request. With code mode, the model writes one TypeScript program that calls the tools it needs, and the program runs in an isolate. The read-only tools leave the tool list and become functions in that program. + +## 1. Add the plugin + + + +react: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +vue: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +solid: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +svelte: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +preact: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +angular: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +octane: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs +vanilla: @tanstack/ai-code-mode @tanstack/ai-isolate-quickjs + + + +```ts group=harness-code-mode +import { toolDefinition } from '@tanstack/ai' +import { codeMode } from '@tanstack/ai-code-mode/harness' +import { defineHarness } from '@tanstack/ai-harness' +import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs' +import { openaiText } from '@tanstack/ai-openai' +import { z } from 'zod' + +const issues = [{ id: 'ENG-1', title: 'Fix login', state: 'open' }] + +const listIssues = toolDefinition({ + name: 'list_issues', + description: 'List the issues of the team', + inputSchema: z.object({ state: z.enum(['open', 'closed']) }), +}).server(async ({ state }) => issues.filter((issue) => issue.state === state)) + +const getIssue = toolDefinition({ + name: 'get_issue', + description: 'Get one issue by id', + inputSchema: z.object({ id: z.string() }), +}).server(async ({ id }) => issues.find((issue) => issue.id === id) ?? null) + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: openaiText('gpt-6.1-sol'), + tools: [listIssues, getIssue], + plugins: () => [codeMode({ driver: createQuickJSIsolateDriver() })], +}) +``` + +The model now has an `execute_typescript` tool. Inside the program, each moved tool is an `external_*` function, for example `external_list_issues({ state: 'open' })`. + +## 2. Pick the isolate + +`driver` takes any isolate driver. Change one line to run the programs somewhere else: + +| Driver | Package | Runs in | +| --- | --- | --- | +| `createQuickJSIsolateDriver()` | `@tanstack/ai-isolate-quickjs` | WebAssembly, on any runtime | +| `createQuickJSBunIsolateDriver()` | `@tanstack/ai-isolate-quickjs-bun` | Native QuickJS in Bun | +| `createNodeIsolateDriver()` | `@tanstack/ai-isolate-node` | A V8 isolate in Node | +| `createCloudflareIsolateDriver()` | `@tanstack/ai-isolate-cloudflare` | A Cloudflare Worker | +| `createDaytonaIsolateDriver()` | `@tanstack/ai-isolate-daytona` | A Daytona sandbox | +| `createE2BIsolateDriver()` | `@tanstack/ai-isolate-e2b` | An E2B sandbox | + +[Code mode isolates](../code-mode/code-mode-isolates) has the options of each driver. + +## 3. Choose which tools move + +A call inside the program does not stop for approval. So by default, a tool moves into code mode only when it is safe to run without a question: + +- The tool runs on the server and does not set `needsApproval`. +- The permission rules allow it in plan mode. A tool with an `edit` or `execute` rule stays. +- The tool does not declare `PermissionResources`. Calls inside the program skip the permission checks, so such a tool stays a normal tool call. + +Where the permission rules come from: + +- With [`permissions()`](./permissions) in the plugin list, code mode asks it about each tool. The answer is the one a call gets in plan mode: your `rules` win, and a tool with no rule gets your `default`. So with `default: 'ask'`, a tool with no rule stays. +- Without `permissions()`, code mode reads the rules that tool plugins add. + +Every other tool stays a normal tool call, with its approvals. With this default pick, code mode also adds `allow` rules for `execute_typescript` and `discover_tools`, so a program does not ask when every tool in it is safe. Your own `rules` still win. + +To pick the tools yourself, pass `include`: + +```ts group=harness-code-mode +export const onlyIssues = codeMode({ + driver: createQuickJSIsolateDriver(), + include: (tool) => tool.name.endsWith('_issue') || tool.name.endsWith('_issues'), +}) +``` + +- A tool that `include` picks moves even when it needs approval, and its calls inside the program do not ask. Pick only tools that are safe to run without a question. +- With `include`, a call of `execute_typescript` gets the permission rules like any tool, for example a question with `default: 'ask'`. +- A safe tool with `metadata.codeMode: true` moves even when `include` does not pick it. + +The other options of `createCodeMode` pass through: `timeout`, `memoryLimit`, `lazyToolsConfig`, and more. + +## 4. Send only the tool names + +The prompt of code mode lists each moved tool with its full TypeScript type. With 100 tools, that is a lot of tokens on every turn, and most turns use a few tools. Set `lazy: true`, and the prompt lists only the names: + +```ts group=harness-code-mode +export const lazyCodeMode = codeMode({ + driver: createQuickJSIsolateDriver(), + lazy: true, + lazyToolsConfig: { includeDescription: 'first-sentence' }, +}) +``` + +The model then works in three steps: + +1. It reads the tool names in the prompt. +2. It calls `discover_tools` with the names it needs, and gets their TypeScript signatures. +3. It calls those tools in `execute_typescript`. + +`includeDescription` sets what the prompt shows with each name: `'none'` (the default, the name only), `'first-sentence'`, or `'full'`. Tools that stay outside code mode keep their full schema. + +## Tools that appear while the session runs + +The plugin chooses the tools again before every turn. A plugin that finds tools at run time with `discoverTools` (for example after a sign-in) adds them first, and the safe ones move into code mode on the next turn. Tool names that are not JavaScript identifiers, for example `notion-search`, get a safe name inside the program: `external_notion_search`. When two tools get the same safe name, only the first one moves. The other one stays a normal tool call. + +## What you have now + +- One `execute_typescript` tool in place of many read-only tools. +- With `lazy: true`, a prompt that lists only the tool names. +- Programs that run in the isolate you choose. +- Approvals kept for every tool that stays a normal tool call. diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 65096431d9..8f3cf34f99 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -322,6 +322,7 @@ export const planTools = definePlugin({ - Pass what a call touches as a third argument, for example `decide('read_file', 'default', { paths: [file] })`. Then the path rules and `root` apply too. - Read it when the plugin runs, not in `setup`. Then `permissions()` can come before or after your plugin. - Saved `always` answers are not part of it. They can only allow, so a check that needs `'allow'` stays strict. +- [Code mode](./code-mode) uses this answer to pick the tools that move into a program. With `default: 'ask'`, a tool with no rule stays a normal tool call. - `PermissionRules`, `PermissionResources`, `PermissionDecisionCapability`, and `decidePermission` come from `@tanstack/ai-harness`. They load no Node module, so a plugin for an edge runtime can use them. The plugins themselves come from `@tanstack/ai-harness/plugins`, which is Node only. ## Known limits diff --git a/packages/ai-code-mode/package.json b/packages/ai-code-mode/package.json index 91435016a9..b95778ae69 100644 --- a/packages/ai-code-mode/package.json +++ b/packages/ai-code-mode/package.json @@ -24,6 +24,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./harness": { + "types": "./dist/esm/harness.d.ts", + "import": "./dist/esm/harness.js" } }, "sideEffects": false, @@ -65,10 +69,18 @@ }, "peerDependencies": { "@tanstack/ai": "workspace:^", + "@tanstack/ai-harness": "workspace:^", "zod": "^3.0.0 || ^4.0.0" }, + "peerDependenciesMeta": { + "@tanstack/ai-harness": { + "optional": true + } + }, "devDependencies": { "@tanstack/ai": "workspace:*", + "@tanstack/ai-harness": "workspace:*", + "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", "zod": "^4.2.0" } diff --git a/packages/ai-code-mode/src/create-code-mode-tool.ts b/packages/ai-code-mode/src/create-code-mode-tool.ts index 0d145b4efa..cb005ed0b9 100644 --- a/packages/ai-code-mode/src/create-code-mode-tool.ts +++ b/packages/ai-code-mode/src/create-code-mode-tool.ts @@ -365,6 +365,7 @@ function buildToolDescription(tools: Array): string { `Execute TypeScript code in a secure sandbox environment. ` + `The code can use these external API functions: ${externalFunctions}.${discoverable} ` + `All external_* calls are async and must be awaited. ` + - `Return a value to pass results back. Use console.log() for debugging.` + `Return a value to pass results back. Use console.log() for debugging. ` + + `The tool returns { success, result, logs, error }, where result is the value your code returned.` ) } diff --git a/packages/ai-code-mode/src/create-system-prompt.ts b/packages/ai-code-mode/src/create-system-prompt.ts index b878abf5bd..e2c72ebb52 100644 --- a/packages/ai-code-mode/src/create-system-prompt.ts +++ b/packages/ai-code-mode/src/create-system-prompt.ts @@ -102,7 +102,7 @@ return { warmestCity: warmest.location, temperature: warmest.temperature }; ### Important Notes - All \`external_*\` calls are async - always use \`await\` -- Return a value to pass results back to you +- Return a value to pass results back to you. The tool returns \`{ success, result, logs, error }\`, where \`result\` is the value your code returned - Use \`console.log()\` for debugging (logs are captured) - The sandbox is isolated - no network access or file system - Each execution is independent (no shared state between calls) diff --git a/packages/ai-code-mode/src/harness.ts b/packages/ai-code-mode/src/harness.ts new file mode 100644 index 0000000000..973c711d5f --- /dev/null +++ b/packages/ai-code-mode/src/harness.ts @@ -0,0 +1,152 @@ +import { + PermissionDecisionCapability, + PermissionResources, + PermissionRules, + decidePermission, + definePlugin, +} from '@tanstack/ai-harness' +import { createCodeMode } from './create-code-mode' +import type { AnyTool } from '@tanstack/ai' +import type { CodeModeTool, CodeModeToolConfig } from './types' + +export interface CodeModePluginOptions extends Omit< + CodeModeToolConfig, + 'tools' +> { + /** + * Which tools move into code mode. Default: every server tool that does + * not need approval and that `permissions()` allows in plan mode (so no + * edits, no commands, nothing that asks first). Without `permissions()`, + * the rules of tool plugins decide. A safe tool with + * `metadata.codeMode: true`, for example a tool of an `mcp()` server with + * `codeMode: true`, moves even when `include` does not pick it. + */ + include?: (tool: CodeModeTool) => boolean + /** + * Make each tool that moves into code mode lazy. The prompt then lists only + * the tool names, and `lazyToolsConfig.includeDescription` adds part of + * each description. The model calls `discover_tools` for the signatures of + * the tools it needs, then calls them in `execute_typescript`. This saves + * tokens on each turn when there are many tools. Default `false`. + */ + lazy?: boolean +} + +/** + * A name that works as a JavaScript identifier. Code mode turns each tool + * into an `external_` function, and MCP tool names often have `-`. + */ +function identifierOf(name: string): string { + const safe = name.replace(/[^A-Za-z0-9_$]/g, '_') + return /^[0-9]/.test(safe) ? `_${safe}` : safe +} + +/** A tool that runs on the server: not a client tool, and it has `execute`. */ +function isServerTool(tool: AnyTool): tool is CodeModeTool { + const side = '__toolSide' in tool ? tool.__toolSide : 'server' + return side === 'server' && typeof tool.execute === 'function' +} + +/** + * A harness plugin that gives the model an `execute_typescript` tool. The + * model writes one TypeScript program that calls several tools, and the + * program runs in the isolate of `driver` (any `@tanstack/ai-isolate-*` + * driver). The tools it calls leave the model's tool list, including tools + * that plugins find at run time, such as MCP tools after `/connect`. + * + * Calls inside the isolate do not stop for approval, so by default only + * tools that are safe to run without a question move into code mode. The + * other tools stay normal tool calls. A tool with `metadata.codeMode: true` + * moves even when `include` does not pick it, but only when it is safe in + * the same way. Calls inside the isolate also skip the permission checks, + * so a tool that declares `PermissionResources` never moves. With + * `lazy: true`, the model gets only the names of the moved tools and asks + * `discover_tools` for the signatures. + * + * @example + * ```ts + * import { codeMode } from '@tanstack/ai-code-mode/harness' + * import { createQuickJSIsolateDriver } from '@tanstack/ai-isolate-quickjs' + * + * defineHarness({ + * name: 'acme/agent', + * adapter, + * plugins: () => [codeMode({ driver: createQuickJSIsolateDriver() })], + * }) + * ``` + */ +export function codeMode(options: CodeModePluginOptions) { + const { include, lazy = false, ...config } = options + return definePlugin({ + name: 'tanstack/code-mode', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => { + const rules = ctx.collect(PermissionRules) + const resources = ctx.collect(PermissionResources) + // The system prompt for the tools of the current turn. + let prompt = '' + // Without `include` or snippets, only tools that plan mode allows + // move, so a program is as safe as its tools. Allow the code mode + // tools, so `default: 'ask'` does not ask for (or plan mode deny) + // every program. The user's rules come later and still win. + const isDefaultPick = !include && !config.getSnippetBindings + const allowCodeMode = isDefaultPick + ? ['execute_typescript', 'discover_tools'].map((tool) => + PermissionRules.item({ tool, decision: 'allow', kind: 'read' }), + ) + : [] + return { + prompts: [{ id: 'tanstack/code-mode', text: () => prompt }], + contribute: allowCodeMode, + prepareTools: ({ tools }) => { + // A failed run below must not leave the prompt of a past turn. + prompt = '' + // Ask permissions() when it is mounted: it knows all rules in its + // order and its default. Read here, so it can come after code mode. + const decide = ctx.getOptional(PermissionDecisionCapability) + const safe = (tool: CodeModeTool) => + !tool.needsApproval && + (decide?.(tool.name, 'plan') ?? + decidePermission(rules, tool.name, 'plan')) === 'allow' + const moves = (tool: CodeModeTool) => { + // The permission checks run on tool calls of the model. A call + // inside the isolate skips them, so such a tool stays a tool call. + const hasResources = resources.some((map) => + Object.hasOwn(map, tool.name), + ) + if (hasResources) return false + // A marked tool still needs the checks of the default pick. + const isMarked = tool.metadata?.codeMode === true + if (isMarked && safe(tool)) return true + return include ? include(tool) : safe(tool) + } + // Two tools that map to the same identifier: the first one moves. + const byIdentifier = new Map() + for (const tool of tools.filter(isServerTool)) { + if (!moves(tool)) continue + const identifier = identifierOf(tool.name) + if (!byIdentifier.has(identifier)) + byIdentifier.set(identifier, tool) + } + if (byIdentifier.size === 0) return tools + const created = createCodeMode({ + ...config, + tools: [...byIdentifier].map(([identifier, tool]) => + identifier === tool.name && !lazy + ? tool + : { ...tool, name: identifier, ...(lazy ? { lazy } : {}) }, + ), + }) + prompt = created.systemPrompt + const moved = new Set( + [...byIdentifier.values()].map((tool) => tool.name), + ) + return [ + ...tools.filter((tool) => !moved.has(tool.name)), + ...created.tools, + ] + }, + } + }, + }) +} diff --git a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts index 20115fb8c8..6ea087117e 100644 --- a/packages/ai-code-mode/tests/create-code-mode-tool.test.ts +++ b/packages/ai-code-mode/tests/create-code-mode-tool.test.ts @@ -57,6 +57,7 @@ describe('createCodeModeTool', () => { }) expect(tool.description).toContain('external_fetchWeather') expect(tool.description).toContain('external_dbQuery') + expect(tool.description).toContain('{ success, result, logs, error }') }) it('tool description says discover_tools is a separate tool', () => { diff --git a/packages/ai-code-mode/tests/create-system-prompt.test.ts b/packages/ai-code-mode/tests/create-system-prompt.test.ts index 4f91c88187..1cbd07ee09 100644 --- a/packages/ai-code-mode/tests/create-system-prompt.test.ts +++ b/packages/ai-code-mode/tests/create-system-prompt.test.ts @@ -56,6 +56,7 @@ describe('createCodeModeSystemPrompt', () => { expect(prompt).toContain('When to Use') expect(prompt).toContain('Available External APIs') expect(prompt).toContain('Important Notes') + expect(prompt).toContain('`{ success, result, logs, error }`') }) it('documents multiple tools', () => { diff --git a/packages/ai-code-mode/tests/harness.test.ts b/packages/ai-code-mode/tests/harness.test.ts new file mode 100644 index 0000000000..fb5cf9de0a --- /dev/null +++ b/packages/ai-code-mode/tests/harness.test.ts @@ -0,0 +1,587 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { + PermissionRules, + createHarnessHost, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' +import { PermissionResources, permissions } from '@tanstack/ai-harness/plugins' +import { codeMode } from '../src/harness' +import type { AnyTextAdapter, AnyTool, StreamChunk } from '@tanstack/ai' +import type { HarnessPlugin } from '@tanstack/ai-harness' +import type { ExecutionResult, IsolateDriver, ToolBinding } from '../src/types' + +const now = () => Date.now() + +/** A model whose turns are scripted. Records the options of every call. */ +function scripted(turns: Array>) { + const calls: Array = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': { + providerOptions: {} as Record, + inputModalities: ['text'] as readonly ['text'], + messageMetadataByModality: { + text: undefined as unknown, + image: undefined as unknown, + audio: undefined as unknown, + video: undefined as unknown, + document: undefined as unknown, + }, + toolCapabilities: [] as ReadonlyArray, + toolCallMetadata: undefined as unknown, + systemPromptMetadata: undefined as never, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options: any) => { + calls.push(options) + const chunks = turns[calls.length - 1] ?? [] + return (async function* () { + yield* chunks + })() + }, + } + return { adapter, calls } +} + +const textTurn = (content: string): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'm', + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm', + delta: content, + timestamp: now(), + }, + { type: EventType.TEXT_MESSAGE_END, messageId: 'm', timestamp: now() }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, +] + +/** A model turn that calls one tool with `args`. */ +const callTurn = (toolName: string, args: unknown): Array => [ + { type: EventType.RUN_STARTED, runId: 'r', threadId: 't', timestamp: now() }, + { + type: EventType.TOOL_CALL_START, + toolCallId: `call-${toolName}`, + toolCallName: toolName, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: `call-${toolName}`, + delta: JSON.stringify(args), + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_END, + toolCallId: `call-${toolName}`, + timestamp: now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + }, +] + +const codeTurn = (code: string) => + callTurn('execute_typescript', { typescriptCode: code }) + +/** + * A stand-in isolate: it records the bindings it got and, for any program, + * calls `external_lookup` and `external_notion_search` and returns both + * results. Real drivers run the program itself. + */ +function fakeDriver() { + const seen: Array> = [] + const driver: IsolateDriver = { + createContext: async ({ bindings }) => { + seen.push(Object.keys(bindings).sort()) + const call = (name: string, args: unknown) => + (bindings[name] as ToolBinding).execute(args) + return { + execute: async (): Promise> => { + const value = { + lookup: await call('external_lookup', { q: 'x' }), + notion: await call('external_notion_search', { query: 'y' }), + } + // The caller picks T. This stand-in always returns the two results. + return { success: true, value: value as T, logs: [] } + }, + dispose: async () => {}, + } + }, + } + return { driver, seen } +} + +const lookup = vi.fn(async () => 'fact') +const notionSearch = vi.fn(async () => ({ results: ['page'] })) +const tools = { + lookup: toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({ q: z.string() }), + }).server(lookup), + remove: toolDefinition({ + name: 'remove', + description: 'Remove a file', + needsApproval: true, + inputSchema: z.object({ path: z.string() }), + }).server(async () => 'removed'), + writeFile: toolDefinition({ + name: 'write_file', + description: 'Write a file', + inputSchema: z.object({ path: z.string() }), + }).server(async () => 'written'), + notionSearch: toolDefinition({ + name: 'notion_search', + description: 'Search Notion', + inputSchema: z.object({ query: z.string() }), + }).server(notionSearch), + askUser: toolDefinition({ + name: 'ask_user', + description: 'Ask the user in the browser', + inputSchema: z.object({}), + }).client(), +} + +// A tool plugin that marks write_file as an edit, like workspaceTools does. +const writes = definePlugin({ + name: 'test/writes', + setup: () => ({ + tools: [tools.writeFile], + contribute: [ + PermissionRules.item({ + tool: 'write_file', + decision: 'ask', + kind: 'edit', + }), + ], + }), +}) +// A connector whose tool appears at run time. +const notion = definePlugin({ + name: 'test/notion', + setup: () => ({ discoverTools: () => [tools.notionSearch] }), +}) + +/** A server tool marked for code mode, like the tools of `mcp({ codeMode })`. */ +function markedTool(name: string, needsApproval = false) { + return toolDefinition({ + name, + description: 'A marked tool', + inputSchema: z.object({}), + needsApproval, + metadata: { codeMode: true }, + }).server(async () => name) +} + +/** Opens a code mode session and returns the tool names of the first model call. */ +async function firstToolNames(options: { + tools: ReadonlyArray + plugins: ReadonlyArray +}) { + const { adapter, calls } = scripted([textTurn('ok')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-marked', + adapter, + tools: options.tools, + plugins: () => options.plugins, + }), + { threadId: 't' }, + ) + await session.prompt('hi') + await host.close() + return calls[0].tools.map((tool: { name: string }) => tool.name) +} + +describe('codeMode', () => { + it('moves safe tools, including discovered ones, behind execute_typescript', async () => { + const { driver, seen } = fakeDriver() + const { adapter, calls } = scripted([ + codeTurn('return { a: await external_lookup({ q: "x" }) }'), + textTurn('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode', + adapter, + tools: [tools.lookup, tools.remove, tools.askUser], + plugins: () => [writes, notion, codeMode({ driver })], + }), + { threadId: 't' }, + ) + const turn = await session.prompt('look it up') + expect(turn.text).toBe('done') + + const names = calls[0].tools.map((tool: { name: string }) => tool.name) + expect(names).toEqual([ + 'remove', + 'ask_user', + 'write_file', + 'execute_typescript', + ]) + const prompts = JSON.stringify(calls[0].systemPrompts) + expect(prompts).toContain('external_lookup') + expect(prompts).toContain('external_notion_search') + expect(prompts).not.toContain('external_write_file') + + expect(seen[0]).toEqual(['external_lookup', 'external_notion_search']) + expect(lookup).toHaveBeenCalledTimes(1) + expect(notionSearch).toHaveBeenCalledTimes(1) + expect(JSON.stringify(calls[1].messages)).toContain('page') + await host.close() + }) + + it('uses include to pick the tools, and leaves the list alone when none match', async () => { + const { driver } = fakeDriver() + const { adapter, calls } = scripted([textTurn('one'), textTurn('two')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const picked = await host.open( + defineHarness({ + name: 'test/code-mode-include', + adapter, + tools: [tools.lookup, tools.remove], + plugins: () => [ + codeMode({ driver, include: (tool) => tool.name === 'remove' }), + ], + }), + { threadId: 'picked' }, + ) + await picked.prompt('hi') + expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'lookup', + 'execute_typescript', + ]) + + const none = await host.open( + defineHarness({ + name: 'test/code-mode-none', + adapter, + tools: [tools.remove], + plugins: () => [codeMode({ driver })], + }), + { threadId: 'none' }, + ) + await none.prompt('hi') + expect(calls[1].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'remove', + ]) + expect(JSON.stringify(calls[1].systemPrompts ?? [])).not.toContain( + 'execute_typescript', + ) + await host.close() + }) + + it('with lazy, lists only the tool names and gives signatures on discover_tools', async () => { + const { driver } = fakeDriver() + const { adapter, calls } = scripted([ + callTurn('discover_tools', { toolNames: ['external_lookup'] }), + textTurn('found it'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-lazy', + adapter, + tools: [tools.lookup, tools.remove], + plugins: () => [codeMode({ driver, lazy: true })], + }), + { threadId: 't' }, + ) + expect((await session.prompt('what can you call?')).text).toBe('found it') + + expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'remove', + 'execute_typescript', + 'discover_tools', + ]) + // The prompt names the tool, but its input shape comes from discovery. + const prompts = JSON.stringify(calls[0].systemPrompts) + expect(prompts).toContain('Discoverable APIs') + expect(prompts).toContain('external_lookup') + expect(prompts).not.toContain('q: string') + expect(JSON.stringify(calls[1].messages)).toContain('q: string') + await host.close() + }) + + it('gives tools with names that are not identifiers a safe name', async () => { + const seen: Array> = [] + const driver: IsolateDriver = { + createContext: async ({ bindings }) => { + seen.push(Object.keys(bindings).sort()) + return { + execute: async (): Promise> => { + const binding = bindings['external_web_search'] as ToolBinding + const value: unknown = await binding.execute({}) + // The caller picks T. This stand-in returns the one result. + return { success: true, value: value as T, logs: [] } + }, + dispose: async () => {}, + } + }, + } + const named = (name: string, reply: string) => + toolDefinition({ + name, + description: 'A tool', + inputSchema: z.object({}), + }).server(async () => reply) + const { adapter, calls } = scripted([codeTurn('return 1'), textTurn('ok')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-names', + adapter, + tools: [ + named('web-search', 'from web-search'), + named('web_search', 'from web_search'), + named('2fa', 'code'), + ], + plugins: () => [codeMode({ driver })], + }), + { threadId: 't' }, + ) + await session.prompt('search') + // web_search clashes with web-search after the rename, so it stays direct. + expect(calls[0].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'web_search', + 'execute_typescript', + ]) + expect(seen[0]).toEqual(['external__2fa', 'external_web_search']) + expect(JSON.stringify(calls[1].messages)).toContain('from web-search') + await host.close() + }) + + it('moves a marked tool that include does not pick, unless a rule or an approval stops it', async () => { + const rules = definePlugin({ + name: 'test/marked-rules', + setup: () => ({ + contribute: [ + PermissionRules.item({ tool: 'docs_allowed', decision: 'allow' }), + PermissionRules.item({ tool: 'docs_ask', decision: 'ask' }), + PermissionRules.item({ tool: 'docs_deny', decision: 'deny' }), + ], + }), + }) + const names = await firstToolNames({ + tools: [ + markedTool('docs_allowed'), + markedTool('docs_ask'), + markedTool('docs_deny'), + markedTool('docs_write', true), + ], + plugins: [ + rules, + codeMode({ driver: fakeDriver().driver, include: () => false }), + ], + }) + expect(names).toEqual([ + 'docs_ask', + 'docs_deny', + 'docs_write', + 'execute_typescript', + ]) + }) + + it('keeps a tool that a rule of permissions() asks for as a tool call', async () => { + const names = await firstToolNames({ + tools: [ + markedTool('docs_search'), + markedTool('notes_search'), + tools.lookup, + ], + plugins: [ + permissions({ + rules: [ + { tool: 'docs_*', decision: 'ask' }, + { tool: 'lookup', decision: 'ask' }, + ], + }), + codeMode({ driver: fakeDriver().driver }), + ], + }) + // A call inside the isolate skips the question, so both tools stay. + expect(names).toEqual(['docs_search', 'lookup', 'execute_typescript']) + }) + + it.each(['before', 'after'])( + 'keeps a tool with no rule as a tool call when permissions() asks by default, mounted %s code mode', + async (order) => { + const guard = permissions({ + default: 'ask', + rules: [{ tool: 'lookup', decision: 'allow' }], + }) + const code = codeMode({ driver: fakeDriver().driver }) + const names = await firstToolNames({ + tools: [markedTool('docs_search'), tools.lookup], + plugins: order === 'before' ? [guard, code] : [code, guard], + }) + expect(names).toEqual(['docs_search', 'execute_typescript']) + }, + ) + + it('keeps a tool that a rule of permissions() asks for, when a later plugin allows it', async () => { + const later = definePlugin({ + name: 'test/later-allow', + setup: () => ({ + contribute: [ + PermissionRules.item({ tool: 'docs_search', decision: 'allow' }), + ], + }), + }) + const names = await firstToolNames({ + tools: [markedTool('docs_search'), tools.lookup], + plugins: [ + permissions({ rules: [{ tool: 'docs_search', decision: 'ask' }] }), + later, + codeMode({ driver: fakeDriver().driver }), + ], + }) + // permissions() checks its own rules last, so the call would ask. + expect(names).toEqual(['docs_search', 'execute_typescript']) + }) + + it('keeps a tool that declares permission resources as a tool call', async () => { + const resources = definePlugin({ + name: 'test/resources', + setup: () => ({ + contribute: [ + PermissionResources.item({ read_file: { paths: () => ['a.txt'] } }), + ], + }), + }) + const names = await firstToolNames({ + tools: [tools.lookup, markedTool('read_file')], + plugins: [resources, codeMode({ driver: fakeDriver().driver })], + }) + expect(names).toEqual(['read_file', 'execute_typescript']) + }) + + it.each(['default', 'plan'])( + 'runs a program with no question when permissions() asks by default, in %s mode', + async (mode) => { + const { driver, seen } = fakeDriver() + const { adapter } = scripted([codeTurn('return 1'), textTurn('done')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-allow', + adapter, + tools: [tools.lookup], + plugins: () => [ + permissions({ + default: 'ask', + rules: [ + { tool: 'lookup', decision: 'allow' }, + { tool: 'notion_search', decision: 'allow' }, + ], + }), + notion, + codeMode({ driver }), + ], + }), + { threadId: 't' }, + ) + if (mode === 'plan') await session.command('mode', 'plan') + expect((await session.prompt('look it up')).text).toBe('done') + // The program ran: a question would wait, and plan mode would deny it. + expect(seen).toHaveLength(1) + await host.close() + }, + ) + + it('asks before a program when a rule of the user asks for execute_typescript', async () => { + const { driver, seen } = fakeDriver() + const { adapter } = scripted([codeTurn('return 1'), textTurn('done')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-user-rule', + adapter, + tools: [tools.lookup], + plugins: () => [ + permissions({ + rules: [{ tool: 'execute_typescript', decision: 'ask' }], + }), + notion, + codeMode({ driver }), + ], + }), + { threadId: 't' }, + ) + const turn = session.prompt('look it up') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + if (!question) throw new Error('No question.') + expect(question.message).toContain('execute_typescript') + await session.answer({ + questionId: question.questionId, + value: { answer: 'reject' }, + }) + expect((await turn).text).toBe('done') + expect(seen).toEqual([]) + await host.close() + }) + + it('drops the prompt of the last turn when the pick fails', async () => { + const { driver } = fakeDriver() + const { adapter, calls } = scripted([textTurn('one'), textTurn('two')]) + let isBroken = false + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/code-mode-failed-pick', + adapter, + tools: [tools.lookup], + plugins: () => [ + codeMode({ + driver, + include: () => { + if (isBroken) throw new Error('The pick failed.') + return true + }, + }), + ], + }), + { threadId: 't' }, + ) + await session.prompt('first') + expect(JSON.stringify(calls[0].systemPrompts)).toContain('external_lookup') + + isBroken = true + await session.prompt('second') + // The harness keeps the tool list as it was, without the old prompt. + expect(calls[1].tools.map((tool: { name: string }) => tool.name)).toEqual([ + 'lookup', + ]) + expect(JSON.stringify(calls[1].systemPrompts ?? [])).not.toContain( + 'external_lookup', + ) + await host.close() + }) +}) diff --git a/packages/ai-code-mode/vite.config.ts b/packages/ai-code-mode/vite.config.ts index 77bcc2e60b..525f4daf73 100644 --- a/packages/ai-code-mode/vite.config.ts +++ b/packages/ai-code-mode/vite.config.ts @@ -29,7 +29,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/harness.ts'], srcDir: './src', cjs: false, }), diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index f64634dd3d..3bf31b5d36 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2275,6 +2275,12 @@ importers: '@tanstack/ai': specifier: workspace:* version: link:../ai + '@tanstack/ai-harness': + specifier: workspace:* + version: link:../ai-harness + '@tanstack/ai-persistence': + specifier: workspace:* + version: link:../ai-persistence '@vitest/coverage-v8': specifier: 4.1.10 version: 4.1.10(vitest@4.1.11) @@ -3706,6 +3712,9 @@ importers: '@tanstack/ai-cloudflare': specifier: workspace:* version: link:../../packages/ai-cloudflare + '@tanstack/ai-code-mode': + specifier: workspace:* + version: link:../../packages/ai-code-mode '@tanstack/ai-compaction': specifier: workspace:* version: link:../../packages/ai-compaction diff --git a/testing/e2e/fixtures/harness-first-party/basic.json b/testing/e2e/fixtures/harness-first-party/basic.json index 77195c081f..22bb740ed8 100644 --- a/testing/e2e/fixtures/harness-first-party/basic.json +++ b/testing/e2e/fixtures/harness-first-party/basic.json @@ -102,6 +102,14 @@ "response": { "content": "Here is a plan for Rome." } + }, + { + "match": { + "userMessage": "[first-party-code-mode] which tools" + }, + "response": { + "content": "I can run TypeScript." + } } ] } diff --git a/testing/e2e/package.json b/testing/e2e/package.json index 2262abbbba..c026281d73 100644 --- a/testing/e2e/package.json +++ b/testing/e2e/package.json @@ -25,6 +25,7 @@ "@tanstack/ai-claude-code": "workspace:*", "@tanstack/ai-client": "workspace:*", "@tanstack/ai-cloudflare": "workspace:*", + "@tanstack/ai-code-mode": "workspace:*", "@tanstack/ai-compaction": "workspace:*", "@tanstack/ai-elevenlabs": "workspace:*", "@tanstack/ai-fal": "workspace:*", diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts index 8dc25c4ef9..814b622241 100644 --- a/testing/e2e/src/routes/api.harness-first-party.ts +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -4,6 +4,7 @@ import { tmpdir } from 'node:os' import { join } from 'node:path' import { promisify } from 'node:util' import { createFileRoute } from '@tanstack/react-router' +import { codeMode } from '@tanstack/ai-code-mode/harness' import { defineAgent, toolDefinition } from '@tanstack/ai' import { createHarnessHost, @@ -194,9 +195,41 @@ async function sessionTitle( return { text: turn.text, title: entry?.title, failures } } +/** + * `codeMode()` from `@tanstack/ai-code-mode/harness`: the safe `list_issues` + * tool leaves the tool list and becomes `external_list_issues` in the code + * mode prompt. The model answers with text, so no program runs and the + * driver is never used. The spec reads the request in aimock's journal. + */ +async function codeModeTools(host: HarnessHost, model: () => AnyTextAdapter) { + const listIssues = toolDefinition({ + name: 'list_issues', + description: 'List the open issues', + inputSchema: z.object({}), + }).server(async () => [{ id: 'ENG-1', title: 'Fix login' }]) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-code-mode', + adapter: model(), + tools: [listIssues], + plugins: () => [ + codeMode({ + driver: { + createContext: () => + Promise.reject(new Error('This scenario runs no program.')), + }, + }), + ], + }), + { threadId: 'e2e-code-mode' }, + ) + const turn = await session.prompt('[first-party-code-mode] which tools') + return { text: turn.text } +} + /** * First-party harness plugins, one scenario per request: - * `?scenario=agents|plan-subagent|env|title&testId=...`. The models are + * `?scenario=agents|plan-subagent|env|title|code-mode&testId=...`. The models are * OpenAI adapters against aimock. */ export const Route = createFileRoute('/api/harness-first-party')({ @@ -228,6 +261,9 @@ export const Route = createFileRoute('/api/harness-first-party')({ ).adapter return Response.json(await sessionTitle(host, model, titleModel)) } + if (scenario === 'code-mode') { + return Response.json(await codeModeTools(host, model)) + } return Response.json( { error: `Unknown scenario ${scenario}` }, { status: 400 }, diff --git a/testing/e2e/tests/harness-first-party.spec.ts b/testing/e2e/tests/harness-first-party.spec.ts index 058cc9419d..4c2b950613 100644 --- a/testing/e2e/tests/harness-first-party.spec.ts +++ b/testing/e2e/tests/harness-first-party.spec.ts @@ -3,7 +3,7 @@ import type { APIRequestContext } from '@playwright/test' /** * First-party harness plugins in a session: `agents()`, `permissions()`, - * `projectInstructions()`, and `title()`. Each test runs one scenario of + * `projectInstructions()`, `title()`, and `codeMode()`. Each test runs one scenario of * `/api/harness-first-party`. The model is the OpenAI adapter against * aimock, and aimock's journal shows what each model call sent. */ @@ -12,6 +12,7 @@ type JournalEntry = { headers?: Record body: { messages?: Array<{ role: string; content: unknown }> + tools?: Array<{ function?: { name?: string } }> tool_choice?: unknown } | null } @@ -122,4 +123,20 @@ test.describe('harness first-party plugins', () => { failures: [], }) }) + + test('codeMode(): a safe tool moves behind execute_typescript', async ({ + request, + testId, + aimockPort, + }) => { + expect(await runScenario(request, 'code-mode', testId)).toEqual({ + text: 'I can run TypeScript.', + }) + + const [call] = await modelCalls(request, aimockPort, testId) + const tools = (call?.body?.tools ?? []).map((tool) => tool.function?.name) + expect(tools).toContain('execute_typescript') + expect(tools).not.toContain('list_issues') + expect(systemOf(call)).toContain('external_list_issues') + }) })