From 1cdc116aa3c94b3b9433c084dd7fef34139d9f35 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:01:15 +0200 Subject: [PATCH 1/9] feat(ai-harness): first-party plugins in @tanstack/ai-harness/plugins Add permissions, agents, question, todos, title, modelPicker, projectInstructions, fileCommands, compact, usage, and boundToolOutput under the ./plugins entry. Export the portable permission names from the package root again, so edge-safe plugins can read the rules. --- .changeset/ai-harness-plugins.md | 11 + docs/config.json | 14 +- docs/harness/agents.md | 229 +++++ docs/harness/permissions.md | 345 ++++++++ docs/harness/plugins.md | 2 + docs/harness/usage.md | 53 +- packages/ai-harness/package.json | 4 + packages/ai-harness/src/first-party/agents.ts | 380 +++++++++ .../src/first-party/bound-output.ts | 247 ++++++ packages/ai-harness/src/first-party/files.ts | 319 +++++++ packages/ai-harness/src/first-party/glob.ts | 14 + packages/ai-harness/src/first-party/index.ts | 31 + .../src/first-party/model-picker.ts | 68 ++ .../ai-harness/src/first-party/permissions.ts | 728 ++++++++++++++++ .../ai-harness/src/first-party/question.ts | 189 +++++ .../src/first-party/session-tools.ts | 272 ++++++ packages/ai-harness/src/first-party/title.ts | 90 ++ packages/ai-harness/src/first-party/todos.ts | 100 +++ .../src/first-party/workspace-hooks.ts | 35 + packages/ai-harness/src/index.ts | 15 + .../ai-harness/tests/agent-middleware.test.ts | 73 +- packages/ai-harness/tests/agents.test.ts | 446 ++++++++++ .../ai-harness/tests/bound-output.test.ts | 314 +++++++ .../ai-harness/tests/discover-tools.test.ts | 109 +++ packages/ai-harness/tests/edge-safety.test.ts | 4 +- packages/ai-harness/tests/first-party.test.ts | 178 ++++ .../ai-harness/tests/instructions.test.ts | 234 ++++++ packages/ai-harness/tests/permissions.test.ts | 780 ++++++++++++++++++ packages/ai-harness/tests/question.test.ts | 190 +++++ .../ai-harness/tests/session-plugins.test.ts | 190 ++++- packages/ai-harness/tests/title-usage.test.ts | 310 +++++++ packages/ai-harness/vite.config.ts | 2 +- .../fixtures/harness-first-party/basic.json | 107 +++ testing/e2e/src/routeTree.gen.ts | 21 + .../e2e/src/routes/api.harness-first-party.ts | 245 ++++++ testing/e2e/tests/harness-first-party.spec.ts | 125 +++ 36 files changed, 6466 insertions(+), 8 deletions(-) create mode 100644 .changeset/ai-harness-plugins.md create mode 100644 docs/harness/agents.md create mode 100644 docs/harness/permissions.md create mode 100644 packages/ai-harness/src/first-party/agents.ts create mode 100644 packages/ai-harness/src/first-party/bound-output.ts create mode 100644 packages/ai-harness/src/first-party/files.ts create mode 100644 packages/ai-harness/src/first-party/glob.ts create mode 100644 packages/ai-harness/src/first-party/index.ts create mode 100644 packages/ai-harness/src/first-party/model-picker.ts create mode 100644 packages/ai-harness/src/first-party/permissions.ts create mode 100644 packages/ai-harness/src/first-party/question.ts create mode 100644 packages/ai-harness/src/first-party/session-tools.ts create mode 100644 packages/ai-harness/src/first-party/title.ts create mode 100644 packages/ai-harness/src/first-party/todos.ts create mode 100644 packages/ai-harness/src/first-party/workspace-hooks.ts create mode 100644 packages/ai-harness/tests/agents.test.ts create mode 100644 packages/ai-harness/tests/bound-output.test.ts create mode 100644 packages/ai-harness/tests/discover-tools.test.ts create mode 100644 packages/ai-harness/tests/first-party.test.ts create mode 100644 packages/ai-harness/tests/instructions.test.ts create mode 100644 packages/ai-harness/tests/permissions.test.ts create mode 100644 packages/ai-harness/tests/question.test.ts create mode 100644 packages/ai-harness/tests/title-usage.test.ts create mode 100644 testing/e2e/fixtures/harness-first-party/basic.json create mode 100644 testing/e2e/src/routes/api.harness-first-party.ts create mode 100644 testing/e2e/tests/harness-first-party.spec.ts diff --git a/.changeset/ai-harness-plugins.md b/.changeset/ai-harness-plugins.md new file mode 100644 index 0000000000..f0812eb769 --- /dev/null +++ b/.changeset/ai-harness-plugins.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai-harness': minor +--- + +Add first-party plugins in `@tanstack/ai-harness/plugins` (a Node entry): + +- `permissions()`: rules for tools, paths, and commands, the `default`, `plan`, `acceptEdits`, and `bypass` modes, questions answered with once, always, or reject, saved "always" rules with `/permissions` and `/permissions forget `, and `listSavedPermissions` / `deleteSavedPermission` for server code. The rules apply in subagent runs too. +- `agents()`: agent profiles with their own prompt, tools, model, rules, and step limit, in code or in Markdown files, `/agent` to switch the primary agent, and the built-in `build`, `plan`, `general`, and `explore` agents. +- `question()`, `todos()`, `title()`, `modelPicker()`, `projectInstructions()`, `fileCommands()`, `compact()`, `usage()` (with your own prices for calls that have no provider cost), and `boundToolOutput()`. + +`PermissionRules`, `PermissionResources`, `PermissionDecisionCapability`, and `decidePermission` are also exported from the package root. They load no Node module, so an edge-safe plugin can read the permission rules. diff --git a/docs/config.json b/docs/config.json index 0d0de346cd..df3ace69e3 100644 --- a/docs/config.json +++ b/docs/config.json @@ -981,11 +981,23 @@ { "label": "Write a plugin", "to": "harness/plugins", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" }, { "label": "Track usage and cost", "to": "harness/usage", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" + }, + { + "label": "Ask before risky tool calls", + "to": "harness/permissions", + "addedAt": "2026-10-13" + }, + { + "label": "Switch between agent profiles", + "to": "harness/agents", "addedAt": "2026-10-13" }, { diff --git a/docs/harness/agents.md b/docs/harness/agents.md new file mode 100644 index 0000000000..8d918b519d --- /dev/null +++ b/docs/harness/agents.md @@ -0,0 +1,229 @@ +--- +title: Switch between agent profiles +id: harness-agents +order: 10 +description: "Give one harness several agents, such as a read-only planner and a builder. Each profile has its own prompt, tools, model, rules, and step limit, in code or in Markdown files." +keywords: + - tanstack ai + - harness + - agents + - agent profiles + - plan + - subagents +--- + +Your agent has two jobs. First it must plan a change without touching a file. Then it must build the change with every tool. The main model also needs a fast helper that searches the code. `agents()` gives the harness named agent profiles. Each profile has its own system prompt, tools, model, permission rules, and step limit. + +## 1. Add the plugin + +```ts group=harness-agents +import { readFile, writeFile } from 'node:fs/promises' +import { resolve } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { agents, permissions } from '@tanstack/ai-harness/plugins' +import { anthropicText } from '@tanstack/ai-anthropic' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' + +const root = process.cwd() +const sonnet = anthropicText('claude-sonnet-5-5') + +const readTool = toolDefinition({ + name: 'read_file', + description: 'Read a file in the project', + inputSchema: z.object({ path: z.string() }), +}).server(({ path }) => readFile(resolve(root, path), 'utf8')) + +const writeTool = toolDefinition({ + name: 'write_file', + description: 'Write a file in the project', + inputSchema: z.object({ path: z.string(), content: z.string() }), +}).server(async ({ path, content }) => { + await writeFile(resolve(root, path), content) + return `Wrote ${path}.` +}) + +const coder = defineHarness({ + name: 'acme/coder', + adapter: sonnet, + tools: [readTool, writeTool], + plugins: () => [ + permissions({ root, rules: [{ tool: 'write_file', decision: 'ask', kind: 'edit' }] }), + agents({ adapter: () => sonnet }), + ], +}) +``` + +`adapter` turns a model id into an adapter. Here every agent uses the same model. + +## 2. Plan, then build + +```ts group=harness-agents +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(coder, { threadId: 'dark-mode' }) + +await session.command('agent', 'plan') +const plan = await session.prompt('How do I add a dark mode?') +console.log(plan.text) + +await session.command('agent', 'build') +await session.prompt('Do step 1.') +``` + +1. The `plan` agent reads the code and gives numbered steps. It has no tool that changes a file. +2. The `build` agent changes the files, after your answer to each [permission question](./permissions). + +- `session.command('agent')` with no input shows the current agent and the list: `Agent: build. Agents: build, plan.` +- A switch applies at the next turn. + +The same switch is the `agent` setting. `session.describe().config` lists its options for a menu: + +```ts group=harness-agents +const setting = session.describe().config.find((entry) => entry.key === 'agent') +if (setting?.option.type === 'select') { + console.log(setting.value, setting.option.options) // 'build' ['build', 'plan'] +} + +await session.setConfig('agent', 'plan') +``` + +A client can change the setting only when the harness exposes it. Add `expose: { config: ['agent'] }` to `defineHarness`. + +## The built-in agents + +| Agent | Mode | Tools | What it does | +| --- | --- | --- | --- | +| `build` | `primary`, the default | Every tool | Reads code, changes files, and runs commands. | +| `plan` | `primary` | The read tools and `todo_write` | Reads the code and gives a plan in numbered steps. | +| `general` | `subagent` | Every tool | Does a task in many steps and reports back. | +| `explore` | `subagent` | The read tools | Searches the code and reports the paths that it finds. | + +The read tools are `read_file`, `list_files`, `grep`, `webfetch`, `websearch`, and `question`. An agent gets only the tools that the harness has. + +The mode of a profile says who starts it: + +- `primary`: it answers the user's turns. Pick it with the `agent` command or setting. +- `subagent`: the main model starts it as a child. +- `all`: both. + +The `plan` agent and the `plan` mode of [permissions](./permissions#pick-a-mode) are separate. The agent has only the read tools and `todo_write`. The mode denies edits and commands for every agent. To drop the built-in agents, set `builtIns: false`. + +## Add your own agent + +Give `agents()` a profile. This one reviews the changes with a stronger model: + +```ts group=harness-agents +import type { AnyTextAdapter } from '@tanstack/ai' +import type { AgentProfile } from '@tanstack/ai-harness/plugins' + +const review: AgentProfile = { + name: 'review', + description: 'Reviews the changes and lists the risks. Does not change files.', + mode: 'all', + model: 'claude-opus-5-5', + system: 'You review code. Read the diff, then list each risk with its file and line.', + tools: ['read_file', 'bash'], + permissions: [{ tool: 'bash', resource: 'git diff*', decision: 'allow' }], + steps: 20, +} + +const models: Record = { + 'claude-opus-5-5': anthropicText('claude-opus-5-5'), +} + +export const withReview = agents({ + adapter: (model) => models[model] ?? sonnet, + agents: [review], +}) +``` + +`name`, `description`, and `mode` are required. The other fields: + +| Field | What it does | +| --- | --- | +| `model` | A model id. `adapter` turns it into an adapter. Without it, a primary agent keeps the harness model, and a subagent uses the model of the main turn. | +| `system` | Text for the system prompt of each turn of this agent. | +| `tools` | The tools that this agent can call, by name. A `*` matches any text, for example `todo_*`. Without it, the agent gets every tool. | +| `permissions` | [Permission rules](./permissions#write-rules) that count while this agent is the primary agent. | +| `steps` | The step limit. See [Limit the steps](#limit-the-steps). | +| `hidden` | `true` leaves the agent out of the `agent` command and setting. The model can still start it as a subagent. | + +- A profile replaces an earlier profile with the same name. The order is the built-in agents, then `dirs`, then `agents`. So a profile named `plan` replaces the built-in `plan`. +- `default` names the primary agent of a new session. Without it, the first primary agent is the default, which is `build`. +- In the main turn, a model that the thread stored with `session.configure` wins over the profile `model`. See [What wins](./thread-settings#what-wins). +- A `tools` list also filters the subagent tools. To let the agent start subagents, add their tool names, for example `explore`. + +## Load agents from Markdown files + +Keep your agents next to your code as Markdown files. Save this file as `.agents/agents/review.md`: + +```md +--- +description: Reviews the changes and lists the risks. Does not change files. +mode: all +model: claude-opus-5-5 +tools: [read_file] +steps: 20 +--- + +You review code. Read the diff, then list each risk with its file and line. +``` + +Then give `agents()` the folder: + +```ts group=harness-agents +export const fromFiles = agents({ + adapter: (model) => models[model] ?? sonnet, + dirs: [`${root}/.agents/agents`], +}) +``` + +- The file name is the agent name, here `review`. +- The body is the system prompt. +- The frontmatter sets `description`, `mode`, `model`, `tools`, `steps`, and `hidden: true`. The default `mode` is `all`. +- An unknown `mode`, or a `steps` that is not a whole number, stops `host.open` with an error, for example `Agent file review.md: "mode" is not valid.` +- The plugin reads the folder when a session opens, and again at each `session.reload()`. A folder that it cannot read gives no agents. + +## Limit the steps + +An agent can call tools for a long time. `steps` sets the most model calls with tools in one run. With `steps: 20`: + +1. The agent makes up to 20 model calls with tools. +2. Call 21 gets `toolChoice: 'none'` and a note: answer now, and say what is left to do. +3. Then the run stops. + +The limit counts again in each turn of a primary agent and in each run of a subagent. + +Some providers do not obey `'none'`. On Amazon Bedrock, `'none'` after a tool call goes to the model as `auto`. See [Provider notes](../tools/tools#provider-notes). If the model calls a tool on the last call anyway: + +- The tool does not run. +- The model gets this tool result: `{ "error": "Step limit reached. Answer without calling tools." }` +- The run stops. That turn can end with no text answer. + +## Let the model start a subagent + +A profile with the mode `subagent` or `all` is also a subagent. The main model gets a tool for each one, with the name of the agent, for example `explore`. Ask the build agent: `use explore to find where the theme is set`. + +A subagent run gets these parts of its profile: + +- The tools of the last main turn, filtered by its `tools`. +- Its `system` prompt. +- Its `model`, or the model of the main turn. +- Its step limit. + +A subagent run stops after 50 model calls, or when an answer is cut at the output limit. The harness also limits the depth and the number of children. See [Stay within limits](./subagents#stay-within-limits). + +## Known limits + +- The `permissions` of a profile count only while it is the primary agent. When the profile runs as a subagent, its own rules do not apply. Put rules that must hold in every run in `permissions({ rules })`. See [Rules in subagents](./permissions#rules-in-subagents). +- A Markdown agent cannot set `permissions`. Put those rules in a profile in code. +- On Amazon Bedrock, the model can still call a tool on the last call. The tool does not run, but the turn can end with no text answer. +- A subagent that starts before the first turn of the session has no tools. It also needs its own `model`. + +## What you have now + +- A read-only `plan` agent, and a `build` agent to do the work. +- Your own agents in code or in Markdown files, each with its own prompt, tools, model, and rules. +- A step limit that ends a run with a text answer. +- Subagents that the main model starts to search and to research. diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md new file mode 100644 index 0000000000..15b2a9aa97 --- /dev/null +++ b/docs/harness/permissions.md @@ -0,0 +1,345 @@ +--- +title: Ask before risky tool calls +id: harness-permissions +order: 9 +description: "Make a harness agent ask before it changes a file or runs a command. Pick a mode, write rules for paths and commands, answer the questions, and keep the rules in subagents." +keywords: + - tanstack ai + - harness + - permissions + - approvals + - plan mode + - rules +--- + +Your agent has tools that change files or run commands. It must not do that without a yes from the user, and some files it must never touch. `permissions()` checks each tool call against rules. Each call runs, waits for the user, or does not run. + +## 1. Add the plugin + +```ts group=harness-permissions +import { writeFile } from 'node:fs/promises' +import { resolve } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { openaiText } from '@tanstack/ai-openai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' + +const root = process.cwd() + +const saveNote = toolDefinition({ + name: 'save_note', + description: 'Write a note file in the project', + inputSchema: z.object({ path: z.string(), text: z.string() }), +}).server(async ({ path, text }) => { + await writeFile(resolve(root, path), text) + return `Saved ${path}.` +}) + +export const notes = defineHarness({ + name: 'acme/notes', + adapter: openaiText('gpt-6.1-sol'), + tools: [saveNote], + plugins: () => [ + permissions({ + root, + rules: [{ tool: 'save_note', decision: 'ask', kind: 'edit' }], + }), + ], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(notes, { threadId: 'notes-1' }) +``` + +The rule makes each `save_note` call wait for the user. + +## 2. Answer the first question + +The question shows up as a `harness.question` event and in `session.snapshot().pendingQuestions`. Answer it with `session.answer`: + +```ts group=harness-permissions +const turn = session.prompt('Save a note with the time of the meeting.') + +for await (const { event } of session.events()) { + if (event.type === 'CUSTOM' && event.name === HARNESS_EVENTS.question) { + const [question] = session.snapshot().pendingQuestions + // For example: Allow save_note {"path":"meeting.md","text":"..."}? + if (question) await session.answer(question.questionId, 'once') + break + } +} + +await turn +``` + +The agent writes the file and continues. + +## Pick a mode + +The `mode` setting decides how much runs without a question: + +| Mode | What happens | +| --- | --- | +| `default` | The rules apply as written. A call that asks waits for the user. | +| `plan` | Read-only. Edits, commands, and calls that ask do not run. The model gets a note to plan and to change no files. | +| `acceptEdits` | Edits run without a question. Commands still ask. | +| `bypass` | Every call runs, also a call that a rule denies. | + +Switch the mode from server code: + +```ts group=harness-permissions +await session.setConfig('mode', 'plan') +// Or run the command. It returns 'Mode: plan.' +await session.command('mode', 'plan') +``` + +`session.command('mode')` with no input shows the current mode. In `acceptEdits` mode, an edit outside `root` or of a `.env` file still asks. + +### Keep the mode on the server + +A client that can set `mode` can set `bypass`. Then every call runs with no question. So by default, a client cannot change the mode: + +- A client `config` input for `mode` gets `{ status: 'rejected', reason: 'not_exposed' }`. The mode stays. +- A client `/mode` command gets the same receipt. +- Server code can still change the mode. + +A client can set only the config keys in `expose.config`. It can run only the commands in `expose.commands`. This harness lets a client pick the model, but not the mode: + +```ts group=harness-permissions +import { modelPicker } from '@tanstack/ai-harness/plugins' + +export const team = defineHarness({ + name: 'acme/team-notes', + adapter: openaiText('gpt-6.1-sol'), + tools: [saveNote], + plugins: () => [ + permissions({ root }), + modelPicker({ + choices: { sol: openaiText('gpt-6.1-sol'), pro: openaiText('gpt-6.1-sol-pro') }, + }), + ], + expose: { config: ['model'], commands: ['model'] }, +}) +``` + +Add `mode` to `expose.config` or `expose.commands` only for clients that you trust. Each of these clients can turn the questions off. + +## Write rules + +Add your own rules with `rules`: + +```ts group=harness-permissions +export const strict = permissions({ + root, + rules: [ + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + { tool: 'edit_file', resource: 'src/**', decision: 'allow' }, + { tool: 'bash', resource: 'pnpm test*', decision: 'allow' }, + { tool: 'deploy_*', decision: 'deny' }, + ], +}) +``` + +A rule has these fields: + +- `tool`: a tool name. A `*` at the end matches every tool that starts with that text. +- `decision`: `allow`, `ask`, or `deny`. +- `resource`: optional. A glob on the path or the command that the call touches. Without it, the rule matches every call of the tool. The tool must say what it touches, see [Rules for your own tool](#rules-for-your-own-tool). +- `kind`: optional. `read`, `edit`, or `execute`. The modes read it: `plan` denies `edit` and `execute`, and `acceptEdits` allows `edit`. + +The plugin reads the rules in this order, and the last rule that matches a call wins: + +1. The rules that other plugins add for their tools. +2. Your `rules`. +3. The saved `always` answers. A saved answer never wins over a `deny`. + +Two more rules apply to each tool: + +- A tool with no rule runs without a question. To ask for each tool that has no rule, set `default: 'ask'`. +- When the last rule of a tool is a `deny` with no `resource`, the model does not get the tool. + +### Paths + +- Inside `root`, a rule matches the path from `root`, for example `src/app.ts`. Outside `root`, it matches the full path, for example `/etc/hosts`. +- `**` matches any part of a path. `*` matches text in one name. `?` matches one character. +- On macOS and Windows, letter case does not count. So a `deny` for `secrets/**` also stops `Secrets/key`. +- A call with more than one path runs only when the rules allow every path. + +These paths ask, unless a rule with a `resource` allows them: + +- A path outside `root`. +- A `.env` file, for example `.env`, `.env.local`, or `prod.env`. +- A path that the plugin cannot resolve without the file system, for example `C:notes.txt` or `notes.txt:hidden`. + +This rule lets the agent read `.env.example` without a question: + +```ts group=harness-permissions +export const exampleEnv = permissions({ + root, + rules: [{ tool: 'read_file', resource: '.env.example', decision: 'allow' }], +}) +``` + +### Commands + +- The plugin splits a command at `&&`, `||`, `;`, `|`, and new lines. The rules must allow every part. So when only `ls*` is allowed, `ls && rm -rf build` asks. +- In a command rule, `*` matches any text, also `/`. So `git *` allows `git add src/app.ts`. +- A part with shell syntax that can hide a command never gets an `allow`. This includes `$(`, a backtick, `&`, `<`, `>`, and a heredoc. Such a part asks, or a matching `deny` stops it. + +## Answer a question + +A call that asks waits for one of three answers: + +| Answer | What happens | +| --- | --- | +| `once` | This call runs. | +| `always` | This call runs, and the plugin saves an `allow` rule for it. | +| `reject` | The call does not run. The model gets a tool error. | + +- Send the answer as a string, `'once'`, or as an object, `{ answer: 'once' }`. +- A reject can carry a `message` for the model: `{ answer: 'reject', message: 'Use pnpm, not npm.' }`. The model gets your message as the tool error. +- Any other answer rejects the call. This includes `y` and `yes`. + +### What `always` saves + +- One `allow` rule for each path or command part of the call, for example `{ tool: 'bash', resource: 'pnpm test', decision: 'allow' }`. The rule names that exact text, not a glob. +- One `allow` for the whole tool, when the tool does not say what it touches. +- With a metadata store (`stores.metadata`), the rules belong to `root`. Every session with that `root` uses them, for every user, also after a restart. +- Without a metadata store, the rules stay in memory until the session closes. +- No rule for a part with shell syntax, `*`, or `?`. The next call like it asks again. + +### See and forget saved answers + +1. Run `session.command('permissions')`. The command lists the saved rules of the project: + +```text +Saved rules: +1. bash pnpm test +2. read_file .env +``` + +2. Run `session.command('permissions', 'forget 2')` to delete rule 2. The open session asks again at the next `read_file` call for `.env`. + +A client can run the command only when `expose.commands` has `permissions`. + +Server code can also read and delete the saved rules without a session. Pass the stores of your host and the same `root` that `permissions()` gets: + +```ts group=harness-permissions +import { deleteSavedPermission, listSavedPermissions } from '@tanstack/ai-harness/plugins' + +// The persistence that you give to createHarnessHost. +const persistence = memoryPersistence() + +const saved = await listSavedPermissions(persistence.stores, root) +const [first] = saved +if (first) await deleteSavedPermission(persistence.stores, root, first) +``` + +- `deleteSavedPermission` finds the rule by its fields. An unknown rule does nothing. +- New sessions stop using a deleted rule. A session that is open now can keep it until it opens again. +- Without a metadata store, both functions find no rules. The command then lists the answers of the open session. + +## Rules in subagents + +The same check runs in every agent run of the session: + +- Each call in a child run gets the same rules and the same mode. This includes the agents that the model calls, the agents that plugins start, and their children. +- A question from a child waits in the same session. Answer it the same way. +- In `plan` mode, a child cannot change a file. Its call gets the tool error `This tool is not allowed in plan mode.` +- The `permissions` of an [agent profile](./agents) count only while that profile is the primary agent. They do not apply when the profile runs as a subagent. + +## Rules for your own tool + +A rule with a `resource` can match the calls of your tool only when the tool says what it touches. Add the tool to `PermissionResources`: + +```ts +import { copyFile } from 'node:fs/promises' +import { resolve } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { PermissionResources, PermissionRules, definePlugin } from '@tanstack/ai-harness' +import { z } from 'zod' + +const root = process.cwd() + +const copy = toolDefinition({ + name: 'copy_file', + description: 'Copy a file in the project', + inputSchema: z.object({ from: z.string(), to: z.string() }), +}).server(async ({ from, to }) => { + await copyFile(resolve(root, from), resolve(root, to)) + return `Copied ${from} to ${to}.` +}) + +/** A string field of the raw input. It throws when the field is missing. */ +function field(input: unknown, key: string) { + const value = + typeof input === 'object' && input !== null ? Reflect.get(input, key) : undefined + if (typeof value !== 'string') throw new Error(`No ${key}.`) + return value +} + +export const fileCopy = definePlugin({ + name: 'acme/file-copy', + setup: () => ({ + tools: [copy], + contribute: [ + PermissionRules.item({ tool: 'copy_file', decision: 'ask', kind: 'edit' }), + PermissionResources.item({ + copy_file: { paths: (input) => [field(input, 'from'), field(input, 'to')] }, + }), + ], + }), +}) +``` + +Now `{ tool: 'copy_file', resource: 'docs/**', decision: 'allow' }` copies files in `docs` without a question. + +- Each function gets the raw input of the model, before the schema check. If it throws, the call does not run, and the model gets `Cannot check the permissions of this call: No from.` +- `paths` are relative to `root`, or absolute. `commands` are command parts, one for each simple command. +- `kind: 'edit'` puts the tool under the modes. `plan` denies it, and `acceptEdits` runs it. +- If a tool does not say what it touches, a rule with a `resource` counts for every call. So a `deny` for one folder denies every call of the tool. + +## Ask for a decision from your plugin + +Your plugin must know what `permissions()` decides for a tool before the model calls it. Read the answer with `PermissionDecisionCapability`: + +```ts +import { PermissionDecisionCapability, definePlugin } from '@tanstack/ai-harness' + +export const planTools = definePlugin({ + name: 'acme/plan-tools', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => ({ + // Give the model only the tools that permissions() allows in plan mode. + prepareTools: ({ tools }) => { + const decide = ctx.getOptional(PermissionDecisionCapability) + if (!decide) return tools + return tools.filter((tool) => decide(tool.name, 'plan') === 'allow') + }, + }), +}) +``` + +- The answer is `'allow'`, `'ask'`, or `'deny'`. It uses all rules in their order, your `rules` last, and your `default`. +- 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. +- `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 + +- A rule that allows a tool with no `kind` also lets it run in `plan` mode. `plan` denies only edits, commands, and calls that ask. +- A client that can send inputs to the thread can answer the questions, also with `always`. Control who can open a thread. +- `permissions()` does not follow links. Your own tool can check the real path of a link with `PermissionDecisionCapability`. + +## What you have now + +- A question before each risky tool call. +- Modes to plan, to edit fast, or to run everything. +- A mode that a client cannot change, unless you expose it. +- Rules for folders, files, and commands, and saved answers for each project that you can list and forget. +- The same rules in every agent run. + +Next: give the agent a planner and a builder with [agent profiles](./agents). diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index 05a2143851..2d8398e213 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -395,3 +395,5 @@ If a `setup` throws, the session keeps working with no plugins. The promise reje - A picker that chooses the model of each turn from its message, context, and sender. - Plugins that share services, lists, and events without knowing each other. - Commands, agents, and whole plugins that change while the session runs. + +Next: add the first-party plugins from `@tanstack/ai-harness/plugins`. [Ask before risky tool calls](./permissions) and [Switch between agent profiles](./agents) are two of them. diff --git a/docs/harness/usage.md b/docs/harness/usage.md index 2acdbbbd16..3a936f03d9 100644 --- a/docs/harness/usage.md +++ b/docs/harness/usage.md @@ -2,7 +2,7 @@ title: Track usage and cost id: harness-usage order: 7 -description: "Read the tokens and the cost of a thread, in total, by model, and by sender. Show a live total, and stop a person who is over budget." +description: "Read the tokens and the cost of a thread, in total, by model, and by sender. Show a live total, stop a person who is over budget, and add /usage with your own prices." keywords: - tanstack ai - harness @@ -111,6 +111,56 @@ for (const entry of page.entries) { A background agent can end after the turn. Then the index entry counts its tokens at the end of the next turn. +## Show usage and price the calls + +Many providers report no cost, so `cost` stays empty. The `usage()` plugin adds a `/usage` command, and can price those calls with your own price list: + +```ts group=harness-usage +import { usage } from '@tanstack/ai-harness/plugins' + +// USD per 1M tokens, by model id. +const prices: Record = { + 'gpt-6.1-sol': { input: 2, output: 8, cacheRead: 0.2 }, +} + +export const priced = defineHarness({ + name: 'acme/priced', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + usage({ + model: (id) => { + const cost = prices[id] + return cost === undefined ? undefined : { cost } + }, + }), + ], +}) +``` + +`session.command('usage')` reads `session.usage()` and shows the total. When the thread used more than one model, it also shows a line for each model: + +```text +3 model calls, 41200 input tokens (30000 cache read, 0 cache write), 2100 output tokens, 43300 total. Cost: $0.0432. +openai/gpt-6.1-sol: 2 model calls, 40000 input tokens (30000 cache read, 0 cache write), 2000 output tokens, 42000 total. +openai/gpt-6-luna: 1 model calls, 1200 input tokens, 100 output tokens, 1300 total. +``` + +How the plugin finds the cost: + +- A call with a cost from its provider keeps that cost, as it is. +- `model` prices only the calls with no provider cost. It gets the model id without the provider, like `gpt-6.1-sol`. It returns `{ cost }` with prices in USD per 1M tokens, or `undefined`. +- Without `model`, `/usage` shows `Cost:` only when a provider reported a cost. +- A call with no price adds a note, for example `(cost unknown for 1 call)`. + +The plugin also writes the cost to `usage.cost` of the session index entry, whenever it has one. `session.usage()` keeps only the costs that the providers report. + +The plugin state, `snapshot().plugins['tanstack/usage']`, has a copy of `session.usage().total` and `contextTokens`: the input tokens of the latest call of the lead model. Compare it with the context window of the model to show how full the context is. + +Limits of the plugin: + +- For one model, when some calls have a provider cost and others do not, `model` does not price the others. +- On a host with a log, the cost in the index entry and the state copy can miss the last call of a run. They get it at the end of the next run. + ## What you have now - The tokens and the cost of a thread, in total, by model, and by sender. @@ -118,3 +168,4 @@ A background agent can end after the turn. Then the index entry counts its token - A live total from the `harness.usage` event. - A budget check before each prompt. - The token counts of each session in the session index. +- `/usage` with your own prices, and a cost for each session in the session index. diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 5d45fa9b9d..1272c358e4 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -29,6 +29,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./plugins": { + "types": "./dist/esm/first-party/index.d.ts", + "import": "./dist/esm/first-party/index.js" } }, "scripts": { diff --git a/packages/ai-harness/src/first-party/agents.ts b/packages/ai-harness/src/first-party/agents.ts new file mode 100644 index 0000000000..7e4a2926e9 --- /dev/null +++ b/packages/ai-harness/src/first-party/agents.ts @@ -0,0 +1,380 @@ +import { + combineStrategies, + defineAgent, + maxIterations, + untilFinishReason, +} from '@tanstack/ai' +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { definePlugin } from '../plugins' +import { readMarkdownFiles } from './files' +import { globToRegExp } from './glob' +import { PermissionRules } from './permissions' +import type { + AnyTextAdapter, + AnyTool, + ChatMiddleware, + ChatMiddlewareConfig, + KeyedAdapter, +} from '@tanstack/ai' +import type { FrontmatterValue } from './files' +import type { PermissionRule } from './permissions' + +/** A named agent: its prompt, tools, model, rules, and step limit. */ +export interface AgentProfile { + name: string + /** What the agent does. The model reads it to pick a subagent. */ + description: string + /** + * - `primary`: it answers the user's turns. Pick it with `/agent`. + * - `subagent`: the main model starts it with a tool named after it. + * - `all`: both. + */ + mode: 'primary' | 'subagent' | 'all' + /** + * A model id. `agents({ adapter })` turns it into an adapter. Without it, + * a primary agent keeps the harness model, and a subagent uses the model + * of the main turn. + */ + model?: string + /** System prompt text added to each turn of this agent. */ + system?: string + /** + * The tools this agent can call, by name. A `*` matches any text, for + * example `todo_*`. Without it, the agent gets every tool. + */ + tools?: ReadonlyArray + /** + * Permission rules that count while this agent is the primary agent. + * `permissions()` reads them. They do not apply in a subagent run. + */ + permissions?: ReadonlyArray + /** + * The most model calls with tools in one run. The next call gets no tool + * calls and a note to answer now, and then the run stops. + */ + steps?: number + /** + * Leave it out of the `/agent` list and the `agent` setting. The model + * can still start it as a subagent. + */ + hidden?: boolean +} + +/** The tools that only read: files, search, the web, and questions. */ +const READ_TOOLS = [ + 'read_file', + 'list_files', + 'grep', + 'webfetch', + 'websearch', + 'question', +] + +/** The built-in agents: `build`, `plan`, `general`, and `explore`. */ +export const builtInAgents: ReadonlyArray = [ + { + name: 'build', + description: 'Does the work: reads code, changes files, and runs commands.', + mode: 'primary', + system: + 'You are the build agent. Do the task the user gives you. Read the code you need, change files, and run commands. Check your work before you say it is done.', + }, + { + name: 'plan', + description: 'Reads the code and makes a plan. Does not change files.', + mode: 'primary', + tools: [...READ_TOOLS, 'todo_write'], + system: + 'You are the plan agent. You can only read and search. Do not change files or run commands. Study the code, then give the user a clear plan in numbered steps.', + }, + { + name: 'general', + description: + 'Does a research or work task in many steps and reports back. Give it the full task.', + mode: 'subagent', + system: + 'You are a general agent. Do the task you get in as many steps as it needs. Use your tools to search, read, and act. At the end, reply with a short report of what you found or did.', + }, + { + name: 'explore', + description: + 'Searches the code fast and reports what it finds. Read-only. Give it a question.', + mode: 'subagent', + tools: READ_TOOLS, + system: + 'You are a fast search agent. Find the files and code that answer the question you get. Only read and search. Reply with the paths and a short note for each.', + }, +] + +const MODES = ['primary', 'subagent', 'all'] as const + +const MAX_STEPS_PROMPT = + 'You reached the step limit. Do not call more tools. Answer now with what you have, and say what is left to do.' + +const MAX_STEPS_RESULT = 'Step limit reached. Answer without calling tools.' + +/** The tools whose name matches one of `patterns`. Without patterns, all of them. */ +function allowedTools( + tools: ReadonlyArray, + patterns: ReadonlyArray | undefined, +) { + if (patterns === undefined) return [...tools] + const matchers = patterns.map((pattern) => globToRegExp(pattern)) + return tools.filter((tool) => + matchers.some((matcher) => matcher.test(tool.name)), + ) +} + +/** + * Apply a profile to each model call of a chat run: keep only its tools, + * and at its step limit, make one last call with no tool calls. + */ +function profileMiddleware(current: () => AgentProfile) { + const isLastStep = (iteration: number) => { + const { steps } = current() + return steps !== undefined && iteration >= steps + } + const middleware: ChatMiddleware = { + name: 'tanstack/agents', + onConfig: (run, config) => { + if (run.phase !== 'init' && run.phase !== 'beforeModel') return + const { tools } = current() + const update: Partial = {} + if (tools !== undefined) update.tools = allowedTools(config.tools, tools) + if (run.phase === 'beforeModel' && isLastStep(run.iteration)) { + update.toolChoice = 'none' + // The config keeps the prompts, so a retried call has it already. + if (!config.systemPrompts.includes(MAX_STEPS_PROMPT)) { + update.systemPrompts = [...config.systemPrompts, MAX_STEPS_PROMPT] + } + } + return update + }, + // Some providers do not follow `toolChoice: 'none'`. Bedrock Converse + // sends it as `auto` after tool calls. So refuse a call of the last step + // here. In the tool phase, `iteration` is still the one of that call. + onBeforeToolCall: (run) => + isLastStep(run.iteration) + ? { type: 'skip', result: { error: MAX_STEPS_RESULT } } + : undefined, + // After the last call, stop. The harness loop limit still applies. + onShouldContinue: (_run, state) => { + const { steps } = current() + return steps === undefined || state.iterationCount <= steps + }, + } + return middleware +} + +/** A profile from a Markdown file. The body is its system prompt. */ +function profileOfFile(file: { + name: string + fields: Record + body: string +}) { + const text = (key: string) => { + const value = file.fields[key] + return typeof value === 'string' && value !== '' ? value : undefined + } + const fail = (key: string) => { + throw new Error(`Agent file ${file.name}.md: "${key}" is not valid.`) + } + const modeText = text('mode') + const mode = + modeText === undefined + ? 'all' + : (MODES.find((candidate) => candidate === modeText) ?? fail('mode')) + const stepsText = text('steps') + const steps = stepsText === undefined ? undefined : Number(stepsText) + if (steps !== undefined && !(Number.isInteger(steps) && steps >= 0)) { + fail('steps') + } + const model = text('model') + const { tools } = file.fields + const system = file.body.trim() + // ponytail: no `permissions` from files, the parser reads no nested values + const profile: AgentProfile = { + name: file.name, + description: text('description') ?? `The ${file.name} agent`, + mode, + ...(model !== undefined && { model }), + ...(system !== '' && { system }), + // `tools: a, b` is a list too, not one tool named `a, b`. + ...(tools !== undefined && { + tools: + typeof tools === 'string' + ? tools + .split(',') + .map((name) => name.trim()) + .filter((name) => name !== '') + : tools, + }), + ...(steps !== undefined && { steps }), + ...(text('hidden') === 'true' && { hidden: true }), + } + return profile +} + +/** + * Named agents, each with its own system prompt, tools, model, permission + * rules, and step limit. + * + * - A primary agent answers the user's turns. Switch it with + * `/agent ` or the `agent` setting. It applies at the next turn. + * - A subagent is a child that the main model starts with a tool named + * after the agent. + * + * The built-in agents are `build` (primary, the default, every tool), + * `plan` (primary, read-only tools), `general` (subagent, every tool), and + * `explore` (subagent, read-only tools). See {@link builtInAgents}. + * + * @param options.adapter - Turns a model id into an adapter: the `model` of + * a profile, or for a subagent without one, the model id of the main turn. + * A `keyedAdapter(...)` is built with the user's key. + * @param options.agents - More profiles. A profile replaces an earlier one + * with the same name: built-ins first, then `dirs`, then `agents`. + * @param options.dirs - Folders of `*.md` agent files. The file name is the + * agent name and the body is its system prompt. The frontmatter sets + * `description`, `mode` (default `all`), `model`, `tools` (a list), + * `steps`, and `hidden`. + * @param options.builtIns - `false` drops the built-in agents. + * @param options.default - The primary agent of a new session. Default: the + * first primary agent, `build` with the built-ins. + * + * @example + * ```ts + * const models: Record = { + * 'claude-sonnet-4-5': anthropicText('claude-sonnet-4-5'), + * } + * agents({ + * adapter: (model) => models[model] ?? openaiText('gpt-5.5'), + * dirs: ['.agents/agents'], + * }) + * ``` + */ +export function agents(options: { + adapter: (model: string) => AnyTextAdapter | KeyedAdapter + agents?: ReadonlyArray + dirs?: ReadonlyArray + builtIns?: boolean + default?: string +}) { + return definePlugin({ + name: 'tanstack/agents', + setup: async (ctx) => { + const dirs = options.dirs ?? [] + const files = (await Promise.all(dirs.map(readMarkdownFiles))).flat() + const all = [ + ...(options.builtIns === false ? [] : builtInAgents), + ...files.map(profileOfFile), + ...(options.agents ?? []), + ] + // A later profile replaces an earlier one and keeps its place. + const profiles = [ + ...new Map(all.map((profile) => [profile.name, profile])).values(), + ] + const visible = profiles.filter( + (profile) => profile.mode !== 'subagent' && !profile.hidden, + ) + const names = visible.map((profile) => profile.name) + const first = visible.find( + (profile) => profile.name === (options.default ?? names[0]), + ) + if (!first) { + throw new Error( + 'agents(): `default` must name a primary agent that is not hidden.', + ) + } + const active = () => { + const name = ctx.config.get('agent') + return visible.find((profile) => profile.name === name) ?? first + } + // The tools and the model id of the last main turn. Subagents use them. + // ponytail: a subagent started before any turn has no tools and needs a `model` + let lead: { tools: ReadonlyArray; model: string } | undefined + + const subagentOf = (profile: AgentProfile) => + defineAgent({ + name: profile.name, + description: profile.description, + run: async (child) => { + const model = profile.model ?? lead?.model + if (model === undefined) { + throw new Error( + `Agent ${profile.name} has no model. Give it a model, or start it from a turn.`, + ) + } + return child.chat({ + adapter: await child.keys.adapter(options.adapter(model)), + tools: allowedTools(lead?.tools ?? [], profile.tools), + systemPrompts: + profile.system === undefined ? [] : [profile.system], + middleware: [profileMiddleware(() => profile)], + // The harness main loop has the same limits. chat() alone stops + // at 5 calls. An answer cut at the output limit ends the run. + agentLoopStrategy: combineStrategies([ + maxIterations(50), + untilFinishReason(['length']), + ]), + }) + }, + }) + + // A profile's rules count only while it is the primary agent. At other + // times the getter gives a tool name that matches no tool. + // ponytail: subagent runs get no profile rules, that needs the run in PermissionRules + const rules = profiles.flatMap((profile) => + (profile.permissions ?? []).map((rule) => + PermissionRules.item({ + ...rule, + get tool() { + return active().name === profile.name ? rule.tool : '' + }, + }), + ), + ) + + return { + config: { + agent: configOption.select({ + options: names, + default: first.name, + description: 'The primary agent', + }), + }, + commands: { + agent: defineCommand({ + description: `Show or switch the primary agent (${names.join(', ')})`, + run: async (input: unknown) => { + const name = typeof input === 'string' ? input.trim() : '' + const list = `Agents: ${names.join(', ')}.` + if (name === '') return `Agent: ${active().name}. ${list}` + const receipt = await ctx.session.setConfig('agent', name) + return receipt.status === 'rejected' + ? `Unknown agent "${name}". ${list}` + : `Agent: ${name}. It applies at the next turn.` + }, + }), + }, + prompts: [ + { id: 'tanstack/agents:system', text: () => active().system ?? '' }, + ], + adapter: () => { + const { model } = active() + return model === undefined ? undefined : options.adapter(model) + }, + prepareTools: ({ tools, model }) => { + lead = { tools, model } + return tools + }, + // In onConfig, so the subagent tools that chat() adds are filtered too. + middleware: [profileMiddleware(active)], + subagents: profiles + .filter((profile) => profile.mode !== 'primary') + .map(subagentOf), + contribute: rules, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/bound-output.ts b/packages/ai-harness/src/first-party/bound-output.ts new file mode 100644 index 0000000000..0cd811144b --- /dev/null +++ b/packages/ai-harness/src/first-party/bound-output.ts @@ -0,0 +1,247 @@ +import { mkdir, readdir, rm, writeFile } from 'node:fs/promises' +import { join } from 'node:path' +import { normalizeToolResult } from '@tanstack/ai' +import { definePlugin } from '../plugins' +import { isRecord } from '../utils' +import type { AfterToolCallInfo, AnyChatMiddleware } from '@tanstack/ai' + +const HOUR = 60 * 60 * 1000 +const DAY = 24 * HOUR +/** Saved outputs start with this. The cleanup removes only these files. */ +const SAVED_PREFIX = 'tool-output-' +/** Core marks subagent tools with this global symbol. */ +const SUBAGENT_TOOL = Symbol.for('tanstack.ai.subagentTool') +const encoder = new TextEncoder() +// ignoreBOM keeps a leading byte order mark in the kept text. +const decoder = new TextDecoder('utf-8', { ignoreBOM: true }) + +/** The number of lines in `text`. A final newline does not start a new line. */ +function countLines(text: string) { + return text.split('\n').length - (text.endsWith('\n') ? 1 : 0) +} + +/** True for a UTF-8 byte inside a character, not at its start. */ +function isContinuationByte(byte: number | undefined) { + return byte !== undefined && (byte & 0xc0) === 0x80 +} + +/** The start of `text` inside both limits. `bytes` is `text` as UTF-8. */ +function firstPart( + text: string, + bytes: Uint8Array, + limits: { maxLines: number; maxBytes: number }, +) { + let end = limits.maxBytes + // Step back to the start of a character, so the cut never splits one. + while (isContinuationByte(bytes[end])) end-- + const firstBytes = decoder.decode(bytes.subarray(0, end)) + const firstLines = text.split('\n', limits.maxLines).join('\n') + return firstLines.length < firstBytes.length ? firstLines : firstBytes +} + +/** The end of `text` inside both limits. `bytes` is `text` as UTF-8. */ +function lastPart( + text: string, + bytes: Uint8Array, + limits: { maxLines: number; maxBytes: number }, +) { + let start = Math.max(0, bytes.length - limits.maxBytes) + // Step forward to the start of a character, so the cut never splits one. + while (isContinuationByte(bytes[start])) start++ + const lastBytes = decoder.decode(bytes.subarray(start)) + const lines = text.split('\n') + // A final newline ends the last line. It does not start a new one. + const finalNewline = text.endsWith('\n') ? 1 : 0 + const from = Math.max(0, lines.length - limits.maxLines - finalNewline) + const lastLines = lines.slice(from).join('\n') + return lastLines.length < lastBytes.length ? lastLines : lastBytes +} + +/** + * Keep the part of `text` that fits in `maxLines` lines and `maxBytes` UTF-8 + * bytes: the first lines (`keep: 'head'`, the default) or the last lines + * (`keep: 'tail'`). When something is cut, `truncated` is `true` and a note + * tells how much is shown: after the text for `head`, before it for `tail`. + * A cut never splits a character. Text inside both limits comes back + * unchanged. + * + * @example + * ```ts + * const { text, truncated } = boundText(output, { + * maxLines: 2000, + * maxBytes: 50 * 1024, + * keep: 'tail', + * }) + * ``` + */ +export function boundText( + text: string, + limits: { maxLines: number; maxBytes: number; keep?: 'head' | 'tail' }, +) { + const lines = countLines(text) + const bytes = encoder.encode(text) + const isOverLimit = lines > limits.maxLines || bytes.length > limits.maxBytes + if (!isOverLimit) return { text, truncated: false } + const isTail = limits.keep === 'tail' + const kept = isTail + ? lastPart(text, bytes, limits) + : firstPart(text, bytes, limits) + const shown = `${countLines(kept)} of ${lines} lines, ${encoder.encode(kept).length} of ${bytes.length} bytes` + return isTail + ? { + text: `[Output cut. Showing the last ${shown}.]\n\n${kept}`, + truncated: true, + } + : { + text: `${kept}\n\n[Output cut. Showing the first ${shown}.]`, + truncated: true, + } +} + +/** + * The time a saved output was written, from its file name, so no stat call + * is needed. `NaN` for a file this plugin did not save. + */ +function savedAt(name: string) { + if (!name.startsWith(SAVED_PREFIX)) return Number.NaN + return Number(name.slice(SAVED_PREFIX.length).split('-')[0]) +} + +/** Remove the saved outputs in `dir` that are older than `maxAgeMs`. */ +async function removeOldOutputs(dir: string, maxAgeMs: number) { + const cutoff = Date.now() - maxAgeMs + const names = await readdir(dir) + const old = names.filter((name) => savedAt(name) < cutoff) + await Promise.all(old.map((name) => rm(join(dir, name), { force: true }))) +} + +/** + * Save the full output of one tool call in `dir`. Gives back the note for + * the model: the file path, or why the save failed. + */ +async function saveFullOutput(dir: string, toolCallId: string, text: string) { + const safeId = toolCallId.replace(/[^\w-]/g, '_') + const path = join(dir, `${SAVED_PREFIX}${Date.now()}-${safeId}.txt`) + try { + await mkdir(dir, { recursive: true }) + await writeFile(path, text) + return `[Full output saved to ${path}.]` + } catch (error) { + // The cut result is still useful, so a failed save does not fail the call. + return `[The full output was not saved: ${String(error)}]` + } +} + +/** + * The result the model sees now. A failed call shows `{ error: message }`, + * unless an earlier middleware replaced it. `undefined` when there is + * nothing to cut. + */ +function currentResult(info: AfterToolCallInfo) { + if (info.ok || info.result !== undefined) return info.result + // A failed subagent call also shows its subagentRunId, which `info` does + // not carry. Leave it alone, so the model can still continue the child. + const isSubagent = info.tool !== undefined && SUBAGENT_TOOL in info.tool + if (isSubagent || !(info.error instanceof Error)) return undefined + return { error: info.error.message } +} + +/** + * The long part of a result, and a function that puts the cut text back. + * An error result keeps its other fields. A subagent result keeps its + * `subagentRunId`, so the model can still continue the child. + */ +function longPart(result: unknown, ok: boolean) { + if (isRecord(result)) { + if (!ok && 'error' in result) { + return { + long: result.error, + put: (cut: string) => ({ ...result, error: cut }), + } + } + if (typeof result.subagentRunId === 'string' && 'result' in result) { + return { + long: result.result, + put: (cut: string) => ({ ...result, result: cut }), + } + } + } + return { long: result, put: (cut: string) => cut } +} + +/** + * Keep every tool result small enough for the model. A result over + * `maxLines` lines (default 2000) or `maxBytes` UTF-8 bytes (default 50 KiB) + * is cut to its first lines, with a note. + * + * - A string result is cut as is. Any other value is cut as JSON text. + * - A failed call has its error text cut. It stays an error. + * - A subagent result has the child's answer cut. Its `subagentRunId` stays. + * - Content parts (images, audio, and so on) are not changed. + * + * It runs for every tool, MCP tools too, in the lead turn and in every agent + * run. With `dir`, the full output is saved to a file in `dir` and the note + * gives its path. Saved files older than `retentionDays` (default 7) are + * removed, at most once an hour. Node only. + * + * @example + * ```ts + * defineHarness({ + * name: 'acme/coder', + * adapter, + * plugins: () => [boundToolOutput({ dir: '.agent/tool-output' })], + * }) + * ``` + */ +export function boundToolOutput( + options: { + maxLines?: number + maxBytes?: number + dir?: string + retentionDays?: number + } = {}, +) { + const { + maxLines = 2000, + maxBytes = 50 * 1024, + dir, + retentionDays = 7, + } = options + return definePlugin({ + name: 'tanstack/bound-tool-output', + setup: () => { + let lastCleanup = 0 + const cleanUp = async (folder: string) => { + const now = Date.now() + if (now - lastCleanup < HOUR) return + lastCleanup = now + // ponytail: best effort. A failed cleanup tries again an hour later. + await removeOldOutputs(folder, retentionDays * DAY).catch(() => {}) + } + const bound = { + name: 'tanstack/bound-tool-output', + onAfterToolCall: async (_ctx, info) => { + const result = currentResult(info) + if (result === undefined) return + const { long, put } = longPart(result, info.ok) + const full = normalizeToolResult(long) + // Content parts (images, audio, and so on) stay as they are. + if (typeof full !== 'string') return + const bounded = boundText(full, { maxLines, maxBytes }) + if (!bounded.truncated) return + if (dir === undefined) { + return { type: 'replaceResult', result: put(bounded.text) } + } + const saved = await saveFullOutput(dir, info.toolCallId, full) + await cleanUp(dir) + return { + type: 'replaceResult', + result: put(`${bounded.text}\n${saved}`), + } + }, + } satisfies AnyChatMiddleware + // The same middleware for the lead turn and every agent run. + return { middleware: [bound], agentMiddleware: [bound] } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/files.ts b/packages/ai-harness/src/first-party/files.ts new file mode 100644 index 0000000000..789d464e06 --- /dev/null +++ b/packages/ai-harness/src/first-party/files.ts @@ -0,0 +1,319 @@ +import { execFile } from 'node:child_process' +import { existsSync, readFileSync, statSync } from 'node:fs' +import { readFile, readdir } from 'node:fs/promises' +import { homedir } from 'node:os' +import { + basename, + dirname, + extname, + isAbsolute, + join, + relative, + resolve, + sep, +} from 'node:path' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import { WorkspaceHooks } from './workspace-hooks' +import type { AnyCommand } from '../commands' + +/** Prompt slots kept free for instruction files that change. */ +const SPARE_SLOTS = 8 + +/** The mtime of the file at `path`, or `undefined` when there is none. */ +function mtimeOf(path: string) { + try { + return statSync(path).mtimeMs + } catch { + return undefined + } +} + +/** The trimmed text of the file at `path`, or `''` when there is none. */ +function textOf(path: string) { + try { + return readFileSync(path, 'utf8').trim() + } catch { + return '' + } +} + +/** + * The folders from the repo root (the nearest folder with `.git`) down to + * `root`, outermost first. Only `root` when no folder has `.git`. + */ +function repoFolders(root: string) { + const folders = [root] + let dir = root + while (!existsSync(join(dir, '.git'))) { + const parent = dirname(dir) + if (parent === dir) return { folders: [root], isRepo: false } + dir = parent + folders.unshift(dir) + } + return { folders, isRepo: true } +} + +/** The date, platform, working folder, and git facts, as a prompt. */ +async function environment(root: string, isRepo: boolean) { + const now = new Date() + // toISOString is UTC. Move by the offset to get the local date. + const date = new Date(now.getTime() - now.getTimezoneOffset() * 60_000) + .toISOString() + .slice(0, 10) + const lines = [ + `Date: ${date}`, + `Platform: ${process.platform}`, + `Working folder: ${root}`, + `Git repository: ${isRepo ? 'yes' : 'no'}`, + ] + if (isRepo) { + const branch = await new Promise((done) => + execFile( + 'git', + ['rev-parse', '--abbrev-ref', 'HEAD'], + { cwd: root, timeout: 5000 }, + (error, stdout) => done(error ? '' : stdout.trim()), + ), + ) + if (branch !== '') lines.push(`Git branch: ${branch}`) + } + return `Environment:\n${lines.map((line) => `- ${line}`).join('\n')}` +} + +/** The note for an instruction file whose text changed from `before`. */ +function changeNote(title: string, before: string, text: string) { + if (text === '') return `${title} were removed.` + if (before === '') return `${title}:\n${text}` + return `${title} changed. Follow this text instead:\n${text}` +} + +/** + * Add instruction files (AGENTS.md, CLAUDE.md) and an environment block to + * the system prompt. + * + * - `global`: files that come first, for example `~/.config/AGENTS.md`. `~` + * is the home folder. + * - Then `files` from each folder between the repo root (the folder with + * `.git`) and `root`, outermost first. With no repo root, only `root`. + * - `env` (default `true`): the date, platform, working folder, and git + * repository and branch, read when the session opens. + * + * At the start of each turn, the plugin checks the mtime of each file. A note + * with the new text of a changed file goes after this plugin's other prompts. + * Put this plugin last: when no prompt comes after its prompts, an adapter + * with mid-conversation changes adds the note to the conversation, and the + * prompt cache holds. Else the change is not additive, and the prompt cache + * starts again. + * + * When `read_file` reads a file below `root`, the instruction files in the + * folders between `root` and that file are added to the result, once per + * file in the session. + * + * @example + * ```ts + * projectInstructions({ root: process.cwd(), global: ['~/.config/AGENTS.md'] }) + * ``` + */ +export function projectInstructions(options: { + root: string + files?: ReadonlyArray + global?: ReadonlyArray + env?: boolean +}) { + const { files: names = ['AGENTS.md', 'CLAUDE.md'], env = true } = options + return definePlugin({ + name: 'tanstack/project-instructions', + setup: async () => { + const root = resolve(options.root) + const { folders, isRepo } = repoFolders(root) + const shownPath = (path: string) => + relative(root, path).split(sep).join('/') + const globals = (options.global ?? []).map((path) => ({ + path: resolve(path.replace(/^~(?=$|[\\/])/, homedir())), + title: `Global instructions from ${path}`, + })) + const projects = folders.flatMap((dir) => + names.map((name) => ({ + path: join(dir, name), + title: `Project instructions from ${shownPath(join(dir, name))}`, + })), + ) + const files = [...globals, ...projects].map((file) => ({ + ...file, + mtimeMs: mtimeOf(file.path), + text: textOf(file.path), + })) + const envBlock = env ? [await environment(root, isRepo)] : [] + /** The prompt with each file as it is now. */ + const current = () => [ + ...files + .filter((file) => file.text !== '') + .map((file) => `${file.title}:\n${file.text}`), + ...envBlock, + ] + const sections = current() + const slots = files.length + envBlock.length + SPARE_SLOTS + + /** Add a note for each file that changed since the last check. */ + const refresh = () => { + const notes: Array = [] + for (const file of files) { + const mtimeMs = mtimeOf(file.path) + if (mtimeMs === file.mtimeMs) continue + const before = file.text + file.mtimeMs = mtimeMs + file.text = textOf(file.path) + if (file.text !== before) + notes.push(changeNote(file.title, before, file.text)) + } + // ponytail: with no spare slot left, the prompt is built again. That + // one change is not added to the conversation. + if (sections.length + notes.length > slots) { + sections.splice(0, sections.length, ...current()) + } else { + sections.push(...notes) + } + } + + /** Nested instruction files the model already got. */ + const given = new Set() + const afterRead = async (path: string) => { + const inside = relative(root, path) + const isBelowRoot = + inside !== '' && !inside.startsWith('..') && !isAbsolute(inside) + if (!isBelowRoot) return undefined + const notes: Array = [] + // The folders below `root`, down to the folder of the file. + const parts = inside.split(sep).slice(0, -1) + let dir = root + for (const part of parts) { + dir = join(dir, part) + for (const name of names) { + const file = join(dir, name) + if (given.has(file)) continue + const text = textOf(file) + if (text === '') continue + given.add(file) + // The model just read this file, so it has the text. + if (file !== path) { + notes.push(`Instructions from ${shownPath(file)}:\n${text}`) + } + } + } + return notes.length > 0 ? notes.join('\n\n') : undefined + } + + return { + // One prompt per slot. Empty slots are left out. The harness reads + // the prompts in order each turn, so slot 0 checks the files first. + prompts: Array.from({ length: slots }, (_, index) => () => { + if (index === 0) refresh() + return sections[index] ?? '' + }), + contribute: [WorkspaceHooks.item({ afterRead })], + } + }, + }) +} + +/** A frontmatter value: text, or a list from `[a, b]` or `- a` lines. */ +export type FrontmatterValue = string | Array + +/** The text without the quotes around it: `"a"` and `'a'` become `a`. */ +function unquote(value: string) { + return value.trim().replace(/^(["'])(.*)\1$/, '$2') +} + +/** + * Split `---` frontmatter from a Markdown file. It reads `key: value` + * lines and lists: `key: [a, b]`, or `key:` with `- a` lines under it. + * Nested values are not read. + */ +export function splitFrontmatter(text: string) { + const fields: Record = {} + const match = /^---\r?\n([\s\S]*?)\r?\n---\r?\n?/.exec(text) + if (!match) return { fields, body: text } + // The list that `- a` lines add to: the last `key:` with no value. + let list: Array | undefined + const lines = (match[1] ?? '').split(/\r?\n/) + for (const line of lines) { + const item = /^\s*-\s+(.*)$/.exec(line) + if (item && list) { + list.push(unquote(item[1] ?? '')) + continue + } + list = undefined + const pair = /^([\w-]+):(.*)$/.exec(line) + if (!pair) continue + const key = pair[1] ?? '' + const value = (pair[2] ?? '').trim() + const inline = /^\[(.*)\]$/.exec(value) + if (value === '') { + list = [] + fields[key] = list + } else if (inline) { + fields[key] = (inline[1] ?? '') + .split(',') + .map(unquote) + .filter((entry) => entry !== '') + } else { + fields[key] = unquote(value) + } + } + return { fields, body: text.slice(match[0].length) } +} + +/** + * The Markdown files in `dir`: the name without `.md`, the frontmatter + * fields, and the body. A folder that cannot be read has no files. + */ +export async function readMarkdownFiles(dir: string) { + let files: Array = [] + try { + files = (await readdir(dir)).filter((file) => extname(file) === '.md') + } catch { + return [] + } + return Promise.all( + files.map(async (file) => ({ + name: basename(file, '.md'), + ...splitFrontmatter(await readFile(resolve(dir, file), 'utf8')), + })), + ) +} + +/** + * One command per Markdown file in `dir` (like `.claude/commands/*.md`). + * Running `/name args` sends the file as a prompt, with `$ARGUMENTS` + * replaced by the args. + */ +export function fileCommands(options: { dir: string }) { + return definePlugin({ + name: 'tanstack/file-commands', + setup: async (ctx) => { + const files = await readMarkdownFiles(options.dir) + const commands: Record = {} + for (const { name, fields, body } of files) { + const { description } = fields + commands[name] = defineCommand({ + description: + typeof description === 'string' && description !== '' + ? description + : `Run ${name}.md`, + run: (input: unknown) => { + const args = + typeof input === 'string' + ? input + : input === undefined + ? '' + : JSON.stringify(input) + ctx.session.prompt(body.split('$ARGUMENTS').join(args).trim()) + return `Sent /${name}.` + }, + }) + } + return { commands } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/glob.ts b/packages/ai-harness/src/first-party/glob.ts new file mode 100644 index 0000000000..8dc2b9e857 --- /dev/null +++ b/packages/ai-harness/src/first-party/glob.ts @@ -0,0 +1,14 @@ +/** A simple glob to RegExp: `**` any path, `*` any name part, `?` one character. */ +export function globToRegExp(glob: string) { + let pattern = '' + for (let index = 0; index < glob.length; index += 1) { + const char = glob.charAt(index) + if (char === '*' && glob[index + 1] === '*') { + pattern += '.*' + index += glob[index + 2] === '/' ? 2 : 1 + } else if (char === '*') pattern += '[^/]*' + else if (char === '?') pattern += '[^/]' + else pattern += char.replace(/[.+^${}()|[\]\\]/g, '\\$&') + } + return new RegExp(`^${pattern}$`) +} diff --git a/packages/ai-harness/src/first-party/index.ts b/packages/ai-harness/src/first-party/index.ts new file mode 100644 index 0000000000..374deb8d06 --- /dev/null +++ b/packages/ai-harness/src/first-party/index.ts @@ -0,0 +1,31 @@ +// First-party plugins. This entry is Node only: some plugins use the file +// system and the shell. +export { modelPicker } from './model-picker' +export { + PERMISSION_MODES, + PermissionDecisionCapability, + PermissionResources, + PermissionRules, + decidePermission, + deleteSavedPermission, + isUnsplittableCommand, + listSavedPermissions, + permissions, +} from './permissions' +export type { + CallResources, + PermissionDecision, + PermissionMode, + PermissionRule, + ToolResources, +} from './permissions' +export { globToRegExp } from './glob' +export { formatTodos, todos } from './todos' +export type { Todo } from './todos' +export { fileCommands, projectInstructions } from './files' +export { compact, usage } from './session-tools' +export { agents, builtInAgents } from './agents' +export type { AgentProfile } from './agents' +export { question } from './question' +export { TitleFailed, title } from './title' +export { boundToolOutput } from './bound-output' diff --git a/packages/ai-harness/src/first-party/model-picker.ts b/packages/ai-harness/src/first-party/model-picker.ts new file mode 100644 index 0000000000..30261ae7f8 --- /dev/null +++ b/packages/ai-harness/src/first-party/model-picker.ts @@ -0,0 +1,68 @@ +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { definePlugin } from '../plugins' +import type { AnyTextAdapter, KeyedAdapter } from '@tanstack/ai' + +/** + * Switch the main model at the next turn with the `model` setting or the + * `/model ` command. A choice can be a `keyedAdapter(...)`: the + * session builds it for each turn with the user's key. + * + * @example + * ```ts + * modelPicker({ + * choices: { + * fast: openaiText('gpt-5.6-luna'), + * claude: keyedAdapter(anthropicByok, (key) => + * createAnthropicChat('claude-sonnet-4-5', key), + * ), + * }, + * default: 'fast', + * }) + * ``` + */ +export function modelPicker(options: { + choices: Record> + default?: string +}) { + const names = Object.keys(options.choices) + const fallback = options.default ?? names[0] + if (fallback === undefined) + throw new Error('modelPicker needs at least one choice.') + return definePlugin({ + name: 'tanstack/model-picker', + setup: (ctx) => ({ + config: { + model: configOption.select({ + options: names, + default: fallback, + category: 'model', + description: 'The main model', + }), + }, + adapter: () => { + const name = ctx.config.get('model') + return typeof name === 'string' ? options.choices[name] : undefined + }, + commands: { + model: defineCommand({ + description: `Show or switch the model (${names.join(', ')})`, + run: async (input: unknown) => { + const name = + typeof input === 'string' + ? input + : typeof input === 'object' && input !== null && 'name' in input + ? String(input.name) + : undefined + if (!name) + return `Model: ${String(ctx.config.get('model'))}. Choices: ${names.join(', ')}.` + if (!names.includes(name)) + return `Unknown model "${name}". Choices: ${names.join(', ')}.` + await ctx.session.setConfig('model', name) + return `Model: ${name}. It applies at the next turn.` + }, + }), + }, + }), + }) +} diff --git a/packages/ai-harness/src/first-party/permissions.ts b/packages/ai-harness/src/first-party/permissions.ts new file mode 100644 index 0000000000..e8dd018975 --- /dev/null +++ b/packages/ai-harness/src/first-party/permissions.ts @@ -0,0 +1,728 @@ +import { createCapability, getMetadata } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { createExtensionPoint } from '../extensions' +import { SessionMetadata, definePlugin } from '../plugins' +import { isRecord } from '../utils' +import { globToRegExp } from './glob' +import type { AnyChatMiddleware, JSONSchema, MetadataStore } from '@tanstack/ai' + +export type PermissionDecision = 'allow' | 'ask' | 'deny' + +/** A rule for one tool. `kind` lets the modes treat edits and commands differently. */ +export interface PermissionRule { + /** A tool name. A trailing `*` matches a prefix, for example `git_*`. */ + tool: string + /** + * A glob on what the call touches: a path for file tools (`src/**`, + * relative to the root), or a command part for shell tools (`git *`). + * Without it, the rule matches every call of the tool. + */ + resource?: string + decision: PermissionDecision + kind?: 'read' | 'edit' | 'execute' +} + +/** + * Tool plugins add their rules here. `permissions()` reads them, and adds its + * own `rules` too, so other readers like `codeMode()` see them. + */ +export const PermissionRules = createExtensionPoint( + 'tanstack/permission-rules', +) + +/** + * What the calls of one tool touch. Each function gets the raw input the + * model sent, before the schema check. A function that throws refuses the + * call. + */ +export interface ToolResources { + /** The files the call reads or writes, relative to the root or absolute. */ + paths?: (input: unknown) => ReadonlyArray + /** The shell commands the call runs, one part for each simple command. */ + commands?: (input: unknown) => ReadonlyArray +} + +/** + * Tool plugins say what each call touches, so rules with a `resource` can + * match it. Each item maps tool names to their resources, for example + * `PermissionResources.item({ read_file: { paths: (input) => [...] } })`. + */ +export const PermissionResources = createExtensionPoint< + Readonly> +>('tanstack/permission-resources') + +export const PERMISSION_MODES = [ + 'default', + 'plan', + 'acceptEdits', + 'bypass', +] as const +export type PermissionMode = (typeof PERMISSION_MODES)[number] + +/** + * What `permissions()` decides for a call of `tool` in `mode`. With + * `resources`, for a call that touches them. Without, for a call that + * declares no `PermissionResources`. It uses all rules in their order, the + * `default`, and the `root`. Other plugins, like `codeMode()` and + * `workspaceTools()`, ask it. Saved answers are not used: they only allow + * more, so an answer here is never looser than a call. + */ +export const PermissionDecisionCapability = createCapability< + ( + tool: string, + mode: PermissionMode, + resources?: CallResources, + ) => PermissionDecision +>()('tanstack/permission-decision') + +/** + * @internal Ask the rules of `permissions()` about `tool`, an action that is + * not a tool call, like `formatter:prettier`. It uses the saved answers. On + * `ask`, it asks the user with `message`, and saves an `always` answer. + * Resolves to `true` when the action can run. The package does not export it. + */ +export const PermissionPrompt = createCapability< + (tool: string, message: string) => Promise +>()('tanstack/permission-prompt') + +/** + * True when a command part still has shell syntax that can run or chain + * more commands (`$(`, a backtick, `;`, `&`, `|`, `<`, `>`, a new line), so + * no allow rule applies to it. + */ +export function isUnsplittableCommand(part: string) { + return /[`;&|<>\r\n]|\$\(/.test(part) +} + +/** What one call touches, as a tool's `ToolResources` returns it. */ +export interface CallResources { + paths?: ReadonlyArray + commands?: ReadonlyArray +} + +/** One thing a call touches, ready to match rule globs against. */ +interface Resource { + type: 'path' | 'command' + /** The text rules match: a path relative to the root, or a command part. */ + match: string + /** Asks even when a rule without a `resource` allows the tool. */ + risky: boolean + /** No allow rule applies. */ + unsplittable: boolean + /** Compare without letter case (a Windows path). */ + foldCase: boolean +} + +function matchesTool(rule: PermissionRule, tool: string) { + return rule.tool.endsWith('*') + ? tool.startsWith(rule.tool.slice(0, -1)) + : rule.tool === tool +} + +/** A path or a drive root in Windows form: `C:\x`, `C:/x`, or `\\server\share`. */ +function isWindowsPath(path: string) { + return /^([a-z]:[\\/]|[\\/]{2})/i.test(path) +} + +/** + * Compare paths without letter case on macOS and Windows, whose file systems + * ignore case by default, and for any Windows-form path. Else `Secrets/key` + * would pass a deny rule for `secrets/**`. + */ +// ponytail: platform default, add a `caseInsensitive` option if a case-sensitive macOS volume needs it +function foldsCase(path: string) { + const platform = typeof process === 'undefined' ? '' : process.platform + return platform === 'darwin' || platform === 'win32' || isWindowsPath(path) +} + +/** `path` with `/` separators, and `.` and `..` resolved. */ +function normalizePath(path: string) { + const slashed = path.replaceAll('\\', '/') + const prefix = + /^[a-z]:\//i.exec(slashed)?.[0] ?? (slashed.startsWith('/') ? '/' : '') + const segments: Array = [] + const parts = slashed.slice(prefix.length).split('/') + for (const segment of parts) { + if (segment === '' || segment === '.') continue + const climbs = + segment === '..' && segments.length > 0 && segments.at(-1) !== '..' + if (climbs) segments.pop() + // `..` above an absolute root stays at the root. + else if (segment !== '..' || prefix === '') segments.push(segment) + } + return prefix + segments.join('/') +} + +/** A `.env` file, or `.env.local` and the like. They often hold secrets. */ +function isEnvFile(path: string) { + const name = path.slice(path.lastIndexOf('/') + 1).toLowerCase() + return /\.env(\.|$)/.test(name) +} + +function pathResource(path: string, root: string | undefined) { + const foldCase = foldsCase(root ?? path) + const slashed = path.replaceAll('\\', '/') + // Forms we cannot resolve without the file system ask: `C:x` (relative to + // the current folder of drive C), a stream (`a.txt:secret`), and a name + // that ends in `.` or a space (Windows drops them). + const cannotResolve = + slashed.replace(/^[a-z]:\//i, '').includes(':') || + slashed + .split('/') + .some((part) => /[. ]$/.test(part) && part !== '.' && part !== '..') + const isAbsolute = /^([a-z]:)?\//i.test(slashed) + const full = normalizePath( + isAbsolute || root === undefined ? slashed : `${root}/${slashed}`, + ) + const resource = { + type: 'path' as const, + match: full, + risky: cannotResolve || isEnvFile(full), + unsplittable: false, + foldCase, + } + if (root === undefined) return resource + const base = normalizePath(root) + const folder = base.endsWith('/') ? base : `${base}/` + const fold = (value: string) => (foldCase ? value.toLowerCase() : value) + const isInside = + fold(full) === fold(base) || fold(full).startsWith(fold(folder)) + // Inside the root, rules match the relative path. Outside, the full path. + return isInside + ? { ...resource, match: full.slice(folder.length) } + : { ...resource, risky: true } +} + +function commandResource(part: string) { + return { + type: 'command' as const, + match: part.trim(), + risky: false, + unsplittable: isUnsplittableCommand(part), + foldCase: false, + } +} + +/** A command glob: `*` matches any text, `/` too. */ +function commandPattern(glob: string) { + const source = glob + .split('*') + .map((part) => part.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')) + .join('.*') + return new RegExp(`^${source}$`) +} + +function matchesResource(glob: string, resource: Resource) { + if (resource.type === 'command') { + return commandPattern(glob).test(resource.match) + } + const pattern = globToRegExp(glob.replaceAll('\\', '/')) + return new RegExp(pattern.source, resource.foldCase ? 'i' : '').test( + resource.match, + ) +} + +function strictest(decisions: ReadonlyArray) { + if (decisions.includes('deny')) return 'deny' + if (decisions.includes('ask')) return 'ask' + return 'allow' +} + +/** + * The decision for a tool call in a mode. The last matching rule wins. + * + * - With `resources`, each path and command part is checked, and every one + * must be allowed. A path outside `root`, a `.env` file, or a path that + * cannot be resolved asks, unless a rule with a `resource` allows it. A + * command part with shell syntax (see `isUnsplittableCommand`) never gets + * an allow. + * - Without `resources`, the call could touch anything, so a later rule + * with a `resource` that asks or denies counts too. + * + * Modes: `bypass` allows all. `plan` denies edits, commands, and every call + * that is not an allow. `acceptEdits` allows edits that would ask. + */ +export function decidePermission( + rules: ReadonlyArray, + tool: string, + mode: PermissionMode, + options: { + fallback?: PermissionDecision + resources?: CallResources + /** The workspace root. Without it, no path counts as outside. */ + root?: string + } = {}, +) { + if (mode === 'bypass') return 'allow' + // An allow rule without `kind` must not hide that a tool edits. + const kind = rules.findLast( + (rule) => matchesTool(rule, tool) && rule.kind !== undefined, + )?.kind + if (mode === 'plan' && (kind === 'edit' || kind === 'execute')) return 'deny' + const decisionOf = (rule: PermissionRule | undefined) => { + const decision = rule?.decision ?? options.fallback ?? 'allow' + const autoEdit = + mode === 'acceptEdits' && + kind === 'edit' && + decision === 'ask' && + rule?.resource === undefined + return autoEdit ? 'allow' : decision + } + const toolRules = rules.filter((rule) => matchesTool(rule, tool)) + const lastToolLevel = toolRules.findLastIndex( + (rule) => rule.resource === undefined, + ) + const toolLevel = lastToolLevel === -1 ? undefined : toolRules[lastToolLevel] + const decideOne = (resource: Resource) => { + const rule = toolRules.findLast( + (candidate) => + candidate.resource === undefined || + matchesResource(candidate.resource, resource), + ) + const decision = decisionOf(rule) + if (decision !== 'allow') return decision + const needsAsk = + resource.unsplittable || (resource.risky && rule?.resource === undefined) + return needsAsk ? 'ask' : 'allow' + } + const resources = + options.resources === undefined + ? undefined + : [ + ...(options.resources.paths ?? []).map((path) => + pathResource(path, options.root), + ), + ...(options.resources.commands ?? []).map(commandResource), + ] + // Unknown resources: the rules with a `resource` after the tool rule could + // match, so the strictest of them counts. + const decisions = + resources === undefined + ? [ + decisionOf(toolLevel), + ...toolRules.slice(lastToolLevel + 1).map((rule) => rule.decision), + ] + : resources.length === 0 + ? [decisionOf(toolLevel)] + : resources.map(decideOne) + const decision = strictest(decisions) + if (mode === 'plan' && decision !== 'allow') return 'deny' + return decision +} + +/** Rules hide a tool when its last rule denies with no `resource`. */ +function isHidden( + rules: ReadonlyArray, + tool: string, + mode: PermissionMode, +) { + if (mode === 'bypass') return false + const last = rules.findLast((rule) => matchesTool(rule, tool)) + return last?.decision === 'deny' && last.resource === undefined +} + +/** The answer to a permission question, for hosts that render a form. */ +const ANSWER_SCHEMA: JSONSchema = { + type: 'object', + properties: { + answer: { + type: 'string', + enum: ['once', 'always', 'reject'], + description: + 'once: allow this call. always: allow it from now on in this project. reject: refuse it.', + }, + message: { + type: 'string', + description: 'With reject: what the model should do instead.', + }, + }, + required: ['answer'], +} + +/** The user's reply. Anything that is not `once` or `always` rejects. */ +function readAnswer(value: unknown) { + const raw = isRecord(value) ? value.answer : value + const choice = typeof raw === 'string' ? raw.trim().toLowerCase() : '' + const text = + isRecord(value) && typeof value.message === 'string' + ? value.message.trim() + : '' + return { + choice: choice === 'once' || choice === 'always' ? choice : 'reject', + message: text === '' ? undefined : text, + } +} + +function isSavedRule(value: unknown): value is PermissionRule { + return ( + isRecord(value) && + typeof value.tool === 'string' && + value.decision === 'allow' && + (value.resource === undefined || typeof value.resource === 'string') + ) +} + +/** A refused call: the model gets a tool error with `error`. */ +function refuse(error: string) { + return { type: 'skip' as const, result: { error } } +} + +const METADATA_NAMESPACE = 'tanstack/permissions' + +/** The root as one string for each folder: `C:\Repo` and `c:/repo/` are the same. */ +function projectId(root: string) { + const path = normalizePath(root) + return foldsCase(root) ? path.toLowerCase() : path +} + +/** Where the metadata store keeps the saved rules of `project`. */ +function savedKeyOf(project: string | undefined) { + return project === undefined + ? 'permissions:saved' + : `permissions:saved:${projectId(project)}` +} + +async function readSaved(metadata: MetadataStore, key: string) { + const stored = await metadata.get(METADATA_NAMESPACE, key) + return Array.isArray(stored) ? stored.filter(isSavedRule) : [] +} + +function isSameRule(a: PermissionRule, b: PermissionRule) { + return ( + a.tool === b.tool && + a.resource === b.resource && + a.decision === b.decision && + a.kind === b.kind + ) +} + +/** + * The rules that `always` answers saved for `project`, in the saved order. + * `project` is the `root` of `permissions()`. `stores` is the persistence + * of the host, for example `persistence.stores`. Without a metadata store, + * it resolves to an empty list. + */ +export async function listSavedPermissions( + stores: { metadata?: MetadataStore }, + project: string | undefined, +) { + return stores.metadata ? readSaved(stores.metadata, savedKeyOf(project)) : [] +} + +/** + * Delete one saved rule of `project`. The rule matches by its fields, as + * `listSavedPermissions` gives them. An unknown rule does nothing. New + * sessions stop using the rule. Open sessions can keep it until they open + * again. + */ +export async function deleteSavedPermission( + stores: { metadata?: MetadataStore }, + project: string | undefined, + rule: PermissionRule, +) { + const { metadata } = stores + if (!metadata) return + const key = savedKeyOf(project) + const rules = await readSaved(metadata, key) + // Remove every copy, so a forgotten rule never stays in force. + const kept = rules.filter((saved) => !isSameRule(saved, rule)) + if (kept.length === rules.length) return + await metadata.set(METADATA_NAMESPACE, key, kept) +} + +/** A saved rule as one line: the tool, then the resource. */ +function describeRule(rule: PermissionRule) { + return rule.resource === undefined + ? rule.tool + : `${rule.tool} ${rule.resource}` +} + +const PLAN_PROMPT = + 'You are in plan mode. Do not change files or run commands. Read what you need, then describe your plan.' + +/** + * Check every tool call against permission rules, with a `mode` setting and a + * `/mode` command: + * + * - `default`: rules apply as written. `ask` asks the user. + * - `plan`: read-only. Edits and commands are denied. + * - `acceptEdits`: edits run without asking. Commands still ask. + * - `bypass`: everything runs. + * + * Rules apply in this order, and the last match wins: rules from tool + * plugins (`PermissionRules`), then `rules`, then saved answers. A saved + * answer never beats a deny. Tool plugins say what a call touches through + * `PermissionResources`. A path outside `root` or a `.env` file asks, unless + * a rule with a `resource` allows it. A tool whose last rule denies with no + * `resource` is removed from the model's tools. + * + * A question takes `once`, `always`, or `reject` with a `message` for the + * model. `always` saves an allow rule for each path or command part of the + * call, for the project (`root`), in the metadata store, else in memory for + * the session. `/permissions` lists the saved rules, and + * `/permissions forget ` deletes one. + */ +export function permissions( + options: { + rules?: ReadonlyArray + default?: PermissionDecision + /** + * The workspace root. A path outside it asks, and saved answers belong + * to it. Without it, no path counts as outside. + */ + root?: string + } = {}, +) { + return definePlugin({ + name: 'tanstack/permissions', + provides: [PermissionDecisionCapability, PermissionPrompt], + setup: (ctx) => { + const contributed = ctx.collect(PermissionRules) + const declared = ctx.collect(PermissionResources) + // The rules of tool plugins, then the user's rules, so the user's win. + const configured = () => [...contributed, ...(options.rules ?? [])] + ctx.provide(PermissionDecisionCapability, (tool, current, resources) => + decidePermission(configured(), tool, current, { + fallback: options.default, + resources, + root: options.root, + }), + ) + const mode = (): PermissionMode => { + const value = ctx.config.get('mode') + return ( + PERMISSION_MODES.find((candidate) => candidate === value) ?? 'default' + ) + } + const root = options.root + const savedKey = savedKeyOf(root) + // Answers saved in this session. Runs without a metadata store use them. + const memory: Array = [] + const loadSaved = async (metadata: MetadataStore | undefined) => { + const rules = metadata ? await readSaved(metadata, savedKey) : [] + return [...rules, ...memory] + } + // ponytail: get + set, a parallel save can drop a rule. It then asks again. + const save = async ( + metadata: MetadataStore | undefined, + rules: ReadonlyArray, + ) => { + memory.push(...rules) + if (!metadata || rules.length === 0) return + const current = await readSaved(metadata, savedKey) + const added = rules.filter( + (rule) => !current.some((saved) => isSameRule(saved, rule)), + ) + if (added.length === 0) return + await metadata.set(METADATA_NAMESPACE, savedKey, [...current, ...added]) + } + // The session's store, for commands. A chat run reads its own. + const sessionMetadata = ctx.getOptional(SessionMetadata) + const savedRules = async () => + sessionMetadata ? readSaved(sessionMetadata, savedKey) : [...memory] + const forget = async (arg: string) => { + const rules = await savedRules() + const rule = /^\d+$/.test(arg) ? rules[Number(arg) - 1] : undefined + if (!rule) { + return `No saved rule ${arg}. Run /permissions to see the list.` + } + await deleteSavedPermission({ metadata: sessionMetadata }, root, rule) + // This session stops using the rule at once. + const kept = memory.filter((saved) => !isSameRule(saved, rule)) + memory.splice(0, memory.length, ...kept) + return `Forgot rule ${arg}: ${describeRule(rule)}.` + } + const resourcesOf = (tool: string, input: unknown) => { + const entry = declared.findLast((item) => Object.hasOwn(item, tool))?.[ + tool + ] + if (!entry) return undefined + const resources = { + paths: entry.paths?.(input) ?? [], + commands: entry.commands?.(input) ?? [], + } + const values = [...resources.paths, ...resources.commands] + if (!values.every((value) => typeof value === 'string')) { + throw new Error('A resource is not a string.') + } + return resources + } + /** The rules an `always` answer saves: one allow for each resource. */ + const rulesToSave = ( + tool: string, + resources: CallResources | undefined, + ) => { + if (resources === undefined) { + return [{ tool, decision: 'allow' as const }] + } + const values = [ + ...(resources.paths ?? []).map( + (path) => pathResource(path, root).match, + ), + ...(resources.commands ?? []) + .filter((part) => !isUnsplittableCommand(part)) + .map((part) => part.trim()), + ] + // A `*` or `?` would turn the saved rule into a wider glob. + return values + .filter((value) => !/[*?]/.test(value)) + .map((resource) => ({ tool, resource, decision: 'allow' as const })) + } + const decide = async ( + tool: string, + input: unknown, + current: PermissionMode, + metadata: MetadataStore | undefined, + ) => { + const resources = resourcesOf(tool, input) + const settings = { fallback: options.default, resources, root } + const rules = configured() + // A saved answer turns an ask into an allow. It never beats a deny, + // so only an ask reads the saved answers. + const first = decidePermission(rules, tool, current, settings) + if (first !== 'ask') return { decision: first, resources } + rules.push(...(await loadSaved(metadata))) + const decision = decidePermission(rules, tool, current, settings) + return { decision, resources } + } + /** Ask the user. An `always` answer saves the rules for the call. */ + const askUser = async ( + message: string, + tool: string, + resources: CallResources | undefined, + metadata: MetadataStore | undefined, + ) => { + // A closed session rejects the question: that is a reject too. + const reply = readAnswer( + await ctx.session + .ask({ + message: `${message} Answer once, always, or reject.`, + schema: ANSWER_SCHEMA, + }) + .catch(() => undefined), + ) + if (reply.choice === 'always') { + // The session keeps the answer in memory, so a failed write does + // not fail the call the user allowed. + await save(metadata, rulesToSave(tool, resources)).catch(() => {}) + } + return reply + } + ctx.provide(PermissionPrompt, async (tool, message) => { + const { decision } = await decide( + tool, + undefined, + mode(), + sessionMetadata, + ) + if (decision !== 'ask') return decision === 'allow' + const reply = await askUser(message, tool, undefined, sessionMetadata) + return reply.choice !== 'reject' + }) + const check = { + name: 'tanstack/permissions', + onBeforeToolCall: async (run, hook) => { + const current = mode() + if (current === 'bypass') return undefined + const metadata = getMetadata(run, { optional: true }) + const checked = await decide( + hook.toolName, + hook.args, + current, + metadata, + ).catch((error: unknown) => + error instanceof Error ? error : new Error(String(error)), + ) + // Fail closed: a call we cannot check does not run. + if (checked instanceof Error) { + return refuse( + `Cannot check the permissions of this call: ${checked.message}`, + ) + } + const { decision, resources } = checked + if (decision === 'allow') return undefined + if (decision === 'deny') { + return refuse( + current === 'plan' + ? 'This tool is not allowed in plan mode.' + : 'A permission rule denies this tool.', + ) + } + const preview = JSON.stringify(hook.args ?? {}).slice(0, 300) + const reply = await askUser( + `Allow ${hook.toolName} ${preview}?`, + hook.toolName, + resources, + metadata, + ) + if (reply.choice === 'reject') { + return refuse(reply.message ?? 'The user denied this tool call.') + } + return undefined + }, + } satisfies AnyChatMiddleware + return { + config: { + mode: configOption.select({ + options: PERMISSION_MODES, + default: 'default', + category: 'mode', + description: 'How tool calls are approved', + }), + }, + prompts: [ + { + id: 'tanstack/permissions:plan', + text: () => (mode() === 'plan' ? PLAN_PROMPT : ''), + }, + ], + commands: { + mode: defineCommand({ + description: `Show or switch the mode (${PERMISSION_MODES.join(', ')})`, + run: async (input: unknown) => { + const next = typeof input === 'string' ? input.trim() : '' + if (!next) return `Mode: ${mode()}.` + const receipt = await ctx.session.setConfig('mode', next) + return receipt.status === 'rejected' + ? `Unknown mode "${next}". Modes: ${PERMISSION_MODES.join(', ')}.` + : `Mode: ${next}.` + }, + }), + permissions: defineCommand({ + description: + 'List the saved "always" rules: /permissions. Delete one: /permissions forget ', + run: async (input: unknown) => { + const arg = typeof input === 'string' ? input.trim() : '' + const forgetArg = /^forget\s+(.+)$/.exec(arg)?.[1] + if (forgetArg !== undefined) return forget(forgetArg.trim()) + if (arg !== '') { + return 'Use /permissions, or /permissions forget .' + } + const rules = await savedRules() + if (rules.length === 0) return 'No saved rules.' + const lines = rules.map( + (rule, index) => `${index + 1}. ${describeRule(rule)}`, + ) + return ['Saved rules:', ...lines].join('\n') + }, + }), + }, + prepareTools: ({ tools }) => { + const rules = configured() + const current = mode() + return tools.filter((tool) => !isHidden(rules, tool.name, current)) + }, + // Other readers, like codeMode(), must see the user's rules too. The + // checks above add them again at the end, so they still win there. + contribute: (options.rules ?? []).map((rule) => + PermissionRules.item(rule), + ), + // The lead turn and every agent run check the same rules. Each run + // gets only one of the two lists, so a call is checked once. + middleware: [check], + agentMiddleware: [check], + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/question.ts b/packages/ai-harness/src/first-party/question.ts new file mode 100644 index 0000000000..7b59088692 --- /dev/null +++ b/packages/ai-harness/src/first-party/question.ts @@ -0,0 +1,189 @@ +import { toolDefinition } from '@tanstack/ai' +import { definePlugin } from '../plugins' +import { isRecord } from '../utils' + +/** One question the model asks, as the tool input has it. */ +interface ToolQuestion { + question: string + header?: string + options: Array<{ label: string; description?: string }> + multiple?: boolean +} + +const NO_ANSWER = 'The user did not answer.' + +function isOption(value: unknown) { + return ( + isRecord(value) && + typeof value.label === 'string' && + (value.description === undefined || typeof value.description === 'string') + ) +} + +function isToolQuestion(value: unknown): value is ToolQuestion { + return ( + isRecord(value) && + typeof value.question === 'string' && + (value.header === undefined || typeof value.header === 'string') && + Array.isArray(value.options) && + value.options.every(isOption) && + (value.multiple === undefined || typeof value.multiple === 'boolean') + ) +} + +/** The text the user sees: the question, the numbered options, and how to answer. */ +function messageOf(item: ToolQuestion) { + const title = item.header ? `${item.header}: ${item.question}` : item.question + const options = item.options.map( + (option, index) => + `${index + 1}. ${option.label}${option.description ? ` - ${option.description}` : ''}`, + ) + const how = + item.multiple === true + ? 'Type one or more numbers or labels, split by commas, or your own answer.' + : 'Type a number or a label, or your own answer.' + return [title, ...options, how].join('\n') +} + +/** + * The answer shape, for hosts that render a form: one label, or a list of + * labels. Hosts that show a text box send the typed text, and the tool reads it. + */ +function answerSchema(item: ToolQuestion) { + const choice = { + type: 'string', + enum: item.options.map((option) => option.label), + } + return item.multiple === true + ? { type: 'array', items: choice, minItems: 1, uniqueItems: true } + : choice +} + +/** The option `text` names, by its label (any letter case) or its number. */ +function labelOf(text: string, options: ToolQuestion['options']) { + const lower = text.toLowerCase() + const byLabel = options.find( + (option) => option.label.trim().toLowerCase() === lower, + ) + if (byLabel) return byLabel.label + return /^\d+$/.test(text) ? options[Number(text) - 1]?.label : undefined +} + +/** + * The answer line for the model: the picked labels, or the user's own + * words. `undefined` when the reply has no text. + */ +function readAnswer(reply: unknown, item: ToolQuestion) { + const isSplit = item.multiple === true && typeof reply === 'string' + const parts = Array.isArray(reply) + ? reply + : isSplit + ? reply.split(',') + : [reply] + const texts = parts + .map((part) => + typeof part === 'string' || typeof part === 'number' + ? String(part).trim() + : '', + ) + .filter((text) => text !== '') + if (texts.length === 0) return undefined + const picked = texts + .map((text) => labelOf(text, item.options)) + .filter((label) => label !== undefined) + const isPicked = + picked.length === texts.length && + (item.multiple === true || picked.length === 1) + if (isPicked) return `Answer: ${[...new Set(picked)].join(', ')}` + const own = typeof reply === 'string' ? reply.trim() : texts.join(', ') + return `Own answer: ${own}` +} + +/** + * Add a `question` tool: the model asks the user one to four questions, + * each with two to six options, and waits for the answers. Each question + * goes to the user through `ctx.session.ask`, one at a time. The user picks + * by number or label, or types their own answer. With no answer, or when + * the session closes, the model gets a tool error. + * + * @example + * ```ts + * plugins: () => [question()] + * ``` + */ +export function question() { + return definePlugin({ + name: 'tanstack/question', + setup: (ctx) => ({ + tools: [ + toolDefinition({ + name: 'question', + description: + 'Ask the user one or more questions, each with options, and wait for the answers. Use it when you need a decision or a preference from the user to continue. The user can also type their own answer.', + inputSchema: { + type: 'object', + properties: { + questions: { + type: 'array', + minItems: 1, + maxItems: 4, + items: { + type: 'object', + properties: { + question: { + type: 'string', + description: 'The full question.', + }, + header: { + type: 'string', + description: 'A short label for the question.', + }, + options: { + type: 'array', + minItems: 2, + maxItems: 6, + items: { + type: 'object', + properties: { + label: { type: 'string' }, + description: { type: 'string' }, + }, + required: ['label'], + }, + }, + multiple: { + type: 'boolean', + description: 'The user can pick more than one option.', + }, + }, + required: ['question', 'options'], + }, + }, + }, + required: ['questions'], + }, + // Asking again after a crash has no side effect. + replay: 'safe', + }).server(async (args: unknown) => { + const items = + isRecord(args) && Array.isArray(args.questions) + ? args.questions + : [] + if (!items.every(isToolQuestion)) + throw new Error('Each question needs a question and options.') + const answers: Array = [] + for (const item of items) { + // A closed session rejects the question: that is no answer too. + const reply = await ctx.session + .ask({ message: messageOf(item), schema: answerSchema(item) }) + .catch(() => undefined) + const answer = readAnswer(reply, item) + if (answer === undefined) throw new Error(NO_ANSWER) + answers.push(`${item.question}\n${answer}`) + } + return answers.join('\n\n') + }), + ], + }), + }) +} diff --git a/packages/ai-harness/src/first-party/session-tools.ts b/packages/ai-harness/src/first-party/session-tools.ts new file mode 100644 index 0000000000..64565ce0e8 --- /dev/null +++ b/packages/ai-harness/src/first-party/session-tools.ts @@ -0,0 +1,272 @@ +import { chat } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import { emptyUsage } from '../usage' +import type { + AnyChatMiddleware, + AnyTextAdapter, + KeyedAdapter, + ModelMessage, +} from '@tanstack/ai' +import type { SessionUsage, UsageCounts } from '../usage' + +const SUMMARY_PROMPT = + 'Summarize the conversation so far for yourself. Keep decisions, open tasks, file names, and facts you still need. Leave out small talk.' + +export function textOf(message: ModelMessage) { + if (typeof message.content === 'string') return message.content + if (Array.isArray(message.content)) { + return message.content + .map((part) => (part.type === 'text' ? part.content : '')) + .join('') + } + return '' +} + +/** One `role: text` line per message. Messages with no text are left out. */ +export function transcriptText(messages: ReadonlyArray) { + return messages + .map((message) => `${message.role}: ${textOf(message)}`) + .filter((line) => !line.endsWith(': ')) + .join('\n') +} + +/** + * `/compact`: replace a long transcript with a summary, so later turns send + * fewer tokens. The summary is written by `adapter`. A `keyedAdapter(...)` + * is built with the user's key when `/compact` runs. + */ +export function compact(options: { + adapter: AnyTextAdapter | KeyedAdapter + keepLast?: number +}) { + return definePlugin({ + name: 'tanstack/compact', + setup: (ctx) => ({ + commands: { + compact: defineCommand({ + description: 'Summarize the conversation to save tokens', + run: async () => { + const messages = await ctx.session.transcript() + const keep = options.keepLast ?? 0 + if (messages.length <= keep + 2) + return 'The conversation is already short.' + // Never keep a tool result without the call before it: the + // provider rejects it. Move the cut back to that call. + let cut = messages.length - keep + while (cut > 0 && messages[cut]?.role === 'tool') cut-- + if (cut < 2) return 'The conversation is already short.' + const older = messages.slice(0, cut) + const transcript = transcriptText(older) + const { text: summary } = await chat({ + adapter: await ctx.keys.adapter(options.adapter), + messages: [ + { role: 'user', content: `${SUMMARY_PROMPT}\n\n${transcript}` }, + ], + stream: false, + }) + await ctx.session.replaceTranscript([ + { + role: 'user', + content: `Summary of our conversation so far:\n${summary}`, + }, + { + role: 'assistant', + content: 'Understood. I will continue from this summary.', + }, + ...messages.slice(cut), + ]) + return `Compacted ${older.length} messages into a summary.` + }, + }), + }, + }), + }) +} + +/** Prices in USD per 1M tokens. */ +interface ModelPrices { + input: number + output: number + cacheRead?: number + cacheWrite?: number +} + +/** The usage of an index entry before the host writes the tokens of a turn. */ +const NO_TOKENS = { + turns: 0, + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + cachedTokens: 0, + cacheWriteTokens: 0, +} + +/** The cost in USD of some model calls of one model. */ +function priceOf(prices: ModelPrices, counts: UsageCounts) { + const { cachedTokens, cacheWriteTokens } = counts + // `promptTokens` counts the cache parts too, and each part has its own price. + const uncached = counts.promptTokens - cachedTokens - cacheWriteTokens + const part = (price: number | undefined, tokens: number) => + ((price ?? 0) / 1_000_000) * tokens + return ( + part(prices.input, uncached) + + part(prices.output, counts.completionTokens) + + part(prices.cacheRead, cachedTokens) + + part(prices.cacheWrite, cacheWriteTokens) + ) +} + +/** One `/usage` line: the calls and tokens of `counts`. */ +function usageLine(counts: UsageCounts) { + // Only a provider that reports its prompt cache has these counts. + const hasCache = counts.cachedTokens > 0 || counts.cacheWriteTokens > 0 + const cache = hasCache + ? ` (${counts.cachedTokens} cache read, ${counts.cacheWriteTokens} cache write)` + : '' + return `${counts.calls} model calls, ${counts.promptTokens} input tokens${cache}, ${counts.completionTokens} output tokens, ${counts.totalTokens} total.` +} + +/** + * Show the token usage of the session (`session.usage()`): the lead turn and + * every agent run (subagents, background agents, and their children). + * `/usage` shows the totals, with the input tokens read from and written to + * the prompt cache, and a line for each model when there is more than one. + * The plugin state, for a UI, has a copy of `session.usage().total` and + * `contextTokens`, the size of the lead model's context at its latest call. + * + * A model call that has a cost from its provider keeps that cost. With + * `model`, the plugin also prices the calls that have no provider cost. + * `model` gets the model id and returns its prices in USD per 1M tokens, or + * `undefined` when it does not know the model. `/usage` shows the cost when + * a call has one, and says how many calls it could not price. When a call + * has a cost, the plugin writes the cost to `usage.cost` of the session + * index entry. + * + * @example + * ```ts + * const prices: Record = { + * 'claude-sonnet-5-5': { input: 3, output: 15 }, + * } + * const priceOf = (id: string) => (prices[id] ? { cost: prices[id] } : undefined) + * plugins: () => [usage({ model: priceOf })] + * ``` + */ +export function usage( + options: { + model?: (modelId: string) => { cost: ModelPrices } | undefined + } = {}, +) { + return definePlugin({ + name: 'tanstack/usage', + setup: (ctx) => { + /** + * A copy of the session totals, and `contextTokens`: the input tokens + * of the latest model call of the lead turn, so how much of the + * model's context the conversation fills now. + */ + const state = ctx.state({ ...emptyUsage().total, contextTokens: 0 }) + // The cost of the model calls of one model, else `undefined`. + const costOf = (model: string, counts: UsageCounts) => { + if (counts.cost !== undefined) return counts.cost + // The key is `provider/model`. The lookup gets the model id only. + const prices = options.model?.(model.slice(model.indexOf('/') + 1)) + return prices && priceOf(prices.cost, counts) + } + /** + * The cost of the session, and how many calls have no cost. `cost` is + * `undefined` when no call has one. + * ponytail: core sums by model. When a model has a provider cost, its + * calls with no provider cost add nothing. Upgrade: core keeps, by + * model, the counts of the calls with no provider cost. + */ + const sessionCost = (totals: SessionUsage) => { + let cost: number | undefined + let unpriced = 0 + for (const [model, counts] of Object.entries(totals.byModel)) { + const modelCost = costOf(model, counts) + if (modelCost === undefined) unpriced += counts.calls + else cost = (cost ?? 0) + modelCost + } + return { cost, unpriced } + } + const show = () => { + const totals = ctx.session.snapshot().usage + const { cost, unpriced } = sessionCost(totals) + const note = + unpriced > 0 + ? ` (cost unknown for ${unpriced} ${unpriced === 1 ? 'call' : 'calls'})` + : '' + const isPriced = options.model !== undefined || cost !== undefined + const priced = isPriced + ? ` Cost: $${(cost ?? 0).toFixed(4)}${note}.` + : '' + const lines = [`${usageLine(totals.total)}${priced}`] + const models = Object.entries(totals.byModel) + // One model has the same counts as the total. + if (models.length > 1) { + for (const [model, counts] of models) { + lines.push(`${model}: ${usageLine(counts)}`) + } + } + return lines.join('\n') + } + // The host writes the token totals to the index entry at the end of + // each turn. This plugin writes only `cost`, and keeps the tokens of + // the entry. + // ponytail: a host with a log counts a call when its log write lands, + // which can be after the run ends. Then the entry and the state get + // that call at the end of the next run or call. Upgrade: the host + // writes the cost at turn end. + const saveCost = async () => { + const { cost } = sessionCost(ctx.session.snapshot().usage) + if (cost === undefined) return + // A change function: it runs again on the new entry when another + // write landed first, so the tokens of that write stay. Before the + // first turn ends, the entry has no usage. + await ctx.session.updateEntry((entry) => ({ + usage: { ...(entry.usage ?? NO_TOKENS), cost }, + })) + } + // Copy the session totals to the state. Only a call of the lead turn + // sets the context size. + const mirror = (contextTokens?: number) => + state.update((current) => ({ + ...ctx.session.snapshot().usage.total, + contextTokens: contextTokens ?? current.contextTokens, + })) + // At the end of each run, the state gets the totals and the index entry + // gets the cost so far. The index is only a list of sessions, and the + // state only a copy: a failed write does not fail the run. The next run + // writes them again. + const save = async () => { + await Promise.all([mirror(), saveCost()]).catch(() => {}) + } + const atRunEnd = { + name: 'tanstack/usage', + onFinish: save, + onAbort: save, + onError: save, + } satisfies AnyChatMiddleware + return { + middleware: [ + { + ...atRunEnd, + onUsage: async (_run, info) => { + await mirror(info.promptTokens) + }, + } satisfies AnyChatMiddleware, + ], + // An agent's middleware runs before the host counts its call, so its + // run copies the totals only at its end. + agentMiddleware: [atRunEnd], + commands: { + usage: defineCommand({ + description: 'Show token usage for this session', + run: show, + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/title.ts b/packages/ai-harness/src/first-party/title.ts new file mode 100644 index 0000000000..929e97dad1 --- /dev/null +++ b/packages/ai-harness/src/first-party/title.ts @@ -0,0 +1,90 @@ +import { chat } from '@tanstack/ai' +import { createPluginEvent } from '../extensions' +import { definePlugin } from '../plugins' +import { textOf } from './session-tools' +import type { + AnyChatMiddleware, + AnyTextAdapter, + KeyedAdapter, + ModelMessage, +} from '@tanstack/ai' + +const TITLE_PROMPT = + 'Write a short title for a conversation that starts with the message below. Use at most 60 characters on one line, with no quotes. Answer with the title only.' + +/** Sent when the title call fails. The turn does not fail. */ +export const TitleFailed = createPluginEvent<{ message: string }>( + 'tanstack/title:failed', +) + +/** The first line of `answer`, with no quotes around it, at most 60 characters. */ +function cleanTitle(answer: string) { + const line = answer.trim().split('\n')[0] ?? '' + return line + .replace(/^["'`]+|["'`]+$/g, '') + .trim() + .slice(0, 60) +} + +/** + * Give each session a title. At the first turn of a session with no title, + * `adapter` writes a short title from the first user message, and the plugin + * saves it as the `title` of the session index entry. The call runs next to + * the turn: it does not slow the turn, and a failure does not fail it. A + * failure sends a {@link TitleFailed} event, and the next turn tries again. + * Without `stores.sessions`, the plugin does nothing. + * + * @example + * ```ts + * plugins: () => [title({ adapter: openaiText('gpt-5.4-nano') })] + * ``` + */ +export function title(options: { + adapter: AnyTextAdapter | KeyedAdapter +}) { + return definePlugin({ + name: 'tanstack/title', + setup: (ctx) => { + let isNaming = false + const name = async (messages: ReadonlyArray) => { + const first = messages.find((message) => message.role === 'user') + const entry = await ctx.session.entry() + // No sessions store, or the session has a title (also after a restart). + if (!first || !entry || entry.title) return + // ponytail: the first 2000 characters are enough for a title. + const message = textOf(first).slice(0, 2000) + const answer = await chat({ + adapter: await ctx.keys.adapter(options.adapter), + messages: [ + { role: 'user', content: `${TITLE_PROMPT}\n\n${message}` }, + ], + stream: false, + }) + const text = cleanTitle(answer.text) + if (text) await ctx.session.updateEntry({ title: text }) + } + return { + middleware: [ + { + name: 'tanstack/title', + onConfig: (run, config) => { + if (run.phase !== 'init' || isNaming) return + isNaming = true + // Not awaited: the title call never slows or fails the turn. + void name(config.messages) + .catch((error: unknown) => + ctx.emit(TitleFailed, { + message: + error instanceof Error ? error.message : String(error), + }), + ) + .finally(() => { + isNaming = false + }) + }, + } satisfies AnyChatMiddleware, + ], + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/todos.ts b/packages/ai-harness/src/first-party/todos.ts new file mode 100644 index 0000000000..eae2c2cce8 --- /dev/null +++ b/packages/ai-harness/src/first-party/todos.ts @@ -0,0 +1,100 @@ +import { toolDefinition } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' + +export interface Todo { + text: string + status: 'pending' | 'in_progress' | 'done' +} + +function isTodo(value: unknown): value is Todo { + return ( + typeof value === 'object' && + value !== null && + 'text' in value && + typeof value.text === 'string' && + 'status' in value && + (value.status === 'pending' || + value.status === 'in_progress' || + value.status === 'done') + ) +} + +const MARK = { pending: '[ ]', in_progress: '[~]', done: '[x]' } as const + +export function formatTodos(items: ReadonlyArray): string { + return items.map((item) => `${MARK[item.status]} ${item.text}`).join('\n') +} + +/** + * A todo list the model keeps for multi-step work (like Claude Code's + * TodoWrite). The list is plugin state, so clients see it change and it + * survives restarts. `/todos` shows it. + */ +export function todos() { + return definePlugin({ + name: 'tanstack/todos', + setup: async (ctx) => { + const state = ctx.state<{ items: Array }>({ items: [] }) + // The prompt runs for each turn, so keep a copy of the list at hand. + let current = (await state.get()).items + return { + tools: [ + toolDefinition({ + name: 'todo_write', + description: + 'Replace the todo list. Use it for work with several steps. Mark one item in_progress at a time.', + inputSchema: { + type: 'object', + properties: { + todos: { + type: 'array', + items: { + type: 'object', + properties: { + text: { type: 'string' }, + status: { + type: 'string', + enum: ['pending', 'in_progress', 'done'], + }, + }, + required: ['text', 'status'], + }, + }, + }, + required: ['todos'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const list = + typeof args === 'object' && + args !== null && + 'todos' in args && + Array.isArray(args.todos) + ? args.todos + : [] + if (!list.every(isTodo)) + throw new Error('Each todo needs text and a status.') + current = (await state.update(() => ({ items: list }))).items + return formatTodos(current) || 'The todo list is empty.' + }), + ], + prompts: [ + { + id: 'tanstack/todos:list', + text: () => + current.length > 0 + ? `Current todo list:\n${formatTodos(current)}` + : '', + }, + ], + commands: { + todos: defineCommand({ + description: 'Show the todo list', + run: () => formatTodos(current) || 'The todo list is empty.', + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/workspace-hooks.ts b/packages/ai-harness/src/first-party/workspace-hooks.ts new file mode 100644 index 0000000000..17421b70fc --- /dev/null +++ b/packages/ai-harness/src/first-party/workspace-hooks.ts @@ -0,0 +1,35 @@ +import { createExtensionPoint } from '../extensions' + +/** + * Hooks for the workspace tools. Each runs inside the tool, while the tool + * holds the lock on that path. A hook that throws fails the tool call. + */ +export interface WorkspaceHooks { + /** + * Runs after a tool wrote the file at `path`, an absolute path. Return + * text to add to the result, or nothing to add nothing. + */ + afterWrite?: (path: string) => Promise + /** + * Runs after `read_file` read the file at `path`, an absolute path. Return + * text to add to the result, or `undefined` to add nothing. + */ + afterRead?: (path: string) => Promise +} + +/** + * Plugins add {@link WorkspaceHooks} here. The workspace tools run them. + * + * @example + * ```ts + * definePlugin({ + * name: 'acme/format', + * setup: () => ({ + * contribute: [WorkspaceHooks.item({ afterWrite: async (path) => format(path) })], + * }), + * }) + * ``` + */ +export const WorkspaceHooks = createExtensionPoint( + 'tanstack/workspace-hooks', +) diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index a6de481a21..41c2a5231b 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -49,6 +49,21 @@ export type { } from './plugins' export { createExtensionPoint, createPluginEvent } from './extensions' +// Portable (no Node imports), so edge-safe plugins such as code mode can read +// the permission rules and resources without the Node-only `./plugins` entry. +export { + PermissionDecisionCapability, + PermissionResources, + PermissionRules, + decidePermission, +} from './first-party/permissions' +export type { + CallResources, + PermissionDecision, + PermissionMode, + PermissionRule, + ToolResources, +} from './first-party/permissions' export type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' export { checkConfigValue, configOption } from './config' diff --git a/packages/ai-harness/tests/agent-middleware.test.ts b/packages/ai-harness/tests/agent-middleware.test.ts index 2a33623009..89c43a0b20 100644 --- a/packages/ai-harness/tests/agent-middleware.test.ts +++ b/packages/ai-harness/tests/agent-middleware.test.ts @@ -1,7 +1,8 @@ import { describe, expect, it } from 'vitest' -import { EventType, defineAgent } from '@tanstack/ai' +import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { permissions, usage } from '../src/first-party' import { mockAdapter, text, toolCall } from './helpers' import type { AnyChatMiddleware, StreamChunk } from '@tanstack/ai' import type { AnyHarness, PluginLifetime } from '../src' @@ -177,4 +178,74 @@ describe('plugin agentMiddleware', () => { ]) await host.close() }) + + it('lets the usage() plugin count the model calls of every agent', async () => { + const researcher = parentAgent( + 'researcher', + textAgent('fetcher', 'fetched', 100), + [10, 20], + ) + const { host, session } = await open( + defineHarness({ + name: 'test/agent-middleware', + adapter: leadCalling('researcher', [1, 2]), + subagents: { agents: [researcher] }, + plugins: () => [usage()], + }), + ) + + await session.prompt('research') + + // The lead (1 + 2), the researcher (10 + 20), and the fetcher (100). + expect(await session.command('usage')).toBe( + '5 model calls, 133 input tokens, 0 output tokens, 133 total.', + ) + // The context size is the input of the lead's latest call, not an agent's. + expect(session.snapshot().plugins['tanstack/usage']).toMatchObject({ + contextTokens: 2, + }) + await host.close() + }) + + it('lets permissions() deny a write a subagent tries in plan mode', async () => { + const writes: Array = [] + const writeFile = toolDefinition({ + name: 'write_file', + description: 'Write a file', + }).server(async (args) => { + writes.push(args) + return 'written' + }) + const writerModel = mockAdapter([ + () => toolCall('write_file', { path: 'x.txt' }), + () => text('writer done'), + ]) + const writer = defineAgent({ + name: 'writer', + description: 'Writes files', + run: (ctx) => + ctx.chat({ adapter: writerModel.adapter, tools: [writeFile] }), + }) + const { host, session } = await open( + defineHarness({ + name: 'test/agent-middleware', + adapter: leadCalling('writer'), + subagents: { agents: [writer] }, + plugins: () => [ + permissions({ + rules: [{ tool: 'write_file', decision: 'ask', kind: 'edit' }], + }), + ], + }), + ) + await session.command('mode', 'plan') + + await session.prompt('write it') + + expect(writes).toEqual([]) + expect(JSON.stringify(writerModel.calls[1].messages)).toContain( + 'This tool is not allowed in plan mode.', + ) + await host.close() + }) }) diff --git a/packages/ai-harness/tests/agents.test.ts b/packages/ai-harness/tests/agents.test.ts new file mode 100644 index 0000000000..c167cbb76b --- /dev/null +++ b/packages/ai-harness/tests/agents.test.ts @@ -0,0 +1,446 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { z } from 'zod' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { permissions } from '../src/first-party' +import { agents } from '../src/first-party/agents' +import { splitFrontmatter } from '../src/first-party/files' +import { messageTexts, mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessPlugin } from '../src' +import type { AgentProfile } from '../src/first-party/agents' + +/** The names of the tools that ran, in order. */ +const ran: Array = [] + +beforeEach(() => { + ran.length = 0 +}) + +const tool = (name: string) => + toolDefinition({ + name, + description: `The ${name} tool`, + inputSchema: z.object({}), + }).server(async () => { + ran.push(name) + return name + }) + +/** + * Open thread `t` with the tools `read_file` and `write_file`. The main + * model starts each subagent with a tool named after it. + */ +async function open(adapter: AnyTextAdapter, plugins: Array) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/agents', + adapter, + tools: [tool('read_file'), tool('write_file')], + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +/** The names of the tools the model got on one call. */ +function toolNames(call: { tools: Array<{ name: string }> }) { + return call.tools.map((entry) => entry.name) +} + +/** The system prompts of one call, as one string. */ +function promptsOf(call: { systemPrompts: unknown }) { + return JSON.stringify(call.systemPrompts) +} + +describe('splitFrontmatter', () => { + it('reads values and both list forms', () => { + const file = [ + '---', + 'description: "Reviews: code"', + 'tools: [read_file, grep]', + 'models:', + ' - fast', + " - 'smart'", + '---', + 'Body', + ].join('\n') + expect(splitFrontmatter(file)).toEqual({ + fields: { + description: 'Reviews: code', + tools: ['read_file', 'grep'], + models: ['fast', 'smart'], + }, + body: 'Body', + }) + }) + + it('keeps a file without frontmatter as the body', () => { + expect(splitFrontmatter('Just text')).toEqual({ + fields: {}, + body: 'Just text', + }) + }) +}) + +describe('agents', () => { + it('offers build and plan as primary agents, and general and explore as subagents', async () => { + const { adapter, calls } = mockAdapter([() => text('hi')]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter }), + ]) + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan.', + ) + await session.prompt('hi') + + expect(toolNames(calls[0])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + expect(promptsOf(calls[0])).toContain('You are the build agent.') + await host.close() + }) + + it('switches to plan with /agent: the next turn gets its prompt and read-only tools', async () => { + const { adapter, calls } = mockAdapter([ + () => text('built'), + () => text('planned'), + ]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter }), + ]) + + await session.prompt('build it') + expect(await session.command('agent', 'plan')).toBe( + 'Agent: plan. It applies at the next turn.', + ) + await session.prompt('plan it') + + expect(toolNames(calls[0])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + expect(toolNames(calls[1])).toEqual(['read_file']) + expect(promptsOf(calls[1])).toContain('You are the plan agent.') + expect(promptsOf(calls[1])).not.toContain('You are the build agent.') + expect(await session.command('agent', 'nope')).toBe( + 'Unknown agent "nope". Agents: build, plan.', + ) + await host.close() + }) + + it('switches the primary agent with the agent setting', async () => { + const { adapter, calls } = mockAdapter([() => text('planned')]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter }), + ]) + + expect((await session.setConfig('agent', 'plan')).status).toBe('accepted') + // A subagent is not a primary agent. + expect((await session.setConfig('agent', 'general')).status).toBe( + 'rejected', + ) + await session.prompt('plan it') + + expect(toolNames(calls[0])).toEqual(['read_file']) + await host.close() + }) + + it('switches the model when the agent has one, and keeps the harness model when not', async () => { + const main = mockAdapter([() => text('from main')]) + const fast = mockAdapter([() => text('from fast')]) + const byModel = (model: string) => { + if (model !== 'fast-model') throw new Error(`Unknown model ${model}`) + return fast.adapter + } + const quick: AgentProfile = { + name: 'quick', + description: 'Answers fast', + mode: 'primary', + model: 'fast-model', + } + const { host, session } = await open(main.adapter, [ + agents({ adapter: byModel, agents: [quick] }), + ]) + + expect(await session.prompt('one')).toEqual({ text: 'from main' }) + await session.command('agent', 'quick') + expect(await session.prompt('two')).toEqual({ text: 'from fast' }) + await host.close() + }) + + it('makes the call after `steps` model calls with no tool calls, then stops', async () => { + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', {}, 'c1'), + () => toolCall('read_file', {}, 'c2'), + () => toolCall('read_file', {}, 'c3'), + () => text('never asked'), + ]) + const short: AgentProfile = { + name: 'short', + description: 'Two steps', + mode: 'primary', + steps: 2, + } + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, agents: [short], default: 'short' }), + ]) + + await session.prompt('go') + + expect(calls.map((call) => call.toolChoice)).toEqual([ + undefined, + undefined, + 'none', + ]) + expect(promptsOf(calls[1])).not.toContain('step limit') + expect(promptsOf(calls[2])).toContain('step limit') + await host.close() + }) + + it('refuses a tool call at the step limit when the provider ignores toolChoice', async () => { + // Bedrock Converse sends `none` as `auto` after tool calls, so the model + // can still call a tool on the last call. + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', {}, 'c1'), + () => toolCall('read_file', {}, 'c2'), + () => text('Here is what I have.'), + ]) + const short: AgentProfile = { + name: 'short', + description: 'One step', + mode: 'primary', + steps: 1, + } + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, agents: [short], default: 'short' }), + ]) + + await session.prompt('go') + + // The run ends after the last call, and only the first call ran. + expect(calls).toHaveLength(2) + expect(calls[1].toolChoice).toBe('none') + expect(ran).toEqual(['read_file']) + + // The next turn shows the model the refusal of the last call. + await session.prompt('and now?') + const refused = calls[2].messages.find( + (message: { role: string; toolCallId?: string }) => + message.role === 'tool' && message.toolCallId === 'c2', + ) + expect(JSON.stringify(refused?.content)).toContain('Step limit reached') + expect(ran).toEqual(['read_file']) + await host.close() + }) + + it('runs a subagent agent as a child of its tool', async () => { + const lead = mockAdapter([ + () => toolCall('explore', { prompt: 'Find the config' }), + () => text('Done.'), + ]) + const child = mockAdapter(() => text('It is in config.ts')) + const models: Array = [] + const { host, session } = await open(lead.adapter, [ + agents({ + adapter: (model) => { + models.push(model) + return child.adapter + }, + }), + ]) + + await session.prompt('Where is the config?') + + // Without its own model, the child uses the model of the main turn. + expect(models).toEqual(['test-model']) + expect(toolNames(child.calls[0])).toEqual(['read_file']) + expect(promptsOf(child.calls[0])).toContain('You are a fast search agent.') + expect(messageTexts(child.calls[0]).at(-1)).toBe('Find the config') + expect(JSON.stringify(lead.calls[1].messages)).toContain( + 'It is in config.ts', + ) + await host.close() + }) + + it('applies the permission rules of an agent only while it is the primary agent', async () => { + const { adapter, calls } = mockAdapter([ + () => text('one'), + () => text('two'), + ]) + const careful: AgentProfile = { + name: 'careful', + description: 'Never writes', + mode: 'primary', + permissions: [{ tool: 'write_file', decision: 'deny' }], + } + const { host, session } = await open(adapter, [ + permissions(), + agents({ adapter: () => adapter, agents: [careful] }), + ]) + + await session.command('agent', 'careful') + await session.prompt('one') + await session.command('agent', 'build') + await session.prompt('two') + + expect(toolNames(calls[0])).toEqual(['read_file', 'general', 'explore']) + expect(toolNames(calls[1])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + await host.close() + }) + + it('replaces a built-in agent with one of the same name', async () => { + const { adapter, calls } = mockAdapter([() => text('planned')]) + const plan: AgentProfile = { + name: 'plan', + description: 'My plan', + mode: 'primary', + system: 'My own plan prompt.', + } + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, agents: [plan] }), + ]) + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan.', + ) + await session.command('agent', 'plan') + await session.prompt('plan it') + + expect(toolNames(calls[0])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + expect(promptsOf(calls[0])).toContain('My own plan prompt.') + await host.close() + }) + + it('drops the built-ins, and leaves hidden agents out of the list', async () => { + const { adapter, calls } = mockAdapter([() => text('hi')]) + const { host, session } = await open(adapter, [ + agents({ + adapter: () => adapter, + builtIns: false, + agents: [ + { name: 'solo', description: 'The only agent', mode: 'primary' }, + { + name: 'secret', + description: 'Not listed', + mode: 'primary', + hidden: true, + }, + ], + }), + ]) + + expect(await session.command('agent')).toBe('Agent: solo. Agents: solo.') + await session.prompt('hi') + + expect(toolNames(calls[0])).toEqual(['read_file', 'write_file']) + await host.close() + }) +}) + +describe('agents from Markdown files', () => { + let root = '' + beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-agents-')) + }) + afterEach(async () => { + await rm(root, { recursive: true, force: true }) + }) + + it('loads an agent with its frontmatter fields', async () => { + await writeFile( + join(root, 'reviewer.md'), + [ + '---', + 'description: Reviews code', + 'mode: all', + 'model: smart-model', + 'tools: [read_file]', + 'steps: 1', + '---', + 'You review code.', + ].join('\n'), + ) + const main = mockAdapter([() => text('built')]) + const smart = mockAdapter([ + () => toolCall('read_file', {}, 'c1'), + () => text('reviewed'), + ]) + const { host, session } = await open(main.adapter, [ + agents({ + adapter: (model) => + model === 'smart-model' ? smart.adapter : main.adapter, + dirs: [root], + }), + ]) + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan, reviewer.', + ) + await session.prompt('build it') + expect( + main.calls[0].tools.find( + (entry: { name: string }) => entry.name === 'reviewer', + )?.description, + ).toBe('Reviews code') + await session.command('agent', 'reviewer') + await session.prompt('review it') + + expect(toolNames(smart.calls[0])).toEqual(['read_file']) + expect(promptsOf(smart.calls[0])).toContain('You review code.') + expect(smart.calls.map((call) => call.toolChoice)).toEqual([ + undefined, + 'none', + ]) + await host.close() + }) + + it('reads the folder again on session.reload()', async () => { + const { adapter } = mockAdapter([]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, dirs: [root] }), + ]) + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan.', + ) + + await writeFile(join(root, 'writer.md'), 'You write docs.') + await session.reload() + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan, writer.', + ) + await host.close() + }) + + it('refuses an agent file with an unknown mode', async () => { + await writeFile(join(root, 'bad.md'), '---\nmode: sometimes\n---\nText') + const { adapter } = mockAdapter([]) + + await expect( + open(adapter, [agents({ adapter: () => adapter, dirs: [root] })]), + ).rejects.toThrow('Agent file bad.md: "mode" is not valid.') + }) +}) diff --git a/packages/ai-harness/tests/bound-output.test.ts b/packages/ai-harness/tests/bound-output.test.ts new file mode 100644 index 0000000000..b4c257e053 --- /dev/null +++ b/packages/ai-harness/tests/bound-output.test.ts @@ -0,0 +1,314 @@ +import { mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { defineAgent, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { boundText, boundToolOutput } from '../src/first-party/bound-output' +import { mockAdapter, text, toolCall } from './helpers' +import type { AnyTool, ModelMessage } from '@tanstack/ai' +import type { AnyAgent, AnyHarness } from '../src' + +const DAY = 24 * 60 * 60 * 1000 +/** Five lines, 34 bytes. */ +const FIVE_LINES = 'line 1\nline 2\nline 3\nline 4\nline 5' +/** FIVE_LINES cut to `maxLines: 3`. */ +const FIRST_THREE = + 'line 1\nline 2\nline 3\n\n[Output cut. Showing the first 3 of 5 lines, 20 of 34 bytes.]' + +let dir: string + +beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-bound-output-')) +}) + +afterEach(async () => { + await rm(dir, { recursive: true, force: true }) +}) + +/** A tool named `lookup` that returns `output`. */ +function lookup(output: unknown) { + return toolDefinition({ + name: 'lookup', + description: 'Look something up', + }).server(async () => output) +} + +/** A model that calls `tool` once, then answers. */ +function callsOnce(tool: string) { + return mockAdapter([() => toolCall(tool, {}), () => text('done')]) +} + +/** The tool message in the messages of one model call. */ +function toolMessageIn(call: { messages: Array }) { + return call.messages.find((message) => message.role === 'tool') +} + +/** Open a session on `harness`, send one prompt, and close it. */ +async function promptOnce(harness: AnyHarness) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 't' }) + await session.prompt('go') + await host.close() +} + +/** The lead model calls `tool`. Gives back the tool message the model saw. */ +async function leadToolMessage( + tool: AnyTool, + options: Parameters[0] = { maxLines: 3 }, +) { + const model = callsOnce(tool.name) + await promptOnce( + defineHarness({ + name: 'test/bound-output', + adapter: model.adapter, + tools: [tool], + plugins: () => [boundToolOutput(options)], + }), + ) + return toolMessageIn(model.calls[1]) +} + +/** The lead model calls `lookup`, which returns `output`. Gives back the result the model saw. */ +async function leadSees( + output: unknown, + options?: Parameters[0], +) { + const message = await leadToolMessage(lookup(output), options) + return message?.content +} + +/** The lead model calls `agent` as a subagent. Gives back the result the model saw, parsed. */ +async function leadSeesAgent(agent: AnyAgent) { + const lead = callsOnce(agent.name) + await promptOnce( + defineHarness({ + name: 'test/bound-output', + adapter: lead.adapter, + subagents: { agents: [agent] }, + plugins: () => [boundToolOutput({ maxLines: 3 })], + }), + ) + return JSON.parse(String(toolMessageIn(lead.calls[1])?.content)) +} + +/** The paths of the outputs saved in `dir`. */ +async function savedOutputs() { + const names = await readdir(dir) + return names + .filter((name) => name.startsWith('tool-output-')) + .map((name) => join(dir, name)) +} + +describe('boundText', () => { + it.each([ + [ + 'too many lines', + 'a\nb\nc\nd', + { maxLines: 2, maxBytes: 100 }, + 'a\nb\n\n[Output cut. Showing the first 2 of 4 lines, 3 of 7 bytes.]', + ], + [ + 'too many bytes', + 'abcdef', + { maxLines: 10, maxBytes: 4 }, + 'abcd\n\n[Output cut. Showing the first 1 of 1 lines, 4 of 6 bytes.]', + ], + [ + 'a two-byte character over the byte limit', + 'aéé', + { maxLines: 10, maxBytes: 4 }, + 'aé\n\n[Output cut. Showing the first 1 of 1 lines, 3 of 5 bytes.]', + ], + [ + 'a surrogate pair over the byte limit', + 'x\u{1F600}', + { maxLines: 10, maxBytes: 4 }, + 'x\n\n[Output cut. Showing the first 1 of 1 lines, 1 of 5 bytes.]', + ], + ])('keeps the head of text with %s', (_case, input, limits, expected) => { + expect(boundText(input, limits)).toEqual({ + text: expected, + truncated: true, + }) + }) + + it.each([ + [ + 'too many lines', + 'a\nb\nc\nd', + { maxLines: 2, maxBytes: 100 }, + '[Output cut. Showing the last 2 of 4 lines, 3 of 7 bytes.]\n\nc\nd', + ], + [ + 'too many lines and a final newline', + 'a\nb\nc\n', + { maxLines: 2, maxBytes: 100 }, + '[Output cut. Showing the last 2 of 3 lines, 4 of 6 bytes.]\n\nb\nc\n', + ], + [ + 'too many bytes', + 'abcdef', + { maxLines: 10, maxBytes: 4 }, + '[Output cut. Showing the last 1 of 1 lines, 4 of 6 bytes.]\n\ncdef', + ], + [ + 'a two-byte character over the byte limit', + 'ééa', + { maxLines: 10, maxBytes: 4 }, + '[Output cut. Showing the last 1 of 1 lines, 3 of 5 bytes.]\n\néa', + ], + [ + 'a surrogate pair over the byte limit', + '\u{1F600}x', + { maxLines: 10, maxBytes: 4 }, + '[Output cut. Showing the last 1 of 1 lines, 1 of 5 bytes.]\n\nx', + ], + ])('keeps the tail of text with %s', (_case, input, limits, expected) => { + expect(boundText(input, { ...limits, keep: 'tail' })).toEqual({ + text: expected, + truncated: true, + }) + }) + + it('leaves text inside both limits, and a final newline is not an extra line', () => { + expect(boundText('a\nb\n', { maxLines: 2, maxBytes: 4 })).toEqual({ + text: 'a\nb\n', + truncated: false, + }) + }) +}) + +describe('boundToolOutput', () => { + it('cuts a long string result to its first lines, with a note', async () => { + expect(await leadSees(FIVE_LINES)).toBe(FIRST_THREE) + }) + + it('cuts any other value as JSON text', async () => { + const seen = await leadSees({ name: 'x'.repeat(200) }, { maxBytes: 20 }) + + expect(seen).toBe( + `{"name":"${'x'.repeat(11)}\n\n[Output cut. Showing the first 1 of 1 lines, 20 of 211 bytes.]`, + ) + }) + + it('leaves a short result alone and saves no file', async () => { + const seen = await leadSees('line 1\nline 2', { maxLines: 3, dir }) + + expect(seen).toBe('line 1\nline 2') + expect(await savedOutputs()).toEqual([]) + }) + + it('leaves content parts alone', async () => { + const image = [ + { + type: 'image', + source: { type: 'data', value: 'A'.repeat(500), mimeType: 'image/png' }, + }, + ] + + expect(await leadSees(image, { maxBytes: 100 })).toEqual(image) + }) + + it('cuts the error text of a failed call, and it stays an error', async () => { + const failing = toolDefinition({ + name: 'lookup', + description: 'Look something up', + }).server(async () => { + throw new Error(FIVE_LINES) + }) + + const message = await leadToolMessage(failing) + + expect(message).toMatchObject({ + content: JSON.stringify({ error: FIRST_THREE }), + error: FIRST_THREE, + }) + }) + + it('saves the full output in dir and gives its path in the note', async () => { + const seen = await leadSees(FIVE_LINES, { maxLines: 3, dir }) + + const [saved = ''] = await savedOutputs() + expect(saved).toMatch(/tool-output-\d+-call-1\.txt$/) + expect(seen).toBe(`${FIRST_THREE}\n[Full output saved to ${saved}.]`) + expect(await readFile(saved, 'utf8')).toBe(FIVE_LINES) + }) + + it('removes saved outputs older than retentionDays and keeps other files', async () => { + const old = `tool-output-${Date.now() - 8 * DAY}-old.txt` + const recent = `tool-output-${Date.now() - DAY}-recent.txt` + const names = [old, recent, 'notes.txt'] + await Promise.all(names.map((name) => writeFile(join(dir, name), 'x'))) + + await leadSees(FIVE_LINES, { maxLines: 3, dir, retentionDays: 7 }) + + const left = await readdir(dir) + expect(left).not.toContain(old) + expect(left).toEqual(expect.arrayContaining([recent, 'notes.txt'])) + expect(left).toHaveLength(3) + }) + + it('still cuts the result when the full output cannot be saved', async () => { + const file = join(dir, 'a-file') + await writeFile(file, 'x') + + const seen = await leadSees(FIVE_LINES, { + maxLines: 3, + dir: join(file, 'out'), + }) + + expect(seen).toEqual( + expect.stringContaining( + `${FIRST_THREE}\n[The full output was not saved: Error:`, + ), + ) + }) + + it('cuts a tool result in a subagent run too', async () => { + const researcherModel = callsOnce('lookup') + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks things up', + run: (ctx) => + ctx.chat({ + adapter: researcherModel.adapter, + tools: [lookup(FIVE_LINES)], + }), + }) + + await leadSeesAgent(researcher) + + expect(toolMessageIn(researcherModel.calls[1])?.content).toBe(FIRST_THREE) + }) + + it('cuts a subagent answer and keeps its subagentRunId', async () => { + const answers = mockAdapter([() => text(FIVE_LINES)]) + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks things up', + run: (ctx) => ctx.chat({ adapter: answers.adapter }), + }) + + expect(await leadSeesAgent(researcher)).toEqual({ + subagentRunId: expect.stringMatching(/.+/), + result: FIRST_THREE, + }) + }) + + it('leaves a failed subagent call alone, so it keeps its subagentRunId', async () => { + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks things up', + run: async () => { + throw new Error(FIVE_LINES) + }, + }) + + expect(await leadSeesAgent(researcher)).toEqual({ + subagentRunId: expect.stringMatching(/.+/), + error: FIVE_LINES, + }) + }) +}) diff --git a/packages/ai-harness/tests/discover-tools.test.ts b/packages/ai-harness/tests/discover-tools.test.ts new file mode 100644 index 0000000000..b5e6b15ec9 --- /dev/null +++ b/packages/ai-harness/tests/discover-tools.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { SessionEvent } from '../src' + +const echoTool = (name: string, reply: string) => + toolDefinition({ + name, + description: `The ${name} tool`, + inputSchema: z.object({}), + }).server(async () => reply) + +describe('discoverTools', () => { + it('adds tools a plugin finds at run time, skips taken names, and warns on failure', async () => { + let online = false + const connector = definePlugin({ + name: 'test/connector', + setup: () => ({ + discoverTools: () => + online + ? [ + echoTool('remote_search', 'remote result'), + echoTool('lookup', 'shadowed'), + ] + : [], + }), + }) + const broken = definePlugin({ + name: 'test/broken', + setup: () => ({ + discoverTools: async () => { + throw new Error('server is down') + }, + }), + }) + const { adapter, calls } = mockAdapter([ + () => text('nothing yet'), + () => toolCall('remote_search', {}), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/discover', + adapter, + tools: [echoTool('lookup', 'local result')], + plugins: () => [connector, broken], + }), + { threadId: 't' }, + ) + const names = (call: number) => + (calls[call].tools ?? []).map((tool: { name: string }) => tool.name) + + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: controller.signal })) + seen.push(entry) + })() + + await session.prompt('before sign-in') + expect(names(0)).toEqual(['lookup']) + + online = true + const turn = await session.prompt('after sign-in') + expect(turn.text).toBe('done') + // The static `lookup` wins over the discovered one with the same name. + expect(names(1)).toEqual(['lookup', 'remote_search']) + expect(JSON.stringify(calls[2].messages)).toContain('remote result') + + controller.abort() + await reading + const warnings = seen.filter( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.plugin.warning', + ) + expect(warnings[0]?.event).toMatchObject({ + value: { plugin: 'test/broken', message: 'server is down' }, + }) + await host.close() + }) +}) + +describe('the default agent loop', () => { + it('keeps calling the model past five steps', async () => { + const steps = Array.from( + { length: 7 }, + (_, index) => () => toolCall('lookup', {}, `call-${index}`), + ) + const { adapter, calls } = mockAdapter([...steps, () => text('finally')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/long-turn', + adapter, + tools: [echoTool('lookup', 'fact')], + }), + { threadId: 't' }, + ) + const turn = await session.prompt('dig deep') + expect(calls).toHaveLength(8) + expect(turn.text).toBe('finally') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/edge-safety.test.ts b/packages/ai-harness/tests/edge-safety.test.ts index e6f443f1e0..f585e5fd8a 100644 --- a/packages/ai-harness/tests/edge-safety.test.ts +++ b/packages/ai-harness/tests/edge-safety.test.ts @@ -17,9 +17,9 @@ const srcDir = resolve(here, '..', 'src') /** * Node-only entry points, relative to `src`. A path that ends with `/` allows - * the whole folder. The core has none. + * the whole folder. */ -const ALLOWED: Array = [] +const ALLOWED = ['first-party/files.ts', 'first-party/bound-output.ts'] const isAllowed = (file: string) => ALLOWED.some((entry) => diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts new file mode 100644 index 0000000000..1b6fc092a1 --- /dev/null +++ b/packages/ai-harness/tests/first-party.test.ts @@ -0,0 +1,178 @@ +import { mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { + compact, + decidePermission, + fileCommands, + globToRegExp, + modelPicker, + projectInstructions, + todos, + usage, +} from '../src/first-party' +import { messageTexts, mockAdapter, text, toolCall } from './helpers' +import type { HarnessPlugin } from '../src' +import type { AnyTextAdapter } from '@tanstack/ai' + +let root = '' +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-')) +}) +afterEach(async () => { + await rm(root, { recursive: true, force: true }) +}) + +async function open(adapter: AnyTextAdapter, plugins: Array) { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/first-party', + adapter, + plugins: () => plugins, + }), + { + threadId: 't', + }, + ) + return { host, session, persistence } +} + +describe('decidePermission', () => { + const rules = [ + { tool: 'read_file', decision: 'allow' as const, kind: 'read' as const }, + { tool: 'write_file', decision: 'ask' as const, kind: 'edit' as const }, + { tool: 'bash', decision: 'ask' as const, kind: 'execute' as const }, + ] + it('applies modes on top of the rules', () => { + expect(decidePermission(rules, 'write_file', 'default')).toBe('ask') + expect(decidePermission(rules, 'write_file', 'acceptEdits')).toBe('allow') + expect(decidePermission(rules, 'bash', 'acceptEdits')).toBe('ask') + expect(decidePermission(rules, 'write_file', 'plan')).toBe('deny') + expect(decidePermission(rules, 'read_file', 'plan')).toBe('allow') + expect(decidePermission(rules, 'bash', 'bypass')).toBe('allow') + }) +}) + +describe('globToRegExp', () => { + it('matches globs like the list tool', () => { + expect(globToRegExp('src/**/*.ts').test('src/a/b.ts')).toBe(true) + expect(globToRegExp('src/**/*.ts').test('src/b.ts')).toBe(true) + expect(globToRegExp('*.md').test('docs/a.md')).toBe(false) + }) +}) + +describe('modelPicker', () => { + it('switches the main model with /model', async () => { + const fast = mockAdapter([() => text('fast answer')]) + const smart = mockAdapter([() => text('smart answer')]) + const { host, session } = await open(smart.adapter, [ + modelPicker({ + choices: { smart: smart.adapter, fast: fast.adapter }, + default: 'smart', + }), + ]) + expect(await session.command('model', 'fast')).toContain('Model: fast') + expect(await session.prompt('hi')).toEqual({ text: 'fast answer' }) + expect(await session.command('model', 'nope')).toContain('Unknown model') + await host.close() + }) +}) + +describe('todos', () => { + it('keeps the list in state and shows it in the next prompt', async () => { + const { adapter, calls } = mockAdapter([ + () => + toolCall('todo_write', { + todos: [{ text: 'write tests', status: 'in_progress' }], + }), + () => text('planned'), + () => text('next'), + ]) + const { host, session, persistence } = await open(adapter, [todos()]) + await session.prompt('plan it') + expect(await session.command('todos')).toBe('[~] write tests') + await session.prompt('continue') + expect(JSON.stringify(calls[2].systemPrompts)).toContain('[~] write tests') + expect( + await persistence.stores.metadata.get('plugin:tanstack/todos', 't'), + ).toEqual({ + items: [{ text: 'write tests', status: 'in_progress' }], + }) + await host.close() + }) +}) + +describe('project files', () => { + it('adds AGENTS.md to the prompt and turns command files into commands', async () => { + await writeFile(join(root, 'AGENTS.md'), 'Use tabs.') + await mkdir(join(root, 'commands')) + await writeFile( + join(root, 'commands', 'review.md'), + '---\ndescription: Review a file\n---\nReview $ARGUMENTS carefully.', + ) + const { adapter, calls } = mockAdapter([() => text('reviewed')]) + const { host, session } = await open(adapter, [ + projectInstructions({ root }), + fileCommands({ dir: join(root, 'commands') }), + ]) + expect(session.commands()).toEqual([ + { + name: 'review', + description: 'Review a file', + owner: 'tanstack/file-commands', + }, + ]) + await session.command('review', 'src/app.ts') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + expect(messageTexts(calls[0])).toEqual(['Review src/app.ts carefully.']) + expect(JSON.stringify(calls[0].systemPrompts)).toContain('Use tabs.') + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await host.close() + }) +}) + +describe('compact and usage', () => { + it('replaces the transcript with a summary and counts tokens', async () => { + const summarizer = mockAdapter([() => text('We planned a trip.')]) + // Real adapters report usage on RUN_FINISHED. + const withUsage = (answer: string) => () => + text(answer).map((chunk) => + chunk.type === 'RUN_FINISHED' + ? { + ...chunk, + usage: [ + { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + cachedInputTokens: 6, + cacheWriteInputTokens: 2, + }, + ], + } + : chunk, + ) + const main = mockAdapter([withUsage('one'), withUsage('two')]) + const { host, session, persistence } = await open(main.adapter, [ + compact({ adapter: summarizer.adapter }), + usage(), + ]) + await session.prompt('first') + await session.prompt('second') + expect(await session.command('compact')).toBe( + 'Compacted 4 messages into a summary.', + ) + const transcript = await persistence.stores.messages.loadThread('t') + expect(transcript).toHaveLength(2) + expect(transcript[0]?.content).toContain('We planned a trip.') + expect(await session.command('usage')).toBe( + '2 model calls, 20 input tokens (12 cache read, 4 cache write), 10 output tokens, 30 total.', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/instructions.test.ts b/packages/ai-harness/tests/instructions.test.ts new file mode 100644 index 0000000000..8f8ef03d1b --- /dev/null +++ b/packages/ai-harness/tests/instructions.test.ts @@ -0,0 +1,234 @@ +import { mkdir, mkdtemp, rm, utimes, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { dirname, join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { projectInstructions } from '../src/first-party' +import { mockAdapter, text } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessPlugin } from '../src' + +// `~` in a global path is the home folder. Each test points it into its own +// folder. +const home = vi.hoisted(() => ({ dir: '' })) +vi.mock('node:os', async (importOriginal) => ({ + ...(await importOriginal()), + homedir: () => home.dir, +})) + +// A fake git: real git can take seconds to start on a busy machine, more +// than the plugin waits. The plugin runs git with `execFile`. The fake names +// `git.branch` for `git rev-parse` in `git.cwd`, and fails for anything else. +const git = vi.hoisted(() => { + const state: { cwd: string; branch: string | undefined } = { + cwd: '', + branch: undefined, + } + return state +}) +vi.mock('node:child_process', async (importOriginal) => ({ + ...(await importOriginal()), + execFile: ( + file: string, + args: Array, + options: { cwd?: string }, + callback: (error: Error | null, stdout: string) => void, + ) => { + const isBranchQuery = + file === 'git' && + args.join(' ') === 'rev-parse --abbrev-ref HEAD' && + options.cwd === git.cwd && + git.branch !== undefined + if (isBranchQuery) callback(null, `${git.branch}\n`) + else callback(new Error('git failed'), '') + }, +})) + +let dir = '' +/** The `root` of the plugin, in the test folder. */ +let root = '' +beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-instructions-')) + home.dir = join(dir, 'home') + root = join(dir, 'project') + await mkdir(root) +}) +afterEach(async () => { + vi.useRealTimers() + await rm(dir, { recursive: true, force: true }) +}) + +/** Write `content` to `path` in the test folder, and make its folders. */ +async function put(path: string, content: string) { + const full = join(dir, path) + await mkdir(dirname(full), { recursive: true }) + await writeFile(full, content) +} + +/** Write `content` to `path`, with an mtime `minutes` in the future. */ +async function change(path: string, content: string, minutes: number) { + await put(path, content) + const time = Date.now() / 1000 + minutes * 60 + await utimes(join(dir, path), time, time) +} + +/** A model that answers 'ok', on a model with mid-conversation changes. */ +function okModel() { + const { adapter, calls } = mockAdapter(() => text('ok')) + const withChannels: AnyTextAdapter = { + ...adapter, + midConversationChannels: { tools: true, systemPrompts: true }, + } + return { adapter: withChannels, calls } +} + +/** Open a session on `adapter` with `plugins`. */ +async function open(adapter: AnyTextAdapter, plugins: Array) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/instructions', + adapter, + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +/** The system prompts of one turn with `projectInstructions(options)`. */ +async function promptsOfOneTurn( + options: Parameters[0], +) { + const { adapter, calls } = okModel() + const { host, session } = await open(adapter, [projectInstructions(options)]) + await session.prompt('hi') + await host.close() + return calls[0].systemPrompts +} + +describe('projectInstructions', () => { + it('puts global files first, then the files from the repo root down to root', async () => { + await put('home/.config/AGENTS.md', 'Global rule.') + await mkdir(join(dir, 'repo', '.git'), { recursive: true }) + await put('repo/AGENTS.md', 'Repo rule.') + await put('repo/pkg/AGENTS.md', 'Package rule.') + await put('repo/pkg/CLAUDE.md', 'Claude rule.') + + expect( + await promptsOfOneTurn({ + root: join(dir, 'repo', 'pkg'), + global: ['~/.config/AGENTS.md'], + env: false, + }), + ).toEqual([ + 'Global instructions from ~/.config/AGENTS.md:\nGlobal rule.', + 'Project instructions from ../AGENTS.md:\nRepo rule.', + 'Project instructions from AGENTS.md:\nPackage rule.', + 'Project instructions from CLAUDE.md:\nClaude rule.', + ]) + }) + + it('reads only root when no folder has .git', async () => { + await put('AGENTS.md', 'Outer rule.') + await put('project/AGENTS.md', 'Inner rule.') + + expect(await promptsOfOneTurn({ root, env: false })).toEqual([ + 'Project instructions from AGENTS.md:\nInner rule.', + ]) + }) + + it('adds a changed file at the end, as a mid-conversation change', async () => { + await put('project/AGENTS.md', 'Use tabs.') + const { adapter, calls } = okModel() + const { host, session } = await open(adapter, [ + projectInstructions({ root, env: false }), + ]) + + await session.prompt('one') + await change('project/AGENTS.md', 'Use spaces.', 1) + await session.prompt('two') + await host.close() + + expect(calls[1].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nUse tabs.', + 'Project instructions from AGENTS.md changed. Follow this text instead:\nUse spaces.', + ]) + expect(calls[1].midConversationChanges).toEqual({ + start: { tools: [], systemPrompts: 1 }, + changes: [{ before: calls[1].messages.length, systemPrompts: 1 }], + }) + }) + + it('notes a removed file and a new file', async () => { + await put('project/AGENTS.md', 'Use tabs.') + const { adapter, calls } = okModel() + const { host, session } = await open(adapter, [ + projectInstructions({ root, env: false }), + ]) + + await session.prompt('one') + await rm(join(root, 'AGENTS.md')) + await put('project/CLAUDE.md', 'Be brief.') + await session.prompt('two') + await host.close() + + expect(calls[1].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nUse tabs.', + 'Project instructions from AGENTS.md were removed.', + 'Project instructions from CLAUDE.md:\nBe brief.', + ]) + }) + + it('builds the prompt again when no spare slot is left', async () => { + await put('project/AGENTS.md', 'Version 0.') + const { adapter, calls } = okModel() + // One file and 8 spare slots: 9 slots in all. + const { host, session } = await open(adapter, [ + projectInstructions({ root, files: ['AGENTS.md'], env: false }), + ]) + + await session.prompt('turn 0') + for (let version = 1; version <= 9; version++) { + await change('project/AGENTS.md', `Version ${version}.`, version) + await session.prompt(`turn ${version}`) + } + await host.close() + + expect(calls[8].systemPrompts.at(-1)).toBe( + 'Project instructions from AGENTS.md changed. Follow this text instead:\nVersion 8.', + ) + expect(calls[9].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nVersion 9.', + ]) + }) + + it('adds an environment block with the date, platform, and folder', async () => { + vi.setSystemTime(new Date('2026-03-04T12:00:00Z')) + + expect(await promptsOfOneTurn({ root })).toEqual([ + `Environment:\n- Date: 2026-03-04\n- Platform: ${process.platform}\n- Working folder: ${root}\n- Git repository: no`, + ]) + }) + + it('names the git branch in a git repository', async () => { + await mkdir(join(root, '.git')) + git.cwd = root + git.branch = 'instructions-test' + + const [environment] = await promptsOfOneTurn({ root }) + expect(environment).toMatch( + /- Git repository: yes\n- Git branch: instructions-test$/, + ) + }) + + it('leaves out the branch when git cannot name it', async () => { + await mkdir(join(root, '.git')) + git.cwd = root + git.branch = undefined + + const [environment] = await promptsOfOneTurn({ root }) + expect(environment).toMatch(/- Git repository: yes$/) + }) +}) diff --git a/packages/ai-harness/tests/permissions.test.ts b/packages/ai-harness/tests/permissions.test.ts new file mode 100644 index 0000000000..d6a64ed6c3 --- /dev/null +++ b/packages/ai-harness/tests/permissions.test.ts @@ -0,0 +1,780 @@ +import { describe, expect, it, vi } from 'vitest' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { + PermissionDecisionCapability, + PermissionPrompt, + PermissionResources, + PermissionRules, + decidePermission, + deleteSavedPermission, + listSavedPermissions, + permissions, +} from '../src/first-party/permissions' +import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessPersistence, HarnessPlugin, HarnessSession } from '../src' +import type { + CallResources, + PermissionDecision, + PermissionMode, + PermissionRule, +} from '../src/first-party/permissions' +import type { Reply } from './helpers' + +const ROOT = '/work/repo' +const WINDOWS_ROOT = 'C:\\work\\repo' +const readAllowed: Array = [ + { tool: 'read_file', decision: 'allow', kind: 'read' }, +] + +function read( + path: string, + options: { root?: string; rules?: Array } = {}, +) { + return decidePermission( + options.rules ?? readAllowed, + 'read_file', + 'default', + { + resources: { paths: [path] }, + root: options.root ?? ROOT, + }, + ) +} + +/** Run `fn` as if the host were `platform`. Letter case rules depend on it. */ +function onPlatform(platform: NodeJS.Platform, fn: () => T) { + const real = process.platform + Object.defineProperty(process, 'platform', { value: platform }) + try { + return fn() + } finally { + Object.defineProperty(process, 'platform', { value: real }) + } +} + +function run(commands: Array, rules: Array) { + return decidePermission(rules, 'bash', 'default', { + resources: { commands }, + }) +} + +describe('decidePermission with resources', () => { + it('lets the last matching rule win, and needs every resource allowed', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + { tool: 'read_file', resource: 'secrets/public/**', decision: 'allow' }, + ] + const decide = (paths: Array) => + decidePermission(rules, 'read_file', 'default', { + resources: { paths }, + root: ROOT, + }) + + expect(decide(['src/a.ts'])).toBe('allow') + expect(decide(['secrets/key'])).toBe('deny') + expect(decide(['secrets/public/logo.png'])).toBe('allow') + expect(decide(['src/a.ts', 'secrets/key'])).toBe('deny') + }) + + it.each([ + ['darwin', 'deny'], + ['win32', 'deny'], + ['linux', 'allow'], + ] as const)( + 'on %s, a deny rule for a posix path matches another letter case: %s', + (platform, expected) => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ] + expect(onPlatform(platform, () => read('Secrets/key', { rules }))).toBe( + expected, + ) + }, + ) + + it('fails closed when the resources of a call are not known', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'ask' }, + ] + + expect(decidePermission(rules, 'read_file', 'default')).toBe('ask') + }) + + it('matches command rules against each part, and * spans slashes', () => { + const rules: Array = [ + { tool: 'bash', decision: 'ask', kind: 'execute' }, + { tool: 'bash', resource: 'ls*', decision: 'allow' }, + { tool: 'bash', resource: 'git *', decision: 'allow' }, + ] + + expect(run(['ls -la'], rules)).toBe('allow') + expect(run(['git add src/a.ts'], rules)).toBe('allow') + expect(run(['ls', 'rm x'], rules)).toBe('ask') + }) + + it.each([ + ['a command that was not split', 'ls && rm x'], + ['command substitution', 'echo $(rm x)'], + ['a backtick', 'echo `rm x`'], + ['process substitution', 'diff <(rm x) y'], + ['a heredoc', 'cat < ~/.bashrc'], + ])('asks for %s even when * is allowed', (_name, part) => { + expect( + run([part], [{ tool: 'bash', resource: '*', decision: 'allow' }]), + ).toBe('ask') + }) + + it('still denies a part with shell syntax when a deny matches', () => { + const rules: Array = [ + { tool: 'bash', resource: '*', decision: 'allow' }, + { tool: 'bash', resource: 'echo *', decision: 'deny' }, + ] + + expect(run(['echo $(rm x)'], rules)).toBe('deny') + }) + + it.each([ + ['..', '../secret.txt'], + ['.. in the middle', 'src/../../secret.txt'], + ['an absolute path', '/etc/passwd'], + ['a folder that only starts like the root', '/work/repo-evil/a.ts'], + ])('asks for a path that leaves the workspace: %s', (_name, path) => { + expect(read(path)).toBe('ask') + }) + + it('asks for another letter case of the root on a case-sensitive system', () => { + expect(onPlatform('linux', () => read('/work/Repo/a.ts'))).toBe('ask') + }) + + it.each([ + ['another drive', 'D:\\other\\x.txt'], + ['a sibling folder in another letter case', 'C:/Work/Other/x.txt'], + ['a folder that only starts like the root', 'C:\\WORK\\REPO-evil\\x.txt'], + ['a path relative to the current folder of a drive', 'C:secret.txt'], + ['a stream of a file', 'notes.txt:hidden'], + ['a name that ends in a dot', 'notes.txt.'], + ])('asks for a Windows path that leaves the workspace: %s', (_name, path) => { + expect(read(path, { root: WINDOWS_ROOT })).toBe('ask') + }) + + it('compares Windows paths without letter case, for allows and denies', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ] + + expect(read('c:\\WORK\\Repo\\src\\a.ts', { root: WINDOWS_ROOT })).toBe( + 'allow', + ) + expect(read('SECRETS\\key', { root: WINDOWS_ROOT, rules })).toBe('deny') + }) + + it('lets a rule with a resource allow a path outside the workspace', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: '/shared/**', decision: 'allow' }, + ] + + expect(read('/shared/notes.md', { rules })).toBe('allow') + }) + + it.each(['.env', 'config/prod.env', '.env.local', 'Prod.ENV'])( + 'asks before reading %s', + (path) => { + expect(read(path)).toBe('ask') + }, + ) + + it('reads other files, and a .env file a rule names', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: '.env.example', decision: 'allow' }, + ] + + expect(read('environment.ts')).toBe('allow') + expect(read('.env.example', { rules })).toBe('allow') + }) + + it('keeps the edit kind in plan mode when a later rule allows the tool', () => { + const rules: Array = [ + { tool: 'write_file', decision: 'ask', kind: 'edit' }, + { tool: 'write_file', decision: 'allow' }, + ] + + expect(decidePermission(rules, 'write_file', 'plan')).toBe('deny') + }) + + it('in acceptEdits mode, still asks for an edit outside the workspace', () => { + const rules: Array = [ + { tool: 'write_file', decision: 'ask', kind: 'edit' }, + ] + const write = (path: string, mode: PermissionMode) => + decidePermission(rules, 'write_file', mode, { + resources: { paths: [path] }, + root: ROOT, + }) + + expect(write('src/a.ts', 'acceptEdits')).toBe('allow') + expect(write('../a.ts', 'acceptEdits')).toBe('ask') + }) +}) + +/** The string `key` of a call input. Throws when it is missing. */ +function field(input: unknown, key: string) { + const value = + typeof input === 'object' && input !== null + ? Reflect.get(input, key) + : undefined + if (typeof value !== 'string') throw new Error(`No ${key}.`) + return value +} + +/** + * `read_file` and `bash` stand-ins that say what each call touches. `ran` + * lists the tools that ran. The `bash` splitter splits on ` && ` only. + */ +function fakeTools() { + const ran: Array = [] + const tool = (name: string) => + toolDefinition({ name, description: name }).server(async () => { + ran.push(name) + return `${name} done` + }) + const plugin = definePlugin({ + name: 'test/tools', + setup: () => ({ + tools: [tool('read_file'), tool('bash'), tool('deploy')], + contribute: [ + PermissionRules.item({ + tool: 'read_file', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ + tool: 'bash', + decision: 'ask', + kind: 'execute', + }), + PermissionResources.item({ + read_file: { paths: (input) => [field(input, 'path')] }, + bash: { commands: (input) => field(input, 'command').split(' && ') }, + }), + ], + }), + }) + return { plugin, ran } +} + +async function open( + persistence: HarnessPersistence, + plugins: Array, + replies: Array, + threadId = 't', +) { + const { adapter, calls } = mockAdapter([...replies, () => text('done')]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/permissions', + adapter, + plugins: () => plugins, + }), + { threadId }, + ) + return { host, session, calls } +} + +/** Wait for the one open question, then answer it. Returns its message. */ +async function answer(session: HarnessSession, value: unknown) { + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + if (!question) throw new Error('No question.') + await session.answer(question.questionId, value) + return question.message +} + +describe('permissions()', () => { + it('saves an always answer for the project, for the next call and a new session', async () => { + const persistence = memoryPersistence() + const tools = fakeTools() + const plugins = [permissions({ root: ROOT }), tools.plugin] + const first = await open(persistence, plugins, [ + () => toolCall('read_file', { path: '.env' }, 'c1'), + () => toolCall('read_file', { path: '.env' }, 'c2'), + ]) + + const turn = first.session.prompt('read it twice') + expect(await answer(first.session, { answer: 'always' })).toContain('.env') + await turn + + expect(tools.ran).toEqual(['read_file', 'read_file']) + expect( + await persistence.stores.metadata.get( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + ), + ).toEqual([{ tool: 'read_file', resource: '.env', decision: 'allow' }]) + await first.host.close() + + const second = await open( + persistence, + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('read_file', { path: '.env' }, 'c3'), + () => toolCall('read_file', { path: 'other.env' }, 'c4'), + ], + 't2', + ) + const next = second.session.prompt('read again') + // `.env` runs at once. `other.env` has no saved answer, so it asks. + expect(await answer(second.session, { answer: 'once' })).toContain( + 'other.env', + ) + await next + + expect(tools.ran).toHaveLength(4) + await second.host.close() + }) + + it('keeps an always answer for the session when there is no metadata store', async () => { + const { stores } = memoryPersistence() + const tools = fakeTools() + const { host, session } = await open( + { stores: { messages: stores.messages, runs: stores.runs } }, + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('read_file', { path: '../shared.txt' }, 'c1'), + () => toolCall('read_file', { path: '../shared.txt' }, 'c2'), + ], + ) + + const turn = session.prompt('read it twice') + await answer(session, { answer: 'always' }) + await turn + + expect(tools.ran).toEqual(['read_file', 'read_file']) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) + + it('gives the model a tool error with the message of a reject', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [() => toolCall('read_file', { path: '../secret.txt' }, 'c1')], + ) + + const turn = session.prompt('read outside') + await answer(session, { + answer: 'reject', + message: 'Stay inside the repo.', + }) + await turn + + expect(tools.ran).toEqual([]) + expect(JSON.stringify(calls[1].messages)).toContain( + '{\\"error\\":\\"Stay inside the repo.\\"}', + ) + await host.close() + }) + + it('asks for ls && rm x when only ls* is allowed', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [ + permissions({ + rules: [{ tool: 'bash', resource: 'ls*', decision: 'allow' }], + }), + tools.plugin, + ], + [ + () => toolCall('bash', { command: 'ls -la' }, 'c1'), + () => toolCall('bash', { command: 'ls && rm x' }, 'c2'), + ], + ) + + const turn = session.prompt('list, then clean') + expect(await answer(session, { answer: 'reject' })).toContain('rm x') + await turn + + expect(tools.ran).toEqual(['bash']) + expect(JSON.stringify(calls[2].messages)).toContain( + 'The user denied this tool call.', + ) + await host.close() + }) + + it('hides a tool that its last rule denies, and keeps one a later rule allows in part', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [ + permissions({ + rules: [ + { tool: 'deploy', decision: 'deny' }, + { tool: 'bash', decision: 'deny' }, + { tool: 'bash', resource: 'ls*', decision: 'allow' }, + ], + }), + tools.plugin, + ], + [], + ) + + await session.prompt('hi') + + const names = calls[0].tools.map((tool: { name: string }) => tool.name) + expect(names.sort()).toEqual(['bash', 'read_file']) + await host.close() + }) + + it('refuses a call whose resources cannot be read', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [() => toolCall('read_file', {}, 'c1')], + ) + + await session.prompt('read nothing') + + expect(tools.ran).toEqual([]) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect(JSON.stringify(calls[1].messages)).toContain( + 'Cannot check the permissions of this call: No path.', + ) + await host.close() + }) + + it('shares its rules with other plugins, and they still beat the rules of tool plugins', async () => { + const tools = fakeTools() + let seen: Array = [] + // Another plugin that reads the rules, like codeMode() does. + const reader = definePlugin({ + name: 'test/reader', + setup: (ctx) => { + const rules = ctx.collect(PermissionRules) + return { + prepareTools: ({ tools: list }) => { + seen = [...rules] + return list + }, + } + }, + }) + const rule: PermissionRule = { tool: 'read_file', decision: 'ask' } + const { host, session } = await open( + memoryPersistence(), + // The tool plugin comes after permissions(), and it allows read_file. + [permissions({ root: ROOT, rules: [rule] }), tools.plugin, reader], + [() => toolCall('read_file', { path: 'src/a.ts' }, 'c1')], + ) + + const turn = session.prompt('read a file') + expect(await answer(session, { answer: 'reject' })).toContain('src/a.ts') + await turn + + expect(tools.ran).toEqual([]) + expect(seen).toContainEqual(rule) + await host.close() + }) + + it('tells other plugins the decision a call gets, for a rule, the default, and plan mode', async () => { + const ran: Array = [] + const tool = (name: string) => + toolDefinition({ name, description: name }).server(async () => { + ran.push(name) + return `${name} done` + }) + let decide: + | ((tool: string, mode: PermissionMode) => PermissionDecision) + | undefined + // Another plugin that asks permissions(), like codeMode() does. + const reader = definePlugin({ + name: 'test/reader', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => { + decide = ctx.getOptional(PermissionDecisionCapability) + return { tools: [tool('lookup'), tool('deploy')] } + }, + }) + const { host, session, calls } = await open( + memoryPersistence(), + [ + permissions({ + default: 'ask', + rules: [{ tool: 'lookup', decision: 'allow' }], + }), + reader, + ], + [ + () => toolCall('lookup', {}, 'c1'), + () => toolCall('deploy', {}, 'c2'), + () => text('asked'), + () => toolCall('lookup', {}, 'c3'), + () => toolCall('deploy', {}, 'c4'), + ], + ) + + // The rule allows lookup, and deploy has no rule, so it asks. + const turn = session.prompt('look up, then deploy') + expect(await answer(session, { answer: 'reject' })).toContain('deploy') + await turn + expect(ran).toEqual(['lookup']) + expect(decide?.('lookup', 'default')).toBe('allow') + expect(decide?.('deploy', 'default')).toBe('ask') + + // Plan mode denies what would ask. + await session.command('mode', 'plan') + await session.prompt('again, in plan mode') + expect(ran).toEqual(['lookup', 'lookup']) + expect(JSON.stringify(calls.at(-1).messages)).toContain( + 'This tool is not allowed in plan mode.', + ) + expect(decide?.('lookup', 'plan')).toBe('allow') + expect(decide?.('deploy', 'plan')).toBe('deny') + await host.close() + }) + + it('tells other plugins the decision for a call that touches resources, from the root', async () => { + let decide: + | (( + tool: string, + mode: PermissionMode, + resources?: CallResources, + ) => PermissionDecision) + | undefined + const reader = definePlugin({ + name: 'test/reader', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => { + decide = ctx.getOptional(PermissionDecisionCapability) + return {} + }, + }) + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ] + const { host } = await open( + memoryPersistence(), + [permissions({ root: ROOT, rules }), reader], + [], + ) + const read = (mode: PermissionMode, ...paths: Array) => + decide?.('read_file', mode, { paths }) + + expect(read('default', `${ROOT}/src/a.ts`)).toBe('allow') + expect(read('default', `${ROOT}/.env`)).toBe('ask') + expect(read('plan', `${ROOT}/.env`)).toBe('deny') + expect(read('default', `${ROOT}/src/a.ts`, `${ROOT}/secrets/key`)).toBe( + 'deny', + ) + await host.close() + }) + + it('never lets a saved answer beat a deny', async () => { + const persistence = memoryPersistence() + await persistence.stores.metadata.set( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + [{ tool: 'read_file', resource: 'secret.txt', decision: 'allow' }], + ) + const tools = fakeTools() + const { host, session } = await open( + persistence, + [ + permissions({ + root: ROOT, + rules: [ + { tool: 'read_file', resource: 'secret.txt', decision: 'deny' }, + ], + }), + tools.plugin, + ], + [() => toolCall('read_file', { path: 'secret.txt' }, 'c1')], + ) + + await session.prompt('read the secret') + + expect(tools.ran).toEqual([]) + await host.close() + }) + + it('lists the saved rules of a project, and deletes one by its fields', async () => { + const persistence = memoryPersistence() + const tools = fakeTools() + const { host, session } = await open( + persistence, + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('read_file', { path: '.env' }, 'c1'), + () => toolCall('bash', { command: 'git status' }, 'c2'), + ], + ) + + const turn = session.prompt('read, then run') + await answer(session, { answer: 'always' }) + await answer(session, { answer: 'always' }) + await turn + await host.close() + + expect(await listSavedPermissions(persistence.stores, ROOT)).toEqual([ + { tool: 'read_file', resource: '.env', decision: 'allow' }, + { tool: 'bash', resource: 'git status', decision: 'allow' }, + ]) + expect(await listSavedPermissions(persistence.stores, '/other')).toEqual([]) + + await deleteSavedPermission(persistence.stores, ROOT, { + tool: 'bash', + resource: 'git *', + decision: 'allow', + }) + await deleteSavedPermission(persistence.stores, ROOT, { + tool: 'read_file', + resource: '.env', + decision: 'allow', + }) + expect(await listSavedPermissions(persistence.stores, ROOT)).toEqual([ + { tool: 'bash', resource: 'git status', decision: 'allow' }, + ]) + }) + + it('lists saved rules with /permissions, and forgets one with /permissions forget', async () => { + const persistence = memoryPersistence() + await persistence.stores.metadata.set( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + [ + { tool: 'bash', resource: 'git status', decision: 'allow' }, + { tool: 'deploy', decision: 'allow' }, + ], + ) + const tools = fakeTools() + const { host, session } = await open( + persistence, + [permissions({ root: ROOT }), tools.plugin], + [], + ) + + expect(await session.command('permissions')).toBe( + 'Saved rules:\n1. bash git status\n2. deploy', + ) + expect(await session.command('permissions', 'forget 3')).toBe( + 'No saved rule 3. Run /permissions to see the list.', + ) + expect(await session.command('permissions', 'forget 1')).toBe( + 'Forgot rule 1: bash git status.', + ) + expect(await listSavedPermissions(persistence.stores, ROOT)).toEqual([ + { tool: 'deploy', decision: 'allow' }, + ]) + await session.command('permissions', 'forget 1') + expect(await session.command('permissions')).toBe('No saved rules.') + await host.close() + }) + + it('asks again in the open session after /permissions forget', async () => { + const tools = fakeTools() + const { host, session } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('bash', { command: 'git status' }, 'c1'), + () => text('ran'), + () => toolCall('bash', { command: 'git status' }, 'c2'), + ], + ) + + const first = session.prompt('run it') + await answer(session, { answer: 'always' }) + await first + expect(await session.command('permissions', 'forget 1')).toBe( + 'Forgot rule 1: bash git status.', + ) + + const second = session.prompt('run it again') + expect(await answer(session, { answer: 'once' })).toContain('git status') + await second + expect(tools.ran).toEqual(['bash', 'bash']) + await host.close() + }) + + it('asks about an action that is not a tool call, like a formatter, by the mode', async () => { + const persistence = memoryPersistence() + let prompt: + | ((tool: string, message: string) => Promise) + | undefined + // A plugin like formatter(): it adds an ask rule, and asks before it runs. + const runner = definePlugin({ + name: 'test/runner', + optionalRequires: [PermissionPrompt], + setup: (ctx) => { + prompt = ctx.getOptional(PermissionPrompt) + return { + contribute: [ + PermissionRules.item({ + tool: 'formatter:*', + decision: 'ask', + kind: 'execute', + }), + ], + } + }, + }) + const { host, session } = await open( + persistence, + [ + permissions({ + root: ROOT, + rules: [{ tool: 'formatter:biome', decision: 'allow' }], + }), + runner, + ], + [], + ) + if (!prompt) throw new Error('No PermissionPrompt.') + + // A user rule allows biome. prettier asks, also in acceptEdits mode. + expect(await prompt('formatter:biome', 'Run biome?')).toBe(true) + await session.command('mode', 'acceptEdits') + const first = prompt('formatter:prettier', 'Run prettier?') + expect(await answer(session, { answer: 'reject' })).toBe( + 'Run prettier? Answer once, always, or reject.', + ) + expect(await first).toBe(false) + + // always saves the rule for the project, so it asks no more. + const second = prompt('formatter:prettier', 'Run prettier?') + await answer(session, { answer: 'always' }) + expect(await second).toBe(true) + expect(await prompt('formatter:prettier', 'Run prettier?')).toBe(true) + expect( + await persistence.stores.metadata.get( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + ), + ).toEqual([{ tool: 'formatter:prettier', decision: 'allow' }]) + + // bypass runs everything. plan runs no command. + await session.command('mode', 'bypass') + expect(await prompt('formatter:oxfmt', 'Run oxfmt?')).toBe(true) + await session.command('mode', 'plan') + expect(await prompt('formatter:prettier', 'Run prettier?')).toBe(false) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/question.test.ts b/packages/ai-harness/tests/question.test.ts new file mode 100644 index 0000000000..dafa3de4b5 --- /dev/null +++ b/packages/ai-harness/tests/question.test.ts @@ -0,0 +1,190 @@ +import { describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { EventType } from '@tanstack/ai' +import { + HARNESS_EVENTS, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { question } from '../src/first-party/question' +import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessSession } from '../src' + +const database = { + question: 'Which database?', + header: 'Database', + options: [ + { label: 'Postgres', description: 'Relational, the default' }, + { label: 'SQLite' }, + ], +} + +const features = { + question: 'Which features?', + options: [{ label: 'Auth' }, { label: 'Search' }, { label: 'Billing' }], + multiple: true, +} + +/** The model calls `question` with `questions`, then ends the turn. */ +async function open(questions: Array) { + const { adapter, calls } = mockAdapter([ + () => toolCall('question', { questions }, 'c1'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/question', + adapter, + plugins: () => [question()], + }), + { threadId: 't' }, + ) + const turn = session.prompt('ask me') + return { host, session, calls, turn } +} + +/** Wait for the one open question, then answer it. Returns the question. */ +async function reply(session: HarnessSession, value: unknown) { + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [pending] = session.snapshot().pendingQuestions + if (!pending) throw new Error('No question.') + await session.answer(pending.questionId, value) + return pending +} + +/** The tool result the model got on its second call. */ +function toolResult( + calls: Array<{ messages: Array<{ role: string; content: unknown }> }>, +) { + return calls[1]?.messages.find((message) => message.role === 'tool')?.content +} + +describe('question()', () => { + it('asks each question in turn and gives the model the picked labels', async () => { + const { host, session, calls, turn } = await open([database, features]) + + const first = await reply(session, '2') + const second = await reply(session, ['billing', 'Auth']) + await turn + + expect(first.message).toBe( + 'Database: Which database?\n1. Postgres - Relational, the default\n2. SQLite\nType a number or a label, or your own answer.', + ) + expect(first.schema).toEqual({ + type: 'string', + enum: ['Postgres', 'SQLite'], + }) + expect(second.message).toContain('3. Billing') + expect(second.schema).toEqual({ + type: 'array', + items: { type: 'string', enum: ['Auth', 'Search', 'Billing'] }, + minItems: 1, + uniqueItems: true, + }) + expect(toolResult(calls)).toBe( + 'Which database?\nAnswer: SQLite\n\nWhich features?\nAnswer: Billing, Auth', + ) + await host.close() + }) + + it.each([ + ['a label in another case', database, ' postgres ', 'Answer: Postgres'], + [ + 'numbers and labels split by commas', + features, + 'auth, 3', + 'Answer: Auth, Billing', + ], + [ + 'text that names no option', + database, + 'We use Mongo', + 'Own answer: We use Mongo', + ], + ])('reads %s', async (_name, asked, value, expected) => { + const { host, session, calls, turn } = await open([asked]) + + await reply(session, value) + await turn + + expect(toolResult(calls)).toBe(`${asked.question}\n${expected}`) + await host.close() + }) + + it('gives the model a tool error when the reply is empty', async () => { + const { host, session, calls, turn } = await open([database, features]) + + await reply(session, ' ') + await turn + + expect(toolResult(calls)).toContain('The user did not answer.') + // The second question is not asked. + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) +}) + +describe('ask with a url', () => { + it('sends the url in the event and the snapshot', async () => { + const asker = definePlugin({ + name: 'test/url-asker', + setup: (ctx) => ({ + commands: { + verify: defineCommand({ + description: 'Asks the user to open a page', + run: () => + ctx.session.ask({ + message: 'Open the page, then type done.', + url: 'https://example.com/verify', + }), + }), + }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/url-question', + adapter: mockAdapter([]).adapter, + plugins: () => [asker], + }), + { threadId: 't' }, + ) + const running = session.command('verify') + + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + expect(session.snapshot().pendingQuestions[0]?.url).toBe( + 'https://example.com/verify', + ) + const pending = await reply(session, 'done') + expect(pending.url).toBe('https://example.com/verify') + expect(await running).toBe('done') + + const controller = new AbortController() + let url: unknown + for await (const { event } of session.events({ + from: '0', + signal: controller.signal, + })) { + if ( + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.question && + event.value !== null && + typeof event.value === 'object' && + 'url' in event.value + ) { + url = event.value.url + controller.abort() + } + } + expect(url).toBe('https://example.com/verify') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/session-plugins.test.ts b/packages/ai-harness/tests/session-plugins.test.ts index 3f4f6f6e26..d0c56f6112 100644 --- a/packages/ai-harness/tests/session-plugins.test.ts +++ b/packages/ai-harness/tests/session-plugins.test.ts @@ -1,4 +1,7 @@ -import { describe, expect, it, vi } from 'vitest' +import { mkdtemp, rm, writeFile, mkdir } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { EventType, toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' @@ -11,7 +14,15 @@ import { defineHarness, definePlugin, } from '../src' -import { gate, mockAdapter, text } from './helpers' +import { + compact, + fileCommands, + modelPicker, + projectInstructions, + todos, + usage, +} from '../src/first-party' +import { gate, mockAdapter, text, toolCall } from './helpers' import type { SessionEvent } from '../src' const settings = definePlugin({ @@ -374,6 +385,181 @@ describe('the plugin session API and inspect', () => { }) }) +describe('first-party plugin edges', () => { + let dir: string + beforeAll(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-files-')) + await mkdir(join(dir, 'commands')) + await writeFile( + join(dir, 'commands', 'review.md'), + '---\ndescription: Review a file\n---\nReview $ARGUMENTS carefully.', + ) + await writeFile(join(dir, 'commands', 'plain.md'), 'Say hi to $ARGUMENTS.') + await writeFile(join(dir, 'commands', 'skip.txt'), 'not a command') + await writeFile(join(dir, 'RULES.md'), 'Use tabs.') + await writeFile(join(dir, 'EMPTY.md'), ' ') + }) + afterAll(async () => { + await rm(dir, { recursive: true, force: true }) + }) + + it('turns markdown files into commands and reads chosen instruction files', async () => { + const { adapter, calls } = mockAdapter([ + () => text('a'), + () => text('b'), + () => text('c'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/files', + adapter, + plugins: () => [ + fileCommands({ dir: join(dir, 'commands') }), + projectInstructions({ + root: dir, + files: ['RULES.md', 'EMPTY.md', 'NONE.md'], + }), + ], + }), + { threadId: 't' }, + ) + const described = session.commands() + expect( + described.map((command) => [command.name, command.description]).sort(), + ).toEqual([ + ['plain', 'Run plain.md'], + ['review', 'Review a file'], + ]) + expect(await session.command('review', 'src/a.ts')).toBe('Sent /review.') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await session.command('plain', { who: 'Ada' }) + await vi.waitFor(() => expect(calls).toHaveLength(2)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await session.command('plain') + await vi.waitFor(() => expect(calls).toHaveLength(3)) + + // A missing folder gives no commands. + const empty = await host.open( + defineHarness({ + name: 'test/files-missing', + adapter, + plugins: () => [fileCommands({ dir: join(dir, 'missing') })], + }), + { threadId: 'empty' }, + ) + expect(empty.commands()).toEqual([]) + + const said = calls.map((call) => JSON.stringify(call.messages.at(-1))) + expect(said[0]).toContain('Review src/a.ts carefully.') + expect(said[1]).toContain('Say hi to {\\"who\\":\\"Ada\\"}.') + expect(said[2]).toContain('Say hi to .') + const prompts = JSON.stringify(calls[0].systemPrompts) + expect(prompts).toContain('Project instructions from RULES.md') + expect(prompts).not.toContain('EMPTY.md') + await host.close() + }) + + it('checks todo input and shows an empty list', async () => { + const { adapter, calls } = mockAdapter([ + () => toolCall('todo_write', { todos: [{ text: 'x', status: 'later' }] }), + () => text('bad input'), + () => toolCall('todo_write', { todos: [] }, 'call-2'), + () => text('cleared'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ name: 'test/todos', adapter, plugins: () => [todos()] }), + { threadId: 't' }, + ) + const before = session.snapshot().plugins + await session.prompt('plan') + // The tool has a plain JSON Schema, which chat() does not check, so the + // tool checks its input itself. + expect(JSON.stringify(calls[1].messages)).toContain( + 'Each todo needs text and a status.', + ) + expect(session.snapshot().plugins).toEqual(before) + expect(await session.command('todos')).toBe('The todo list is empty.') + await session.prompt('clear') + expect(JSON.stringify(calls[3].messages)).toContain( + 'The todo list is empty.', + ) + expect(await session.command('todos')).toBe('The todo list is empty.') + await host.close() + }) + + it('needs a model choice and answers every /model form', async () => { + expect(() => modelPicker({ choices: {} })).toThrow('at least one choice') + const fast = mockAdapter([]).adapter + const smart = mockAdapter([]).adapter + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/models', + adapter: smart, + plugins: () => [modelPicker({ choices: { smart, fast } })], + }), + { threadId: 't' }, + ) + expect(await session.command('model')).toBe( + 'Model: smart. Choices: smart, fast.', + ) + expect(await session.command('model', 'slow')).toBe( + 'Unknown model "slow". Choices: smart, fast.', + ) + expect(await session.command('model', { name: 'fast' })).toBe( + 'Model: fast. It applies at the next turn.', + ) + await host.close() + }) + + it('keeps short conversations, keeps the last messages, and counts usage', async () => { + const summarizer = mockAdapter([() => text('the summary')]) + const withUsage = (content: string) => + text(content).map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 3, completionTokens: 2, totalTokens: 5 }, + } + : chunk, + ) + const { adapter } = mockAdapter([ + () => withUsage('one'), + () => withUsage('two'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/compact', + adapter, + plugins: () => [ + compact({ adapter: summarizer.adapter, keepLast: 1 }), + usage(), + ], + }), + { threadId: 't' }, + ) + expect(await session.command('compact')).toBe( + 'The conversation is already short.', + ) + await session.prompt([{ type: 'text', content: 'first question' }] as never) + await session.prompt('second question') + expect(await session.command('compact')).toBe( + 'Compacted 3 messages into a summary.', + ) + expect(JSON.stringify(summarizer.calls[0].messages)).toContain( + 'first question', + ) + expect(await session.command('usage')).toBe( + '2 model calls, 6 input tokens, 4 output tokens, 10 total.', + ) + await host.close() + }) +}) + describe('session events', () => { it('publishes config changes', async () => { const host = createHarnessHost({ persistence: memoryPersistence() }) diff --git a/packages/ai-harness/tests/title-usage.test.ts b/packages/ai-harness/tests/title-usage.test.ts new file mode 100644 index 0000000000..9e2f6ca709 --- /dev/null +++ b/packages/ai-harness/tests/title-usage.test.ts @@ -0,0 +1,310 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { usage } from '../src/first-party/session-tools' +import { TitleFailed, title } from '../src/first-party/title' +import { messageTexts, mockAdapter, text } from './helpers' +import type { StreamChunk, TokenUsage } from '@tanstack/ai' +import type { AnyAgent, HarnessPersistence, HarnessPlugin } from '../src' +import type { Reply } from './helpers' + +/** + * Open thread `t`. The main model answers `answer` unless `reply` is given. + * `agents` are the background agents of the session. + */ +async function open( + plugins: Array, + { + persistence = memoryPersistence(), + reply = () => text('answer'), + agents = [], + }: { + persistence?: HarnessPersistence + reply?: Reply + agents?: Array + } = {}, +) { + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/title-usage', + adapter: mockAdapter(reply).adapter, + agents, + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +describe('title()', () => { + it('writes a title from the first user message after the first turn', async () => { + const titler = mockAdapter(() => text('"Trip to Rome"\nA second line')) + const { host, session } = await open([title({ adapter: titler.adapter })]) + + const turn = await session.prompt('Plan a trip to Rome for me') + + expect(turn.text).toBe('answer') + await vi.waitFor(async () => + expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome'), + ) + expect(messageTexts(titler.calls[0])[0]).toContain( + 'Plan a trip to Rome for me', + ) + await host.close() + }) + + it('does not ask for a title again once the session has one', async () => { + const titler = mockAdapter([ + () => text('Trip to Rome'), + () => text('Another title'), + ]) + const { host, session } = await open([title({ adapter: titler.adapter })]) + await session.prompt('Plan a trip to Rome for me') + await vi.waitFor(async () => + expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome'), + ) + + await session.prompt('Add a day in Florence') + + expect(titler.calls).toHaveLength(1) + expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome') + await host.close() + }) + + it('keeps the turn when the title call fails', async () => { + const failures: Array<{ message: string }> = [] + const listener = definePlugin({ + name: 'test/listener', + setup: (ctx) => { + ctx.on(TitleFailed, (failure) => failures.push(failure)) + }, + }) + const titler = mockAdapter(() => { + throw new Error('title model down') + }) + const { host, session } = await open([ + listener, + title({ adapter: titler.adapter }), + ]) + + const turn = await session.prompt('Plan a trip to Rome for me') + + expect(turn.text).toBe('answer') + await vi.waitFor(() => + expect(failures).toEqual([{ message: 'title model down' }]), + ) + expect((await host.sessions.get('t'))?.title).toBeUndefined() + await host.close() + }) + + it('does nothing without a sessions store', async () => { + const { messages, runs, metadata } = memoryPersistence().stores + const titler = mockAdapter(() => text('Trip to Rome')) + const { host, session } = await open([title({ adapter: titler.adapter })], { + persistence: { stores: { messages, runs, metadata } }, + }) + + await session.prompt('Plan a trip to Rome for me') + + expect(titler.calls).toHaveLength(0) + await host.close() + }) +}) + +describe('usage()', () => { + /** Prices in USD per 1M tokens. */ + const PRICES = { input: 2, output: 10, cacheRead: 1, cacheWrite: 4 } + /** + * One model call. With PRICES it costs $3.25: 500k uncached input ($1), + * 100k output ($1), 250k cache read ($0.25), 250k cache write ($1). + */ + const CALL: TokenUsage = { + promptTokens: 1_000_000, + completionTokens: 100_000, + totalTokens: 1_100_000, + promptTokensDetails: { cachedTokens: 250_000, cacheWriteTokens: 250_000 }, + } + /** The `/usage` line of one CALL. */ + const CALL_LINE = + '1 model calls, 1000000 input tokens (250000 cache read, 250000 cache write), 100000 output tokens, 1100000 total.' + /** A model call that answers and reports `usage`. */ + const replyWith = (usage: TokenUsage) => () => + text('answer').map( + (chunk): StreamChunk => + chunk.type === EventType.RUN_FINISHED ? { ...chunk, usage } : chunk, + ) + const withUsage = replyWith(CALL) + /** The prices of the mock adapter's model only. */ + const pricesOf = (modelId: string) => + modelId === 'test-model' ? { cost: PRICES } : undefined + /** A background agent. Its model, `worker-model`, reports `usage`. */ + const workerWith = (usage: TokenUsage) => { + const { adapter } = mockAdapter(replyWith(usage)) + return defineAgent({ + name: 'worker', + description: 'Works', + run: (ctx) => + ctx.chat({ adapter: { ...adapter, model: 'worker-model' } }), + }) + } + + it('uses a provider cost as is, and prices only the calls with no provider cost', async () => { + // The lookup knows every model: it would price the worker call at $3.25. + const { host, session } = await open( + [usage({ model: () => ({ cost: PRICES }) })], + { reply: withUsage, agents: [workerWith({ ...CALL, cost: 0.5 })] }, + ) + + await session.prompt('one') + await session.agent('worker')?.start() + + // The lead call costs $3.25 from the lookup, the worker call $0.50 from + // its provider. + expect(await session.command('usage')).toBe( + [ + '2 model calls, 2000000 input tokens (500000 cache read, 500000 cache write), 200000 output tokens, 2200000 total. Cost: $3.7500.', + `mock/test-model: ${CALL_LINE}`, + `mock/worker-model: ${CALL_LINE}`, + ].join('\n'), + ) + // `/usage` shows the totals of the session. + expect(session.usage().total).toEqual({ + calls: 2, + promptTokens: 2_000_000, + completionTokens: 200_000, + totalTokens: 2_200_000, + cachedTokens: 500_000, + cacheWriteTokens: 500_000, + cost: 0.5, + }) + expect((await host.sessions.get('t'))?.usage?.cost).toBeCloseTo(3.75, 10) + await host.close() + }) + + it('shows a provider cost without the model option, and writes it to the index entry', async () => { + const { host, session } = await open([usage()], { + reply: replyWith({ ...CALL, cost: 0.5 }), + }) + + await session.prompt('one') + + expect(await session.command('usage')).toBe(`${CALL_LINE} Cost: $0.5000.`) + expect((await host.sessions.get('t'))?.usage?.cost).toBe(0.5) + await host.close() + }) + + it('counts a provider cost once in the index entry', async () => { + const { host, session } = await open([usage({ model: pricesOf })], { + reply: replyWith({ ...CALL, cost: 0.5 }), + }) + + await session.prompt('one') + await session.prompt('two') + + // Two calls at $0.50. The host writes no cost of its own. + expect((await host.sessions.get('t'))?.usage?.cost).toBeCloseTo(1, 10) + await host.close() + }) + + it('writes the session totals to the index entry, with the calls of a background agent', async () => { + const { host, session } = await open([usage()], { + reply: withUsage, + agents: [workerWith(CALL)], + }) + + await session.prompt('one') + await session.agent('worker')?.start() + await session.prompt('two') + + // Two lead calls and one worker call. + expect((await host.sessions.get('t'))?.usage).toEqual({ + turns: 3, + promptTokens: 3_000_000, + completionTokens: 300_000, + totalTokens: 3_300_000, + cachedTokens: 750_000, + cacheWriteTokens: 750_000, + }) + await host.close() + }) + + it('keeps a copy of the session totals in its state', async () => { + const { host, session } = await open([usage()], { + reply: withUsage, + agents: [workerWith({ ...CALL, promptTokens: 2_000_000 })], + }) + + await session.prompt('one') + await session.agent('worker')?.start() + + // The worker's call is in the copy too. + expect(session.snapshot().plugins['tanstack/usage']).toEqual({ + calls: 2, + promptTokens: 3_000_000, + completionTokens: 200_000, + totalTokens: 2_200_000, + cachedTokens: 500_000, + cacheWriteTokens: 500_000, + // The input of the lead's latest call, not the worker's. + contextTokens: 1_000_000, + }) + await host.close() + }) + + it('prices each model call and shows the cost in /usage', async () => { + const { host, session } = await open([usage({ model: pricesOf })], { + reply: withUsage, + }) + + await session.prompt('one') + await session.prompt('two') + + expect(await session.command('usage')).toBe( + '2 model calls, 2000000 input tokens (500000 cache read, 500000 cache write), 200000 output tokens, 2200000 total. Cost: $6.5000.', + ) + await host.close() + }) + + it('writes the cost to the index entry and keeps the token totals of the host', async () => { + const { host, session } = await open([usage({ model: pricesOf })], { + reply: withUsage, + }) + + await session.prompt('one') + await session.prompt('two') + + expect((await host.sessions.get('t'))?.usage).toEqual({ + turns: 2, + promptTokens: 2_000_000, + completionTokens: 200_000, + totalTokens: 2_200_000, + cachedTokens: 500_000, + cacheWriteTokens: 500_000, + cost: expect.closeTo(6.5, 10), + }) + await host.close() + }) + + it('adds no cost for a model with no prices, and says so in /usage', async () => { + const { host, session } = await open([usage({ model: () => undefined })], { + reply: withUsage, + }) + + await session.prompt('one') + + expect(await session.command('usage')).toBe( + `${CALL_LINE} Cost: $0.0000 (cost unknown for 1 call).`, + ) + expect((await host.sessions.get('t'))?.usage).toEqual({ + turns: 1, + promptTokens: 1_000_000, + completionTokens: 100_000, + totalTokens: 1_100_000, + cachedTokens: 250_000, + cacheWriteTokens: 250_000, + }) + await host.close() + }) +}) diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index 1f3542380f..86b2a7a465 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,7 +28,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/first-party/index.ts'], srcDir: './src', cjs: false, }), diff --git a/testing/e2e/fixtures/harness-first-party/basic.json b/testing/e2e/fixtures/harness-first-party/basic.json new file mode 100644 index 0000000000..77195c081f --- /dev/null +++ b/testing/e2e/fixtures/harness-first-party/basic.json @@ -0,0 +1,107 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[first-party-plan] write x.txt", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "writer", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[first-party-plan] write x.txt", + "sequenceIndex": 1 + }, + "response": { + "content": "The writer could not write in plan mode." + } + }, + { + "match": { + "userMessage": "[first-party-plan-writer] write x.txt", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "write_file", + "arguments": "{\"path\":\"x.txt\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[first-party-plan-writer] write x.txt", + "sequenceIndex": 1 + }, + "response": { + "content": "Plan mode blocked the write." + } + }, + { + "match": { + "userMessage": "[first-party-agents] first" + }, + "response": { + "content": "Built it." + } + }, + { + "match": { + "userMessage": "[first-party-agents] second", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "lookup", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[first-party-agents] second", + "sequenceIndex": 1 + }, + "response": { + "content": "Stopped at the step limit." + } + }, + { + "match": { + "userMessage": "[first-party-env] where am i" + }, + "response": { + "content": "You are in a new repository." + } + }, + { + "match": { + "userMessage": "[first-party-title] plan a trip to Rome", + "model": "gpt-5.4-nano" + }, + "response": { + "content": "\"Trip to Rome\"" + } + }, + { + "match": { + "userMessage": "[first-party-title] plan a trip to Rome", + "model": "gpt-5.5" + }, + "response": { + "content": "Here is a plan for Rome." + } + } + ] +} diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index d7fe95be71..e604264202 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -86,6 +86,7 @@ import { Route as ApiEmbeddingRouteImport } from './routes/api.embedding' import { Route as ApiFakeTextRouteImport } from './routes/api.fake-text' import { Route as ApiHarnessContinueRouteImport } from './routes/api.harness-continue' import { Route as ApiHarnessRetryAfterRouteImport } from './routes/api.harness-retry-after' +import { Route as ApiHarnessFirstPartyRouteImport } from './routes/api.harness-first-party' import { Route as ApiFileSourceWireRouteImport } from './routes/api.file-source-wire' import { Route as ApiForeignChunkEventsRouteImport } from './routes/api.foreign-chunk-events' import { Route as ApiForeignClientToolRouteImport } from './routes/api.foreign-client-tool' @@ -582,6 +583,11 @@ const ApiHarnessRetryAfterRoute = ApiHarnessRetryAfterRouteImport.update({ path: '/api/harness-retry-after', getParentRoute: () => rootRouteImport, } as any) +const ApiHarnessFirstPartyRoute = ApiHarnessFirstPartyRouteImport.update({ + id: '/api/harness-first-party', + path: '/api/harness-first-party', + getParentRoute: () => rootRouteImport, +} as any) const ApiFileSourceWireRoute = ApiFileSourceWireRouteImport.update({ id: '/api/file-source-wire', path: '/api/file-source-wire', @@ -1149,6 +1155,7 @@ export interface FileRoutesByFullPath { '/api/fake-text': typeof ApiFakeTextRoute '/api/harness-continue': typeof ApiHarnessContinueRoute '/api/harness-retry-after': typeof ApiHarnessRetryAfterRoute + '/api/harness-first-party': typeof ApiHarnessFirstPartyRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute '/api/foreign-client-tool': typeof ApiForeignClientToolRoute @@ -1321,6 +1328,7 @@ export interface FileRoutesByTo { '/api/fake-text': typeof ApiFakeTextRoute '/api/harness-continue': typeof ApiHarnessContinueRoute '/api/harness-retry-after': typeof ApiHarnessRetryAfterRoute + '/api/harness-first-party': typeof ApiHarnessFirstPartyRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute '/api/foreign-client-tool': typeof ApiForeignClientToolRoute @@ -1494,6 +1502,7 @@ export interface FileRoutesById { '/api/fake-text': typeof ApiFakeTextRoute '/api/harness-continue': typeof ApiHarnessContinueRoute '/api/harness-retry-after': typeof ApiHarnessRetryAfterRoute + '/api/harness-first-party': typeof ApiHarnessFirstPartyRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute '/api/foreign-client-tool': typeof ApiForeignClientToolRoute @@ -1668,6 +1677,7 @@ export interface FileRouteTypes { | '/api/fake-text' | '/api/harness-continue' | '/api/harness-retry-after' + | '/api/harness-first-party' | '/api/file-source-wire' | '/api/foreign-chunk-events' | '/api/foreign-client-tool' @@ -1840,6 +1850,7 @@ export interface FileRouteTypes { | '/api/fake-text' | '/api/harness-continue' | '/api/harness-retry-after' + | '/api/harness-first-party' | '/api/file-source-wire' | '/api/foreign-chunk-events' | '/api/foreign-client-tool' @@ -2012,6 +2023,7 @@ export interface FileRouteTypes { | '/api/fake-text' | '/api/harness-continue' | '/api/harness-retry-after' + | '/api/harness-first-party' | '/api/file-source-wire' | '/api/foreign-chunk-events' | '/api/foreign-client-tool' @@ -2185,6 +2197,7 @@ export interface RootRouteChildren { ApiFakeTextRoute: typeof ApiFakeTextRoute ApiHarnessContinueRoute: typeof ApiHarnessContinueRoute ApiHarnessRetryAfterRoute: typeof ApiHarnessRetryAfterRoute + ApiHarnessFirstPartyRoute: typeof ApiHarnessFirstPartyRoute ApiFileSourceWireRoute: typeof ApiFileSourceWireRoute ApiForeignChunkEventsRoute: typeof ApiForeignChunkEventsRoute ApiForeignClientToolRoute: typeof ApiForeignClientToolRoute @@ -2817,6 +2830,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiHarnessRetryAfterRouteImport parentRoute: typeof rootRouteImport } + '/api/harness-first-party': { + id: '/api/harness-first-party' + path: '/api/harness-first-party' + fullPath: '/api/harness-first-party' + preLoaderRoute: typeof ApiHarnessFirstPartyRouteImport + parentRoute: typeof rootRouteImport + } '/api/file-source-wire': { id: '/api/file-source-wire' path: '/api/file-source-wire' @@ -3611,6 +3631,7 @@ const rootRouteChildren: RootRouteChildren = { ApiFakeTextRoute: ApiFakeTextRoute, ApiHarnessContinueRoute: ApiHarnessContinueRoute, ApiHarnessRetryAfterRoute: ApiHarnessRetryAfterRoute, + ApiHarnessFirstPartyRoute: ApiHarnessFirstPartyRoute, ApiFileSourceWireRoute: ApiFileSourceWireRoute, ApiForeignChunkEventsRoute: ApiForeignChunkEventsRoute, ApiForeignClientToolRoute: ApiForeignClientToolRoute, diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts new file mode 100644 index 0000000000..8dc25c4ef9 --- /dev/null +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -0,0 +1,245 @@ +import { execFile } from 'node:child_process' +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { promisify } from 'node:util' +import { createFileRoute } from '@tanstack/react-router' +import { defineAgent, toolDefinition } from '@tanstack/ai' +import { + createHarnessHost, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' +import { + TitleFailed, + agents, + permissions, + projectInstructions, + title, +} from '@tanstack/ai-harness/plugins' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessHost } from '@tanstack/ai-harness' +import { createTextAdapter } from '@/lib/providers' + +const run = promisify(execFile) + +/** + * `agents()`: the first turn runs with the default `build` agent. `/agent + * brief` switches to an agent with `steps: 1`. Its turn calls `lookup` once, + * then the step limit makes the last call without tools. The spec reads the + * model calls in aimock's journal. + */ +async function agentProfiles(host: HarnessHost, model: () => AnyTextAdapter) { + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({}), + }).server(async () => 'A fact.') + const session = await host.open( + defineHarness({ + name: 'e2e/harness-agents', + adapter: model(), + tools: [lookup], + plugins: () => [ + agents({ + adapter: () => model(), + agents: [ + { + name: 'brief', + description: 'Answers after one tool step', + mode: 'primary', + system: 'You are the brief agent.', + steps: 1, + }, + ], + }), + ], + }), + { threadId: 'e2e-agents' }, + ) + const first = (await session.prompt('[first-party-agents] first')).text + const switched = await session.command('agent', 'brief') + const second = (await session.prompt('[first-party-agents] second')).text + return { first, switched, second } +} + +/** + * `permissions()` in `plan` mode: the main model starts the `writer` + * subagent, and the subagent's `write_file` call does not run. + */ +async function planSubagent(host: HarnessHost, model: () => AnyTextAdapter) { + let writes = 0 + const results: Array = [] + // A fake write_file, so a failing test changes no files. + const writeFile = toolDefinition({ + name: 'write_file', + description: 'Write a file', + inputSchema: z.object({ path: z.string() }), + }).server(async () => { + writes += 1 + return { written: true } + }) + const writer = defineAgent({ + name: 'writer', + description: 'Writes files', + run: (ctx) => + ctx.chat({ + adapter: model(), + messages: [ + { role: 'user', content: '[first-party-plan-writer] write x.txt' }, + ], + tools: [writeFile], + middleware: [ + { + name: 'e2e/tool-results', + onAfterToolCall: (_ctx, info) => { + results.push(info.result) + }, + }, + ], + stream: false, + }), + }) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-plan', + adapter: model(), + subagents: { agents: [writer] }, + plugins: () => [ + permissions({ + rules: [{ tool: 'write_file', decision: 'ask', kind: 'edit' }], + }), + ], + }), + { threadId: 'e2e-plan' }, + ) + await session.command('mode', 'plan') + const turn = await session.prompt('[first-party-plan] write x.txt') + return { writes, results, text: turn.text } +} + +/** + * `projectInstructions()` in a new git repository on the branch `e2e-env`. + * Returns the folder and the platform, so the spec can find the environment + * block in the system prompt of the model call. + */ +async function environmentPrompt( + host: HarnessHost, + model: () => AnyTextAdapter, +) { + const root = await mkdtemp(join(tmpdir(), 'e2e-env-')) + try { + // The branch line needs a commit. + await run('git', ['init', '-q', '-b', 'e2e-env'], { cwd: root }) + await run( + 'git', + [ + '-c', + 'user.name=e2e', + '-c', + 'user.email=e2e@example.com', + 'commit', + '-q', + '--allow-empty', + '-m', + 'init', + ], + { cwd: root }, + ) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-env', + adapter: model(), + plugins: () => [projectInstructions({ root })], + }), + { threadId: 'e2e-env' }, + ) + const turn = await session.prompt('[first-party-env] where am i') + return { root, platform: process.platform, text: turn.text } + } finally { + await rm(root, { recursive: true, force: true }) + } +} + +/** `title()`: the first turn names the session in the session index. */ +async function sessionTitle( + host: HarnessHost, + model: () => AnyTextAdapter, + titleModel: AnyTextAdapter, +) { + const failures: Array = [] + const listener = definePlugin({ + name: 'e2e/title-failures', + setup: (ctx) => { + ctx.on(TitleFailed, (failure) => failures.push(failure.message)) + }, + }) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-title', + adapter: model(), + plugins: () => [listener, title({ adapter: titleModel })], + }), + { threadId: 'e2e-title' }, + ) + const turn = await session.prompt('[first-party-title] plan a trip to Rome') + const deadline = Date.now() + 5_000 + let entry = await host.sessions.get('e2e-title') + while (!entry?.title && failures.length === 0 && Date.now() < deadline) { + await new Promise((resolve) => setTimeout(resolve, 50)) + entry = await host.sessions.get('e2e-title') + } + return { text: turn.text, title: entry?.title, failures } +} + +/** + * First-party harness plugins, one scenario per request: + * `?scenario=agents|plan-subagent|env|title&testId=...`. The models are + * OpenAI adapters against aimock. + */ +export const Route = createFileRoute('/api/harness-first-party')({ + server: { + handlers: { + POST: async ({ request }) => { + const params = new URL(request.url).searchParams + const scenario = params.get('scenario') + const testId = params.get('testId') ?? undefined + const model = () => + createTextAdapter('openai', 'gpt-5.5', undefined, testId).adapter + const host = createHarnessHost({ persistence: memoryPersistence() }) + try { + if (scenario === 'agents') { + return Response.json(await agentProfiles(host, model)) + } + if (scenario === 'plan-subagent') { + return Response.json(await planSubagent(host, model)) + } + if (scenario === 'env') { + return Response.json(await environmentPrompt(host, model)) + } + if (scenario === 'title') { + const titleModel = createTextAdapter( + 'openai', + 'gpt-5.4-nano', + undefined, + testId, + ).adapter + return Response.json(await sessionTitle(host, model, titleModel)) + } + return Response.json( + { error: `Unknown scenario ${scenario}` }, + { status: 400 }, + ) + } catch (error) { + return Response.json({ + error: error instanceof Error ? error.message : String(error), + }) + } finally { + await host.close() + } + }, + }, + }, +}) diff --git a/testing/e2e/tests/harness-first-party.spec.ts b/testing/e2e/tests/harness-first-party.spec.ts new file mode 100644 index 0000000000..058cc9419d --- /dev/null +++ b/testing/e2e/tests/harness-first-party.spec.ts @@ -0,0 +1,125 @@ +import { test, expect } from './fixtures' +import type { APIRequestContext } from '@playwright/test' + +/** + * First-party harness plugins in a session: `agents()`, `permissions()`, + * `projectInstructions()`, and `title()`. 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. + */ + +type JournalEntry = { + headers?: Record + body: { + messages?: Array<{ role: string; content: unknown }> + tool_choice?: unknown + } | null +} + +/** This test's model calls, in order, from aimock's journal. */ +async function modelCalls( + request: APIRequestContext, + aimockPort: number, + testId: string, +) { + const journal = await request.get( + `http://127.0.0.1:${aimockPort}/v1/_requests`, + ) + const entries: Array = await journal.json() + return entries.filter((entry) => + Object.entries(entry.headers ?? {}).some( + ([key, value]) => key.toLowerCase() === 'x-test-id' && value === testId, + ), + ) +} + +/** The system prompt text of one model call. */ +const systemOf = (call: JournalEntry | undefined) => + (call?.body?.messages ?? []) + .filter((message) => message.role === 'system') + .map((message) => + typeof message.content === 'string' ? message.content : '', + ) + .join('\n') + +/** Run one scenario of the route. */ +async function runScenario( + request: APIRequestContext, + scenario: string, + testId: string, +) { + const response = await request.post( + `/api/harness-first-party?scenario=${scenario}&testId=${encodeURIComponent(testId)}`, + ) + expect(response.ok()).toBe(true) + return response.json() +} + +test.describe('harness first-party plugins', () => { + test('agents(): /agent switches the profile, and the step limit makes a last call without tools', async ({ + request, + testId, + aimockPort, + }) => { + expect(await runScenario(request, 'agents', testId)).toEqual({ + first: 'Built it.', + switched: 'Agent: brief. It applies at the next turn.', + second: 'Stopped at the step limit.', + }) + + const calls = await modelCalls(request, aimockPort, testId) + expect(calls.map((call) => call.body?.tool_choice)).toEqual([ + undefined, + undefined, + 'none', + ]) + expect(systemOf(calls[0])).toContain('You are the build agent.') + expect(systemOf(calls[1])).toContain('You are the brief agent.') + expect(systemOf(calls[1])).not.toContain('You are the build agent.') + expect(systemOf(calls[1])).not.toContain('step limit') + expect(systemOf(calls[2])).toContain('You reached the step limit.') + }) + + test('permissions(): plan mode stops a write in a subagent', async ({ + request, + testId, + }) => { + expect(await runScenario(request, 'plan-subagent', testId)).toEqual({ + writes: 0, + results: [{ error: 'This tool is not allowed in plan mode.' }], + text: 'The writer could not write in plan mode.', + }) + }) + + test('projectInstructions(): the environment block reaches the model', async ({ + request, + testId, + aimockPort, + }) => { + const body = await runScenario(request, 'env', testId) + expect(body.text).toBe('You are in a new repository.') + + const [call] = await modelCalls(request, aimockPort, testId) + const system = systemOf(call) + expect(system).toMatch(/Environment:\n- Date: \d{4}-\d{2}-\d{2}\n/) + expect(system).toContain( + [ + `- Platform: ${body.platform}`, + `- Working folder: ${body.root}`, + '- Git repository: yes', + '- Git branch: e2e-env', + ].join('\n'), + ) + }) + + test('title(): the first turn names the session in the session index', async ({ + request, + testId, + }) => { + expect(await runScenario(request, 'title', testId)).toEqual({ + text: 'Here is a plan for Rome.', + title: 'Trip to Rome', + failures: [], + }) + }) +}) From 84c615396cfc7f6c69aa1160e790aa0675b2a005 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:08:25 +0200 Subject: [PATCH 2/9] docs(ai-harness): say that a tool with no rule runs in plan mode --- docs/harness/permissions.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 15b2a9aa97..9062b3f9bf 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -330,7 +330,7 @@ export const planTools = definePlugin({ ## Known limits -- A rule that allows a tool with no `kind` also lets it run in `plan` mode. `plan` denies only edits, commands, and calls that ask. +- `plan` denies only edits, commands, and calls that ask. A tool with no rule, or with an `allow` rule and no `kind`, still runs in `plan` mode. Give each tool that changes something a rule with a `kind`. - A client that can send inputs to the thread can answer the questions, also with `always`. Control who can open a thread. - `permissions()` does not follow links. Your own tool can check the real path of a link with `PermissionDecisionCapability`. From 3188a70c47933604dd34b527faf74c0531f1b084 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:37:27 +0200 Subject: [PATCH 3/9] feat(ai-code-mode): codeMode() harness plugin in @tanstack/ai-code-mode/harness --- .changeset/code-mode-harness.md | 5 + docs/code-mode/code-mode.md | 1 + docs/config.json | 8 +- docs/harness/code-mode.md | 135 +++++ docs/harness/permissions.md | 1 + packages/ai-code-mode/package.json | 12 + packages/ai-code-mode/src/harness.ts | 152 ++++++ packages/ai-code-mode/tests/harness.test.ts | 482 ++++++++++++++++++ packages/ai-code-mode/vite.config.ts | 2 +- pnpm-lock.yaml | 9 + .../fixtures/harness-first-party/basic.json | 8 + testing/e2e/package.json | 1 + .../e2e/src/routes/api.harness-first-party.ts | 38 +- testing/e2e/tests/harness-first-party.spec.ts | 19 +- 14 files changed, 869 insertions(+), 4 deletions(-) create mode 100644 .changeset/code-mode-harness.md create mode 100644 docs/harness/code-mode.md create mode 100644 packages/ai-code-mode/src/harness.ts create mode 100644 packages/ai-code-mode/tests/harness.test.ts 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/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 df3ace69e3..7fe4b72259 100644 --- a/docs/config.json +++ b/docs/config.json @@ -993,13 +993,19 @@ { "label": "Ask before risky tool calls", "to": "harness/permissions", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" }, { "label": "Switch between agent profiles", "to": "harness/agents", "addedAt": "2026-10-13" }, + { + "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..30c03dabdb --- /dev/null +++ b/docs/harness/code-mode.md @@ -0,0 +1,135 @@ +--- +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 | +| `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 | + +[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`. + +## 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 9062b3f9bf..dfdff21b2c 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -326,6 +326,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/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/harness.test.ts b/packages/ai-code-mode/tests/harness.test.ts new file mode 100644 index 0000000000..14b41a7901 --- /dev/null +++ b/packages/ai-code-mode/tests/harness.test.ts @@ -0,0 +1,482 @@ +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']) + }) +}) 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 7df08d78ac..97782c8bf6 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2269,6 +2269,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) @@ -3687,6 +3693,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') + }) }) From 9bdd8600090271e9aa9ece6e229835f3c434d005 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:39:03 +0200 Subject: [PATCH 4/9] fix(ai-harness): adopt object params from #1779 The first-party plugins and their tests now call setConfig and answer with one object. Prompt results in the tests also have media: []. --- packages/ai-harness/src/first-party/agents.ts | 5 ++++- .../ai-harness/src/first-party/model-picker.ts | 2 +- packages/ai-harness/src/first-party/permissions.ts | 5 ++++- packages/ai-harness/tests/agents.test.ts | 14 ++++++++------ packages/ai-harness/tests/first-party.test.ts | 2 +- packages/ai-harness/tests/permissions.test.ts | 2 +- packages/ai-harness/tests/question.test.ts | 2 +- 7 files changed, 20 insertions(+), 12 deletions(-) diff --git a/packages/ai-harness/src/first-party/agents.ts b/packages/ai-harness/src/first-party/agents.ts index 7e4a2926e9..2b8970c11a 100644 --- a/packages/ai-harness/src/first-party/agents.ts +++ b/packages/ai-harness/src/first-party/agents.ts @@ -350,7 +350,10 @@ export function agents(options: { const name = typeof input === 'string' ? input.trim() : '' const list = `Agents: ${names.join(', ')}.` if (name === '') return `Agent: ${active().name}. ${list}` - const receipt = await ctx.session.setConfig('agent', name) + const receipt = await ctx.session.setConfig({ + key: 'agent', + value: name, + }) return receipt.status === 'rejected' ? `Unknown agent "${name}". ${list}` : `Agent: ${name}. It applies at the next turn.` diff --git a/packages/ai-harness/src/first-party/model-picker.ts b/packages/ai-harness/src/first-party/model-picker.ts index 30261ae7f8..b84943d044 100644 --- a/packages/ai-harness/src/first-party/model-picker.ts +++ b/packages/ai-harness/src/first-party/model-picker.ts @@ -58,7 +58,7 @@ export function modelPicker(options: { return `Model: ${String(ctx.config.get('model'))}. Choices: ${names.join(', ')}.` if (!names.includes(name)) return `Unknown model "${name}". Choices: ${names.join(', ')}.` - await ctx.session.setConfig('model', name) + await ctx.session.setConfig({ key: 'model', value: name }) return `Model: ${name}. It applies at the next turn.` }, }), diff --git a/packages/ai-harness/src/first-party/permissions.ts b/packages/ai-harness/src/first-party/permissions.ts index e8dd018975..923198e4db 100644 --- a/packages/ai-harness/src/first-party/permissions.ts +++ b/packages/ai-harness/src/first-party/permissions.ts @@ -683,7 +683,10 @@ export function permissions( run: async (input: unknown) => { const next = typeof input === 'string' ? input.trim() : '' if (!next) return `Mode: ${mode()}.` - const receipt = await ctx.session.setConfig('mode', next) + const receipt = await ctx.session.setConfig({ + key: 'mode', + value: next, + }) return receipt.status === 'rejected' ? `Unknown mode "${next}". Modes: ${PERMISSION_MODES.join(', ')}.` : `Mode: ${next}.` diff --git a/packages/ai-harness/tests/agents.test.ts b/packages/ai-harness/tests/agents.test.ts index c167cbb76b..68e0684b65 100644 --- a/packages/ai-harness/tests/agents.test.ts +++ b/packages/ai-harness/tests/agents.test.ts @@ -147,11 +147,13 @@ describe('agents', () => { agents({ adapter: () => adapter }), ]) - expect((await session.setConfig('agent', 'plan')).status).toBe('accepted') + expect( + (await session.setConfig({ key: 'agent', value: 'plan' })).status, + ).toBe('accepted') // A subagent is not a primary agent. - expect((await session.setConfig('agent', 'general')).status).toBe( - 'rejected', - ) + expect( + (await session.setConfig({ key: 'agent', value: 'general' })).status, + ).toBe('rejected') await session.prompt('plan it') expect(toolNames(calls[0])).toEqual(['read_file']) @@ -175,9 +177,9 @@ describe('agents', () => { agents({ adapter: byModel, agents: [quick] }), ]) - expect(await session.prompt('one')).toEqual({ text: 'from main' }) + expect(await session.prompt('one')).toEqual({ text: 'from main', media: [] }) await session.command('agent', 'quick') - expect(await session.prompt('two')).toEqual({ text: 'from fast' }) + expect(await session.prompt('two')).toEqual({ text: 'from fast', media: [] }) await host.close() }) diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index 1b6fc092a1..853a4fbd2a 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -77,7 +77,7 @@ describe('modelPicker', () => { }), ]) expect(await session.command('model', 'fast')).toContain('Model: fast') - expect(await session.prompt('hi')).toEqual({ text: 'fast answer' }) + expect(await session.prompt('hi')).toEqual({ text: 'fast answer', media: [] }) expect(await session.command('model', 'nope')).toContain('Unknown model') await host.close() }) diff --git a/packages/ai-harness/tests/permissions.test.ts b/packages/ai-harness/tests/permissions.test.ts index d6a64ed6c3..5efc4f347e 100644 --- a/packages/ai-harness/tests/permissions.test.ts +++ b/packages/ai-harness/tests/permissions.test.ts @@ -300,7 +300,7 @@ async function answer(session: HarnessSession, value: unknown) { ) const [question] = session.snapshot().pendingQuestions if (!question) throw new Error('No question.') - await session.answer(question.questionId, value) + await session.answer({ questionId: question.questionId, value }) return question.message } diff --git a/packages/ai-harness/tests/question.test.ts b/packages/ai-harness/tests/question.test.ts index dafa3de4b5..feadd904e1 100644 --- a/packages/ai-harness/tests/question.test.ts +++ b/packages/ai-harness/tests/question.test.ts @@ -53,7 +53,7 @@ async function reply(session: HarnessSession, value: unknown) { ) const [pending] = session.snapshot().pendingQuestions if (!pending) throw new Error('No question.') - await session.answer(pending.questionId, value) + await session.answer({ questionId: pending.questionId, value }) return pending } From b113c45ab8531da92c7195f5e0b5507cefbce99b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:39:03 +0200 Subject: [PATCH 5/9] docs(ai-harness): answer permission questions with onEvent The permissions page listens for the question with session.onEvent instead of a loop over session.events(). The permissions and agents pages call setConfig and answer with one object. --- docs/config.json | 6 ++++-- docs/harness/agents.md | 2 +- docs/harness/permissions.md | 25 +++++++++++-------------- 3 files changed, 16 insertions(+), 17 deletions(-) diff --git a/docs/config.json b/docs/config.json index 75a1aa5b77..f6d718a6c8 100644 --- a/docs/config.json +++ b/docs/config.json @@ -997,12 +997,14 @@ { "label": "Ask before risky tool calls", "to": "harness/permissions", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-14" }, { "label": "Switch between agent profiles", "to": "harness/agents", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-14" }, { "label": "Run agents from a harness", diff --git a/docs/harness/agents.md b/docs/harness/agents.md index 8d918b519d..b4eab5ae94 100644 --- a/docs/harness/agents.md +++ b/docs/harness/agents.md @@ -85,7 +85,7 @@ if (setting?.option.type === 'select') { console.log(setting.value, setting.option.options) // 'build' ['build', 'plan'] } -await session.setConfig('agent', 'plan') +await session.setConfig({ key: 'agent', value: 'plan' }) ``` A client can change the setting only when the harness exposes it. Add `expose: { config: ['agent'] }` to `defineHarness`. diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 9062b3f9bf..18f945e33a 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -57,24 +57,21 @@ The rule makes each `save_note` call wait for the user. ## 2. Answer the first question -The question shows up as a `harness.question` event and in `session.snapshot().pendingQuestions`. Answer it with `session.answer`: +The question shows up as a `harness.question` event and in `session.snapshot().pendingQuestions`. Listen for the event, then answer with `session.answer`: ```ts group=harness-permissions -const turn = session.prompt('Save a note with the time of the meeting.') - -for await (const { event } of session.events()) { - if (event.type === 'CUSTOM' && event.name === HARNESS_EVENTS.question) { - const [question] = session.snapshot().pendingQuestions - // For example: Allow save_note {"path":"meeting.md","text":"..."}? - if (question) await session.answer(question.questionId, 'once') - break - } -} +const stopAnswering = session.onEvent('CUSTOM', async ({ event }) => { + if (event.name !== HARNESS_EVENTS.question) return + const [question] = session.snapshot().pendingQuestions + // For example: Allow save_note {"path":"meeting.md","text":"..."}? + if (question) await session.answer({ questionId: question.questionId, value: 'once' }) +}) -await turn +await session.prompt('Save a note with the time of the meeting.') +stopAnswering() ``` -The agent writes the file and continues. +The listener says `once` to each question. The agent writes the file and continues. ## Pick a mode @@ -90,7 +87,7 @@ The `mode` setting decides how much runs without a question: Switch the mode from server code: ```ts group=harness-permissions -await session.setConfig('mode', 'plan') +await session.setConfig({ key: 'mode', value: 'plan' }) // Or run the command. It returns 'Mode: plan.' await session.command('mode', 'plan') ``` From aea1f8a8e5be811a3fe543df1c050ec154b6c17d Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:42:59 +0000 Subject: [PATCH 6/9] ci: apply automated fixes --- packages/ai-harness/tests/agents.test.ts | 10 ++++++++-- packages/ai-harness/tests/first-party.test.ts | 5 ++++- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/packages/ai-harness/tests/agents.test.ts b/packages/ai-harness/tests/agents.test.ts index 68e0684b65..09134886ba 100644 --- a/packages/ai-harness/tests/agents.test.ts +++ b/packages/ai-harness/tests/agents.test.ts @@ -177,9 +177,15 @@ describe('agents', () => { agents({ adapter: byModel, agents: [quick] }), ]) - expect(await session.prompt('one')).toEqual({ text: 'from main', media: [] }) + expect(await session.prompt('one')).toEqual({ + text: 'from main', + media: [], + }) await session.command('agent', 'quick') - expect(await session.prompt('two')).toEqual({ text: 'from fast', media: [] }) + expect(await session.prompt('two')).toEqual({ + text: 'from fast', + media: [], + }) await host.close() }) diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index 853a4fbd2a..f61ebc7838 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -77,7 +77,10 @@ describe('modelPicker', () => { }), ]) expect(await session.command('model', 'fast')).toContain('Model: fast') - expect(await session.prompt('hi')).toEqual({ text: 'fast answer', media: [] }) + expect(await session.prompt('hi')).toEqual({ + text: 'fast answer', + media: [], + }) expect(await session.command('model', 'nope')).toContain('Unknown model') await host.close() }) From 5225695c27d08862b63d5c27f0c304721cf279ff Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:56:57 +0200 Subject: [PATCH 7/9] docs(ai-harness): answer each open permission question once Tool calls can run in parallel, so two questions can be open at the same time. The sample now answers each open question once, by id, instead of the first one in the list. --- docs/harness/permissions.md | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 18f945e33a..ef07be211e 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -60,11 +60,16 @@ The rule makes each `save_note` call wait for the user. The question shows up as a `harness.question` event and in `session.snapshot().pendingQuestions`. Listen for the event, then answer with `session.answer`: ```ts group=harness-permissions +// Calls can run in parallel, so answer each open question once. +const answered = new Set() const stopAnswering = session.onEvent('CUSTOM', async ({ event }) => { if (event.name !== HARNESS_EVENTS.question) return - const [question] = session.snapshot().pendingQuestions - // For example: Allow save_note {"path":"meeting.md","text":"..."}? - if (question) await session.answer({ questionId: question.questionId, value: 'once' }) + for (const question of session.snapshot().pendingQuestions) { + if (answered.has(question.questionId)) continue + answered.add(question.questionId) + // For example: Allow save_note {"path":"meeting.md","text":"..."}? + await session.answer({ questionId: question.questionId, value: 'once' }) + } }) await session.prompt('Save a note with the time of the meeting.') From d890fa399558aa3acd4f7d93bb10d46791adace5 Mon Sep 17 00:00:00 2001 From: Jack Herrington Date: Thu, 15 Oct 2026 02:34:42 -0700 Subject: [PATCH 8/9] fix(ai-code-mode): tell the model what execute_typescript returns (#1785) fix(ai-code-mode): describe execute_typescript result shape to the model The tool description and system prompt now say the tool returns { success, result, logs, error }, where result is the program's return value, so models stop guessing at the wrapper. Co-authored-by: Claude Opus 5.5 (1M context) --- .changeset/code-mode-result-shape.md | 5 +++++ packages/ai-code-mode/src/create-code-mode-tool.ts | 3 ++- packages/ai-code-mode/src/create-system-prompt.ts | 2 +- packages/ai-code-mode/tests/create-code-mode-tool.test.ts | 1 + packages/ai-code-mode/tests/create-system-prompt.test.ts | 1 + 5 files changed, 10 insertions(+), 2 deletions(-) create mode 100644 .changeset/code-mode-result-shape.md 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/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/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', () => { From 572f3e37681875aadbeca88ed5992267a553ffd6 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 15 Oct 2026 13:51:26 +0200 Subject: [PATCH 9/9] test(ai-code-mode): cover the code mode allow rules and the prompt reset Add tests: a program runs with no question when permissions() asks by default, in default and plan mode. A user rule that asks for execute_typescript still wins. A failed pick drops the prompt of the last turn. The docs list all six isolate drivers and say what happens when two tools get the same safe name. --- docs/harness/code-mode.md | 6 +- packages/ai-code-mode/tests/harness.test.ts | 105 ++++++++++++++++++++ 2 files changed, 109 insertions(+), 2 deletions(-) diff --git a/docs/harness/code-mode.md b/docs/harness/code-mode.md index 30c03dabdb..266c022390 100644 --- a/docs/harness/code-mode.md +++ b/docs/harness/code-mode.md @@ -67,9 +67,11 @@ The model now has an `execute_typescript` tool. Inside the program, each moved t | 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. @@ -101,7 +103,7 @@ export const onlyIssues = codeMode({ - 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. +The other options of `createCodeMode` pass through: `timeout`, `memoryLimit`, `lazyToolsConfig`, and more. ## 4. Send only the tool names @@ -125,7 +127,7 @@ The model then works in three steps: ## 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`. +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 diff --git a/packages/ai-code-mode/tests/harness.test.ts b/packages/ai-code-mode/tests/harness.test.ts index 14b41a7901..fb5cf9de0a 100644 --- a/packages/ai-code-mode/tests/harness.test.ts +++ b/packages/ai-code-mode/tests/harness.test.ts @@ -479,4 +479,109 @@ describe('codeMode', () => { }) 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() + }) })