From 1cdc116aa3c94b3b9433c084dd7fef34139d9f35 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:01:15 +0200 Subject: [PATCH 01/15] 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 02/15] 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 a8949873ca8fad08a98dc10104ac6dffb65b2030 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:28:13 +0200 Subject: [PATCH 03/15] feat(ai-harness): coding tools in @tanstack/ai-harness/plugins/coding Add workspaceTools (read, write, edit or patch, list, grep, bash, webfetch, websearch), the WorkspaceBackend with hostBackend, snapshots (/undo, /redo, revert file restore), and formatter in a new Node entry. turndown and @vscode/ripgrep are optional peers. Split out of #1555. --- .changeset/ai-harness-coding-tools.md | 13 + docs/config.json | 21 + docs/harness/coding-agent.md | 176 +++ docs/harness/coding-tools-backends.md | 215 ++++ docs/harness/coding-tools.md | 285 +++++ docs/harness/permissions.md | 45 +- docs/harness/snapshots.md | 128 ++ packages/ai-harness/package.json | 18 +- .../src/first-party/coding/backend.ts | 360 ++++++ .../ai-harness/src/first-party/coding/bash.ts | 259 ++++ .../ai-harness/src/first-party/coding/edit.ts | 190 +++ .../src/first-party/coding/formatter.ts | 272 +++++ .../src/first-party/coding/index.ts | 12 + .../src/first-party/coding/patch.ts | 450 +++++++ .../ai-harness/src/first-party/coding/read.ts | 201 ++++ .../src/first-party/coding/search.ts | 428 +++++++ .../src/first-party/coding/snapshots.ts | 398 +++++++ .../ai-harness/src/first-party/coding/web.ts | 472 ++++++++ .../src/first-party/coding/workspace.ts | 446 +++++++ .../ai-harness/tests/coding-backend.test.ts | 183 +++ packages/ai-harness/tests/coding-bash.test.ts | 283 +++++ packages/ai-harness/tests/coding-edit.test.ts | 162 +++ .../ai-harness/tests/coding-formatter.test.ts | 446 +++++++ .../ai-harness/tests/coding-patch.test.ts | 483 ++++++++ packages/ai-harness/tests/coding-read.test.ts | 308 +++++ .../ai-harness/tests/coding-search.test.ts | 502 ++++++++ .../ai-harness/tests/coding-snapshots.test.ts | 736 ++++++++++++ packages/ai-harness/tests/coding-web.test.ts | 384 ++++++ packages/ai-harness/tests/edge-safety.test.ts | 8 +- packages/ai-harness/tests/first-party.test.ts | 52 +- .../ai-harness/tests/instructions.test.ts | 37 +- .../ai-harness/tests/workspace-tools.test.ts | 1037 +++++++++++++++++ packages/ai-harness/vite.config.ts | 6 +- pnpm-lock.yaml | 133 +++ .../e2e/fixtures/tools-test/coding-tools.json | 178 +++ testing/e2e/src/routeTree.gen.ts | 21 + testing/e2e/src/routes/api.coding-tools.ts | 143 +++ .../e2e/tests/tools-test/coding-tools.spec.ts | 193 +++ 38 files changed, 9675 insertions(+), 9 deletions(-) create mode 100644 .changeset/ai-harness-coding-tools.md create mode 100644 docs/harness/coding-agent.md create mode 100644 docs/harness/coding-tools-backends.md create mode 100644 docs/harness/coding-tools.md create mode 100644 docs/harness/snapshots.md create mode 100644 packages/ai-harness/src/first-party/coding/backend.ts create mode 100644 packages/ai-harness/src/first-party/coding/bash.ts create mode 100644 packages/ai-harness/src/first-party/coding/edit.ts create mode 100644 packages/ai-harness/src/first-party/coding/formatter.ts create mode 100644 packages/ai-harness/src/first-party/coding/index.ts create mode 100644 packages/ai-harness/src/first-party/coding/patch.ts create mode 100644 packages/ai-harness/src/first-party/coding/read.ts create mode 100644 packages/ai-harness/src/first-party/coding/search.ts create mode 100644 packages/ai-harness/src/first-party/coding/snapshots.ts create mode 100644 packages/ai-harness/src/first-party/coding/web.ts create mode 100644 packages/ai-harness/src/first-party/coding/workspace.ts create mode 100644 packages/ai-harness/tests/coding-backend.test.ts create mode 100644 packages/ai-harness/tests/coding-bash.test.ts create mode 100644 packages/ai-harness/tests/coding-edit.test.ts create mode 100644 packages/ai-harness/tests/coding-formatter.test.ts create mode 100644 packages/ai-harness/tests/coding-patch.test.ts create mode 100644 packages/ai-harness/tests/coding-read.test.ts create mode 100644 packages/ai-harness/tests/coding-search.test.ts create mode 100644 packages/ai-harness/tests/coding-snapshots.test.ts create mode 100644 packages/ai-harness/tests/coding-web.test.ts create mode 100644 packages/ai-harness/tests/workspace-tools.test.ts create mode 100644 testing/e2e/fixtures/tools-test/coding-tools.json create mode 100644 testing/e2e/src/routes/api.coding-tools.ts create mode 100644 testing/e2e/tests/tools-test/coding-tools.spec.ts diff --git a/.changeset/ai-harness-coding-tools.md b/.changeset/ai-harness-coding-tools.md new file mode 100644 index 0000000000..2ed90ac76a --- /dev/null +++ b/.changeset/ai-harness-coding-tools.md @@ -0,0 +1,13 @@ +--- +'@tanstack/ai-harness': minor +--- + +Add coding tools in the new Node entry `@tanstack/ai-harness/plugins/coding`: + +- `workspaceTools({ root })` gives `read_file` (text pages, images, PDFs), `write_file`, `edit_file` or `patch` (`editStyle`), `list_files`, `grep`, `bash` (with `background`), `webfetch`, and `websearch` (with a `SearchProvider`). With `permissions()`, edits, `bash`, and `webfetch` ask first, and the rules also apply to the real path of a link. `outside: 'ask'` asks before a path outside `root`. The `cwd` thread setting gives each thread its own working folder. +- A `WorkspaceBackend` runs the tools. `hostBackend` is the default. On Windows, commands do not run a program from the working folder by its name. +- `snapshots()` adds `/undo` and `/redo`, which restore only the files of the turn (git 2.26 or newer), file restore for `session.revert()`, and a diff for each model step. +- `formatter()` formats each written file with the formatter of the project. With `permissions()`, each run asks first. +- `WorkspaceHooks` lets other plugins run code after a read or a write. + +`@vscode/ripgrep` and `turndown` are optional peer dependencies. diff --git a/docs/config.json b/docs/config.json index df3ace69e3..943122e6e4 100644 --- a/docs/config.json +++ b/docs/config.json @@ -993,6 +993,27 @@ { "label": "Ask before risky tool calls", "to": "harness/permissions", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" + }, + { + "label": "Give your agent coding tools", + "to": "harness/coding-tools", + "addedAt": "2026-10-13" + }, + { + "label": "Extend the coding tools", + "to": "harness/coding-tools-backends", + "addedAt": "2026-10-13" + }, + { + "label": "Undo a turn", + "to": "harness/snapshots", + "addedAt": "2026-10-13" + }, + { + "label": "Build a coding agent", + "to": "harness/coding-agent", "addedAt": "2026-10-13" }, { diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md new file mode 100644 index 0000000000..c5f88ff349 --- /dev/null +++ b/docs/harness/coding-agent.md @@ -0,0 +1,176 @@ +--- +title: Build a coding agent +id: harness-coding-agent +order: 6 +description: "Turn a harness into a coding agent with file tools, permission modes, undo, a todo list, a model picker, project instructions, session titles, a cost count, and /compact." +keywords: + - tanstack ai + - harness + - coding agent + - permissions + - workspace tools + - project instructions + - undo +--- + +You want your own coding agent in the terminal. It reads and edits files, and runs commands after you approve them. When a turn goes wrong, you undo it. The first-party plugins in `@tanstack/ai-harness/plugins` and `@tanstack/ai-harness/plugins/coding` give you those parts. You pick the model and the rules. + +## 1. Define the agent + +```ts group=harness-coding-agent +import { homedir } from 'node:os' +import { join } from 'node:path' +import { createInterface } from 'node:readline/promises' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, +} from '@tanstack/ai-harness' +import { + compact, + fileCommands, + modelPicker, + permissions, + projectInstructions, + title, + todos, + usage, +} from '@tanstack/ai-harness/plugins' +import { snapshots, workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { openaiText } from '@tanstack/ai-openai' +import { memoryPersistence } from '@tanstack/ai-persistence' + +const root = process.cwd() +const smart = openaiText('gpt-6.1-sol') +const fast = openaiText('gpt-6-luna') + +const coder = defineHarness({ + name: 'acme/coder', + adapter: smart, + systemPrompts: ['You are a careful coding agent. Read before you edit.'], + plugins: () => [ + permissions({ root }), + workspaceTools({ root }), + snapshots({ root, dataDir: join(homedir(), '.acme-coder', 'snapshots') }), + todos(), + modelPicker({ choices: { smart, fast }, default: 'smart' }), + fileCommands({ dir: `${root}/.claude/commands` }), + compact({ adapter: fast }), + title({ adapter: fast }), + usage(), + // Keep this plugin last. See "Project instructions" below. + projectInstructions({ root }), + ], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(coder, { threadId: 'main' }) +const terminal = createInterface({ + input: process.stdin, + output: process.stdout, +}) + +for (;;) { + const line = await terminal.question('\n> ') + // `/undo`, `/model fast`, `/mode plan`, and the other plugin commands. + if (line.startsWith('/')) { + const [name = '', ...args] = line.slice(1).split(' ') + console.log(await session.command(name, args.join(' ') || undefined)) + continue + } + for await (const chunk of session.prompt(line).stream()) { + if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) + // A risky call waits for an answer: once, always, or reject. + if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { + const [question] = session.snapshot().pendingQuestions + if (question) { + const answer = await terminal.question(`\n${question.message} `) + await session.answer(question.questionId, answer) + } + } + } +} +``` + +Both plugin entries run only in Node. The coding tools use the file system and the shell of the machine. A line that starts with `/` runs a command of a plugin. Any other line is a prompt. + +## 2. Run it + +1. Run `npx tsx coder.ts` in your project. +2. Ask for a change. The agent reads files freely, and asks before it writes a file or runs a command. +3. Answer `once` to allow the call, `always` to allow it from now on, or `reject` to refuse it. +4. If you do not like the result, run `/undo`. The files go back to how they were before the turn. + +## What each plugin adds + +| Plugin | Adds | +|---|---| +| `permissions({ root })` | Asks before risky tool calls. `/mode` switches between `default`, `plan` (read-only), `acceptEdits` (edits run without asking), and `bypass`. See [Ask before risky tool calls](./permissions). | +| `workspaceTools({ root })` | Tools that read, search, and change the files in `root`, run commands, and read web pages. A path outside `root` is refused, or asks first with `outside: 'ask'`. See [Give your agent coding tools](./coding-tools). | +| `snapshots({ root, dataDir })` | `/undo` puts the files back as they were before the last turn, and `/redo` brings them back. It needs `git`. See [Undo a turn](./snapshots). | +| `todos()` | A `todo_write` tool the model uses for multi-step work, and `/todos`. | +| `modelPicker({ choices })` | `/model ` switches the model at the next turn. | +| `fileCommands({ dir })` | Each `.md` file becomes a slash command. `$ARGUMENTS` is replaced by what you type after it. | +| `compact({ adapter })` | `/compact` replaces a long conversation with a summary. | +| `title({ adapter })` | Names the session from its first message. | +| `usage()` | `/usage` shows the tokens and the cost of the session, with a line for each model. See [Track usage and cost](./usage). | +| `projectInstructions({ root })` | Adds `AGENTS.md`, `CLAUDE.md`, and facts about the environment to the system prompt. | + +To compact on its own before the context limit, see [Compact a harness session](./compaction). + +## Project instructions + +Your repository has an `AGENTS.md` or a `CLAUDE.md` with its rules. `projectInstructions({ root })` gives them to the model: + +- It reads `AGENTS.md` and `CLAUDE.md` in each folder from the repository root (the folder with `.git`) down to `root`. Set other file names with `files`. +- `global` adds files that come first, for example `['~/.config/AGENTS.md']`. +- When `read_file` reads a file in a subfolder, the model also gets the instruction files of the folders on the way, once each. +- An environment block gives the date, the platform, the working folder, and the git repository and branch. `env: false` turns it off. +- When you change an instruction file, the next turn sends the new text as a note. + +Put `projectInstructions()` last in the plugin list. Then a changed file goes into the conversation as a [mid-conversation change](../advanced/mid-conversation-changes), and the prompt cache stays. When the prompt of another plugin comes after it, a change builds the system prompt again, and the prompt cache starts again. + +## Session titles and cost + +`title({ adapter })` and `usage()` write to the session index, one entry per thread: + +- `title` names the session from its first message. The turn does not wait for it. +- When `usage` knows the cost of the session, it writes the cost. A cost that the provider reports is used as it is. +- To price the calls of a provider that reports no cost, pass `usage({ model })`. [Track usage and cost](./usage#show-usage-and-price-the-calls) shows how. + +A [sidebar of sessions](./sessions) reads the title and the cost from the index. The `usage()` state also has `contextTokens`, the size of the lead model's context at its latest call, for a context meter. + +## Add more + +- `agents({ adapter })`: named agents with their own prompt, tools, and model. `/agent plan` switches to a read-only agent. See [Switch between agent profiles](./agents). +- `question()`: a `question` tool. The model asks the user to pick an option, and waits for the answer. +- `formatter({ root })`: formats each file after a write. See [Format each file after a write](./coding-tools-backends#format-each-file-after-a-write). + +## Add your own rules + +Tool plugins add permission rules to the `PermissionRules` extension point. Add your own for any tool: + +```ts group=harness-coding-agent +import { PermissionRules } from '@tanstack/ai-harness/plugins' +import { definePlugin } from '@tanstack/ai-harness' + +export const noDeploys = definePlugin({ + name: 'acme/no-deploys', + setup: () => ({ + contribute: [PermissionRules.item({ tool: 'deploy_*', decision: 'deny' })], + }), +}) +``` + +A trailing `*` matches every tool that starts with the text. The last matching rule wins. [Ask before risky tool calls](./permissions) explains the rules, the modes, and the saved answers. + +The workspace tools run on your machine with your permissions. To run them on another machine, [write your own backend](./coding-tools-backends#write-your-own-backend). + +## What you have now + +- A terminal coding agent with file tools, approvals, modes, a todo list, and a model picker. +- `/undo` and `/redo` for the file changes of a turn. +- The rules of your repository in every turn. A change to them keeps the prompt cache. +- A title and a cost for each session. + +Next: let the user undo more than one turn with [Undo a turn](./snapshots). diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md new file mode 100644 index 0000000000..da9e702484 --- /dev/null +++ b/docs/harness/coding-tools-backends.md @@ -0,0 +1,215 @@ +--- +title: Extend the coding tools +id: harness-coding-tools-backends +order: 7 +description: "Run the coding tools on your own backend, format each written file, and keep tool output small." +keywords: + - tanstack ai + - harness + - coding agent + - workspace tools + - formatter + - workspace hooks +--- + +Your agent must work on files that are not on this machine. Or you want code to run after each write, like a formatter. The [coding tools](./coding-tools) send every read, write, and command through a backend, and they run hooks after each read and write. This page shows how to change both. + +## Write your own backend + +A `WorkspaceBackend` is where the tools read and write files and run commands. Write one for another place, for example a remote machine. Every path that the tools give it is absolute. + +The backend must have: + +| Method | What it does | +| --- | --- | +| `readFile(path)` | Gives the bytes of a file. Throws when there is no file. | +| `writeFile(path, data)` | Creates or replaces a file. Creates the missing parent folders. | +| `stat(path)` | Gives `{ type, size, mtimeMs }`, or `undefined` when nothing is at `path`. | +| `readdir(path)` | Gives the entries of a folder. A link has the type `'link'`. | +| `exec(command, options)` | Runs a shell command and waits. Gives `{ exitCode, stdout, stderr }`. | + +`exec` resolves with the exit code of a command that fails. It does not throw. After `options.timeoutMs`, it stops the command with every process that it started, and gives exit code 124. + +The backend can also have: + +| Member | What it does | Without it | +| --- | --- | --- | +| `shell` | `'sh'` or `'cmd'`. Sets the path style and how arguments are quoted. | `sh` quotes, and the path style of this machine. | +| `remove(path)` | Removes one file. | `patch` cannot delete or move files. | +| `realpath(path)` | Gives the path with every link resolved. | The tools do not check links. | +| `spawn(command, options)` | Starts a command in the background. | `background: true` gives an error. | + +The fastest start is to wrap `hostBackend`. This backend logs each command before it runs: + +```ts group=harness-coding-tools-backends +import { hostBackend, workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' + +const root = process.cwd() + +const logged: WorkspaceBackend = { + ...hostBackend, + exec: (command, options) => { + console.log(`$ ${command}`) + return hostBackend.exec(command, options) + }, +} + +export const loggedTools = workspaceTools({ root, backend: logged }) +``` + +Only `hostBackend` itself uses the `@vscode/ripgrep` binary. A wrapped backend uses `rg` on the `PATH`. + +## Run code after a read or a write + +Another plugin can add code that runs after the tools read or write a file. Add a `WorkspaceHooks` item: + +```ts group=harness-coding-tools-backends +import { definePlugin } from '@tanstack/ai-harness' +import { WorkspaceHooks } from '@tanstack/ai-harness/plugins/coding' + +export const generatedFiles = definePlugin({ + name: 'acme/generated-files', + setup: () => ({ + contribute: [ + WorkspaceHooks.item({ + afterRead: async (path) => + path.endsWith('.gen.ts') + ? 'This file is generated. Change its source, not this file.' + : undefined, + afterWrite: async (path) => { + console.log(`changed ${path}`) + }, + }), + ], + }), +}) +``` + +- `afterRead` gets the absolute path. The text that it gives is added after the lines that `read_file` shows. +- `afterWrite` runs after `write_file`, `edit_file`, and `patch`. +- A hook runs while the tool holds the lock on that path. A hook that throws fails the tool call. + +Add the plugin next to the tools: `plugins: () => [workspaceTools({ root }), generatedFiles]`. + +## Format each file after a write + +The model writes code that your formatter changes later, and the diff gets noisy. `formatter()` formats each file that the tools write: + +```ts group=harness-coding-tools-backends +import { defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { formatter } from '@tanstack/ai-harness/plugins/coding' +import { openaiText } from '@tanstack/ai-openai' + +export const formatted = defineHarness({ + name: 'acme/formatted-coder', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + permissions({ root }), + workspaceTools({ root }), + formatter({ root }), + ], +}) +``` + +At the first write of a session, `formatter()` looks at the files in `root`. For each file, the first formatter that the project uses for that file type runs. + +For JavaScript, TypeScript, JSON, CSS, and other web files: + +| Formatter | Runs when `root` has | +| --- | --- | +| prettier | `.prettierrc*`, `prettier.config.*`, or a `prettier` field in `package.json` | +| biome | `biome.json` or `biome.jsonc` | +| oxfmt | `.oxfmtrc*`, or `oxfmt` in the `devDependencies` | + +These three run with `npx --no-install`, so the project must have them installed. + +For other languages: + +| Formatter | Files | Runs when `root` has | +| --- | --- | --- | +| ruff | `.py`, `.pyi` | `ruff.toml`, `.ruff.toml`, or `[tool.ruff]` in `pyproject.toml` | +| gofmt | `.go` | `go.mod` | +| rustfmt | `.rs` | `Cargo.toml` | + +Add your own formatter. Your formatters come before the built-in ones: + +```ts group=harness-coding-tools-backends +export const withToml = formatter({ + root, + formatters: [ + { + name: 'taplo', + extensions: ['.toml'], + command: (file) => `taplo fmt ${file}`, + }, + ], +}) +``` + +- `command` gets the absolute path, already quoted for the shell. +- `when(project)` limits a formatter to some projects. Without it, the formatter is used in every project. +- `builtins: false` turns the built-in formatters off. +- `timeoutMs` limits each run. The default is 20 seconds. +- On your own backend, give `formatter()` the same `backend` as `workspaceTools()`. + +A formatter that fails, stops at the time limit, or is not installed does not fail the tool call. The file stays as the tool wrote it, and the plugin sends `FormatFailed`: + +```ts group=harness-coding-tools-backends +import { FormatFailed } from '@tanstack/ai-harness/plugins/coding' + +export const formatAlerts = definePlugin({ + name: 'acme/format-alerts', + setup: (ctx) => { + ctx.on(FormatFailed, (failure) => + console.warn(`${failure.path}: ${failure.message}`), + ) + }, +}) +``` + +Clients get the same event as a `harness.plugin.event` with the name `tanstack/formatter:failed`. + +### Each run asks first + +A formatter runs code from the project, like `prettier.config.js`. The model can change that code with an edit. So with `permissions()`, the formatter asks before each run, also in `acceptEdits` mode. The question names the run `formatter:prettier`, `formatter:biome`, and so on. + +- Answer `always` once, and the formatter runs without a question in this project. +- A `reject` keeps the write. The tool result says that the formatter did not run. +- A rule can allow a formatter: `{ tool: 'formatter:prettier', decision: 'allow' }`. +- Without `permissions()`, the formatter runs without a question. + +[Formatters](./permissions#formatters) shows what each mode does. + +## Keep tool output small + +One tool result with a full log or a big JSON file fills the context of the model. `boundToolOutput()` cuts every tool result that is too long: + +```ts group=harness-coding-tools-backends +import { boundToolOutput } from '@tanstack/ai-harness/plugins' + +export const bounded = defineHarness({ + name: 'acme/bounded-coder', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + workspaceTools({ root }), + boundToolOutput({ dir: '.agent/tool-output' }), + ], +}) +``` + +- A result over 2000 lines or 50 KiB keeps its first lines and gets a note. Set other limits with `maxLines` and `maxBytes`. +- With `dir`, the full output goes to a file in `dir`, and the note gives the path. Files older than `retentionDays` (default 7) are removed. +- It cuts the result of every tool, MCP tools too, in the lead turn and in agent runs. +- An error stays an error. Images and other content parts do not change. + +`read_file` and `bash` keep their output inside the default limits, so `boundToolOutput()` does not cut them again. + +## What you have now + +- A backend of your own for another place. +- Code that runs after each read and each write. +- Formatted files after each write, and tool results that fit the context. + +Next: let the user undo the changes of a turn with [snapshots](./snapshots). diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md new file mode 100644 index 0000000000..55431056e8 --- /dev/null +++ b/docs/harness/coding-tools.md @@ -0,0 +1,285 @@ +--- +title: Give your agent coding tools +id: harness-coding-tools +order: 6 +description: "Let a harness agent read, change, and search files, run commands, and read web pages, with a question before each risky call." +keywords: + - tanstack ai + - harness + - coding agent + - workspace tools + - permissions + - formatter +--- + +Your agent must read and change the files of a project, search them, and run commands like `pnpm test`. It must stay in that project, and it must ask before it changes a file. `workspaceTools()` gives the model these tools for one folder. `permissions()` asks the user before each risky call. + +## 1. Install + + + +react: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +vue: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +solid: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +svelte: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +preact: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +angular: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +octane: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai +vanilla: @tanstack/ai @tanstack/ai-harness @tanstack/ai-persistence @tanstack/ai-openai + + + +Two optional packages make the tools better: + + + +react: @vscode/ripgrep turndown +vue: @vscode/ripgrep turndown +solid: @vscode/ripgrep turndown +svelte: @vscode/ripgrep turndown +preact: @vscode/ripgrep turndown +angular: @vscode/ripgrep turndown +octane: @vscode/ripgrep turndown +vanilla: @vscode/ripgrep turndown + + + +- `@vscode/ripgrep`: a fast `rg` binary for `list_files` and `grep`. Without it, the tools use `rg` on the `PATH`. Without `rg`, they use `git ls-files`, then a walk through the folders. +- `turndown`: `webfetch` gives HTML pages as Markdown. Without it, HTML pages come back as plain text. + +## 2. Add the tools + +```ts group=harness-coding-tools +import { createInterface } from 'node:readline/promises' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, +} from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { openaiText } from '@tanstack/ai-openai' +import { memoryPersistence } from '@tanstack/ai-persistence' + +const root = process.cwd() + +const coder = defineHarness({ + name: 'acme/coder', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [permissions({ root }), workspaceTools({ root })], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(coder, { threadId: 'main' }) +const terminal = createInterface({ + input: process.stdin, + output: process.stdout, +}) + +// Send each line as a prompt. Print the answer. Ask the user each +// question that a risky call waits on. +for (;;) { + const line = await terminal.question('\n> ') + for await (const chunk of session.prompt(line).stream()) { + if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) + if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { + const [question] = session.snapshot().pendingQuestions + if (question) { + const answer = await terminal.question(`\n${question.message} `) + await session.answer(question.questionId, answer) + } + } + } +} +``` + +`@tanstack/ai-harness/plugins/coding` runs only in Node. The tools use the file system and the shell of the machine. + +## 3. Run it + +1. Run `npx tsx coder.ts` in your project. +2. Ask: `find the failing test and fix it`. +3. When the agent wants to change a file or run a command, it asks in the terminal. Answer `once`, `always`, or `reject`. + +The agent reads and searches without a question. For a full coding agent with undo, a todo list, and a model picker, see [Build a coding agent](./coding-agent). + +## The tools + +### Read and search + +These tools run without a question. + +| Tool | What it does | +| --- | --- | +| `read_file` | Reads a text file as numbered lines, at most 2000 lines or 50 KiB a page. The model reads more with `offset` and `limit`. | +| `list_files` | Lists the files in a folder. An optional glob, like `src/**/*.ts`, filters them. | +| `grep` | Searches the file contents with a regular expression. Gives `file:line: text`. | + +- `read_file` gives images (PNG, JPEG, GIF, WebP) and PDFs up to 5 MB to the model as media. +- `read_file` refuses a binary file with a short note. For a missing file, it names up to 3 files with a close name. +- `list_files` and `grep` skip the files that `.gitignore` names, the `.git` and `node_modules` folders, and links. +- `list_files` shows at most 1000 files. `grep` shows at most 200 matches. + +### Change files + +These tools ask first. + +| Tool | What it does | +| --- | --- | +| `write_file` | Creates or replaces a file. | +| `edit_file` | Replaces text in a file. The old text must be in the file once, unless the model sets `replaceAll`. | +| `patch` | Applies a patch that adds, deletes, updates, or moves files. | + +A model gets `edit_file` or `patch`, not both. Every model gets `write_file`. The `editStyle` option picks the edit tool: + +- `'auto'` (default): GPT models (ids that start with `gpt-`, or `o` and a digit, like `o3`) get `patch`. Other models get `edit_file`. +- `'edit'` or `'patch'`: every model gets that tool. + +`edit_file` keeps the line endings and the byte order mark of the file. It also finds text that is different only in Unicode form or in the spaces at the end of a line. + +To let the user undo the changes of a turn, add [snapshots](./snapshots). + +### Commands and the web + +| Tool | What it does | Asks first | +| --- | --- | --- | +| `bash` | Runs a shell command in the working folder. | Yes | +| `webfetch` | Reads a web page by its URL. | Yes | +| `websearch` | Searches the web. It is there only with a search provider. | No | + +For `bash`: + +- The model gets the exit code and the last 1000 lines (20 KiB) of the output. +- A command stops after 2 minutes, with exit code 124. Every process that it started stops too. Set another limit with `bashTimeoutMs`. The model can ask for up to 10 minutes. +- With `background: true`, the call returns at once with a job id. When the job ends, the model gets a note with the output, and a new turn starts. The time limit does not apply to a background job. +- With `spillDir`, long output goes to a file in that folder, and the model gets the path. `spillDir` is relative to `root`, or absolute. +- When the session closes, the background jobs stop. +- On a durable host (one with `stores.log`), the log keeps each background job until it ends. When a crash stops a job, the next host adds the note "Background job bash-1 stopped when the host restarted." Without `stores.log`, a crash stops the job, and no note comes. + +On Windows, the default backend runs commands in `cmd.exe`. A program name like `rg` or `gradlew` comes from the `PATH` only, not from the working folder. So a cloned repo cannot put its own `rg.cmd` in place of `rg`. To run a script in the working folder, give its path: `.\gradlew build`. + +For `webfetch`: + +- It reads `http` and `https` URLs only, and follows at most 5 redirects. +- It refuses `localhost`, private network addresses, and cloud metadata addresses. It checks the DNS answer of each connection. +- HTML comes back as Markdown, or as text without `turndown`. The model can ask for `format: 'html'`. Text and JSON come back as they are. +- It reads at most 5 MB. It stops after 30 seconds, and the model can ask for up to 120 seconds. + +Change the web tools with the `web` option: + +- `web: false`: no web tools. +- `web: { allowPrivateHosts: true }`: `webfetch` can read `localhost` and your private network. Set it only for a model that is allowed to reach your network. +- `web: { search }`: adds `websearch`. + +There is no built-in search engine. Give `websearch` a `SearchProvider`: + +```ts group=harness-coding-tools +import type { SearchProvider } from '@tanstack/ai-harness/plugins/coding' +import { searchTheWeb } from './search' + +const search: SearchProvider = { + search: (query, { limit, signal }) => searchTheWeb(query, { limit, signal }), +} + +export const withSearch = workspaceTools({ root, web: { search } }) +``` + +The provider gives at most `limit` results. Each result is `{ title, url, snippet? }`. The model reads a result with `webfetch`. + +## Permissions and approvals + +Each tool tells `permissions()` what kind of call it makes: + +| Tools | `default` mode | `plan` mode | +| --- | --- | --- | +| `read_file`, `list_files`, `grep`, `websearch` | Runs | Runs | +| `write_file`, `edit_file`, `patch` | Asks | Denied | +| `bash` | Asks | Denied | +| `webfetch` | Asks | Denied | + +- `acceptEdits` mode runs the edits without a question. `bash` still asks. +- `bypass` mode runs every call. +- `read_file` asks for a `.env` file. +- `grep` skips each file that `read_file` asks about or denies, also in `plan` mode. A note at the end gives the number of skipped files. +- `list_files` still shows the names of these files. A name is not the contents. + +Let `webfetch` run without a question with one rule: + +```ts group=harness-coding-tools +export const trustTheWeb = permissions({ + root, + rules: [{ tool: 'webfetch', decision: 'allow' }], +}) +``` + +In `plan` mode, a call that asks is denied. This rule makes `webfetch` an allow, so `webfetch` runs in `plan` mode too. + +Each call also tells `permissions()` the paths or the commands that it touches. So a rule with a `resource` can allow one folder or one command: + +```ts group=harness-coding-tools +export const trustTests = permissions({ + root, + rules: [ + { tool: 'edit_file', resource: 'src/**', decision: 'allow' }, + { tool: 'bash', resource: 'pnpm test*', decision: 'allow' }, + ], +}) +``` + +- Paths match from `root`. Give `permissions()` and `workspaceTools()` the same `root`. +- A rule also matches the real path of a link, and the stricter decision wins. So a link `notes` to `.env` asks, and a `deny` for `secrets/**` also stops a link to `secrets`. +- `bash` splits a command at `&&`, `||`, `;`, `|`, and new lines. Every part must be allowed. +- A command with `$(`, a backtick, or a heredoc never gets an allow. + +[Permissions](./permissions) explains the rules, the modes, and the saved answers. + +## Outside paths and the working folder + +### A path outside `root` + +The tools refuse a path outside `root`. To ask the user instead, set `outside: 'ask'`: + +```ts group=harness-coding-tools +export const askOutside = workspaceTools({ root, outside: 'ask' }) +``` + +- A yes allows that folder, and the folders in it, until the session ends. +- `bypass` mode allows every path without a question. +- A link in `root` that leads out of `root` counts as outside. +- With `permissions()`, `grep` skips the files outside `root`, because `read_file` asks for them. A `read_file` rule for that folder, for example `{ tool: 'read_file', resource: '/shared/**', decision: 'allow' }`, lets `grep` show them. + +### A working folder for each thread + +Each thread can work in its own folder in `root`. Set the `cwd` thread setting: + +```ts group=harness-coding-tools-cwd +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { openaiText } from '@tanstack/ai-openai' + +const repos = defineHarness({ + name: 'acme/repos', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [workspaceTools({ root: '/srv/repos' })], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(repos, { threadId: 'fix-login' }) +await session.configure({ cwd: 'web-app' }) +``` + +- Relative paths start in `/srv/repos/web-app`, and `bash` runs there. +- The system prompt tells the model the working folder. +- Permission rules still match paths from `root`, for example `web-app/src/**`. +- `cwd` is relative to `root`. A `cwd` outside `root` counts as an outside path. + +[Store settings per thread](./thread-settings) shows how a client sets `cwd`. + +## What you have now + +- An agent that reads, searches, and changes the files of one folder, and runs commands in it. +- A question before each edit, command, and web fetch. +- A working folder for each thread, and a question before a path outside `root`. + +Next: run the same tools on your own backend, or format each file after a write. See [Extend the coding tools](./coding-tools-backends). diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 9062b3f9bf..9f72b340a8 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -84,7 +84,7 @@ The `mode` setting decides how much runs without a question: | --- | --- | | `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. | +| `acceptEdits` | Edits run without a question. Commands and formatters still ask. | | `bypass` | Every call runs, also a call that a rule denies. | Switch the mode from server code: @@ -188,6 +188,49 @@ export const exampleEnv = permissions({ - 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. +## The rules of the coding tools + +`workspaceTools()` from `@tanstack/ai-harness/plugins/coding` adds these rules. [Give your agent coding tools](./coding-tools#permissions-and-approvals) shows what each tool does. + +| Tool | Decision | Kind | +| --- | --- | --- | +| `read_file`, `list_files`, `grep` | `allow` | `read` | +| `write_file`, `edit_file`, `patch` | `ask` | `edit` | +| `bash` | `ask` | `execute` | +| `webfetch` | `ask` | none | + +- `websearch` has no rule, so it runs. +- The path rules above also apply. For example, `read_file` asks for a `.env` file. +- The coding tools check the real path of a link too, and the stricter decision wins. So a link `notes` to `.env` asks, like `.env`. +- `grep` skips each file that `read_file` asks about or denies, also after an `always` answer. A note gives the number of skipped files. +- Your `rules` come after these rules, so they win. `{ tool: 'webfetch', decision: 'allow' }` lets `webfetch` run without a question. + +### Formatters + +A formatter runs code from the project, for example `prettier.config.js` or a package in `node_modules`. The model can edit these files. So [`formatter()`](./coding-tools-backends#format-each-file-after-a-write) asks before each run, like `bash`. + +It adds one rule: `{ tool: 'formatter:*', decision: 'ask', kind: 'execute' }`. Each run is named `formatter:` and the formatter name, for example `formatter:prettier`. + +| Mode | What the formatter does | +| --- | --- | +| `default` | It asks before it runs. | +| `acceptEdits` | It asks before it runs. The edit itself runs without a question. | +| `bypass` | It runs without a question. | +| `plan` | It does not run. No file changes in this mode. | + +- An `always` answer saves `formatter:prettier` for the project, so the user answers once. +- A `reject` skips the formatter for that write. The file stays as the tool wrote it, and the tool result says that the formatter did not run. +- Without `permissions()`, the formatter runs without a question. + +This rule lets prettier run without a question: + +```ts group=harness-permissions +export const trustedFormatter = permissions({ + root, + rules: [{ tool: 'formatter:prettier', decision: 'allow' }], +}) +``` + ## Answer a question A call that asks waits for one of three answers: diff --git a/docs/harness/snapshots.md b/docs/harness/snapshots.md new file mode 100644 index 0000000000..2c5fc79dd4 --- /dev/null +++ b/docs/harness/snapshots.md @@ -0,0 +1,128 @@ +--- +title: Undo a turn +id: harness-snapshots +order: 8 +description: "Let the user undo the file changes of the last agent turn with /undo, bring them back with /redo, and show what each model step changed." +keywords: + - tanstack ai + - harness + - coding agent + - snapshots + - undo + - diff +--- + +The agent changed five files, and the result is wrong. You want the files back as they were before that turn, and you do not want a reset in your own git repository. `snapshots()` saves the files of the workspace at each model step, in a git repository of its own. `/undo` puts the files back and removes the turn from the conversation. + +## 1. Add the plugin + +```ts group=harness-snapshots +import { homedir } from 'node:os' +import { join } from 'node:path' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { snapshots, workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { openaiText } from '@tanstack/ai-openai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' + +const root = process.cwd() +const history = snapshots({ + root, + dataDir: join(homedir(), '.acme-coder', 'snapshots'), +}) + +const coder = defineHarness({ + name: 'acme/coder', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [permissions({ root }), workspaceTools({ root }), history], +}) + +const log = memoryLogStore() +const host = createHarnessHost({ + persistence: { stores: { log, runs: memoryPersistence().stores.runs } }, +}) +const session = await host.open(coder, { threadId: 'fix-login' }) +``` + +- `root` is the workspace folder, an absolute path. +- `dataDir` is a folder for the snapshot repositories, an absolute path. Put it outside `root`. Each workspace gets its own repository in it. +- The plugin needs the `git` command. `/undo` and `/redo` need git 2.26 or later. +- Your project does not need to be a git repository, and the plugin never uses the `.git` folder of your project. + +## 2. Undo and redo a turn + +1. Send a prompt that changes files: `await session.prompt('Fix the login form.')`. +2. Run `await session.command('undo')`. It returns the number of files that it put back, for example `Undid the last turn. Files restored: 3.` +3. Look at the files that the turn changed: they are as they were before the turn. The turn is gone from the conversation too. +4. To bring back the files and the turn, run `await session.command('redo')`. It returns `Redid the last turn.` + +In a terminal app, the user types `/undo` and `/redo`. See [Build a coding agent](./coding-agent). + +- `/undo` puts back each file that the turn changed in `root`. This includes the edits of the file tools and the changes of `bash` commands. +- `/undo` and `/redo` do not touch other files. This is also true for a file that you add or edit after the turn. +- If you edit a file that the turn changed, `/undo` and `/redo` still write over it, and your edit is lost. +- The snapshots skip the files that `.gitignore` names, like `node_modules`. `/undo` does not touch them. +- A new turn clears `/redo`. + +In your own UI, add `expose: { commands: ['undo', 'redo'] }` to `defineHarness`. A client can run only the commands that the harness exposes. Then call `client.command('undo')`. The result text arrives as a `harness.command.result` event. + +## See what a step changed + +Each model step that changed files leaves a step: `{ from, to, files, toolCallIds }`. `from` and `to` are the ids of two snapshots. `files` lists the changed paths, and `toolCallIds` lists the tool calls of the step. `history.diff(from, to)` gives the changed files and a unified diff. + +On a [durable host](./durable-sessions), like the host above, the steps are `tanstack/snapshots:step` records in the session log: + +```ts group=harness-snapshots +await session.prompt('Fix the login form.') + +for (const { record } of await log.read('fix-login')) { + if (record.type !== 'tanstack/snapshots:step') continue + if (typeof record.from !== 'string' || typeof record.to !== 'string') continue + const { files, patch } = await history.diff(record.from, record.to) + console.log(files.join(', ')) + console.log(patch) +} +``` + +- Without `to`, `diff` compares with the files now. `history.diff(record.to)` shows what changed after that step. +- On a host without a log, the steps are a list in the plugin state, `snapshot().plugins['tanstack/snapshots']`. +- A step that changed no file is not saved. + +## Go back to an earlier message + +`/undo` goes back one turn. To go back more turns, revert the session to a message of the transcript: + +```ts group=harness-snapshots +const messages = await session.transcript() +const firstAnswer = messages.find((message) => message.role === 'assistant') +if (firstAnswer?.id) await session.revert(firstAnswer.id) + +// Bring back the files and the messages. +await session.unrevert() +``` + +- A revert touches only the files that the tool calls after that message changed. Each file goes back to its state before the first of these calls. A file that did not exist then is removed. +- Other files stay as they are. Your own edits between the hidden turns stay too. +- `unrevert()` touches only the same files. It puts them back as they were when the revert started, with your edits. +- A prompt that you send during `revert()` or `unrevert()` waits until it ends. Then the turn runs. +- `transcript()` hides the messages after that message until `unrevert()` or the next prompt. +- The next prompt drops the hidden messages for good. Then `unrevert()` does not change the files. +- The revert is saved, so `unrevert()` still works after a restart. + +In your own UI, call `client.revert(messageId)` and `client.unrevert()`. The client needs `'undo'` in `expose.commands`. More about reverts: [Go back to an earlier message](./sessions#go-back-to-an-earlier-message). + +## Known limits + +- One level of undo. A second `/undo` answers `Nothing to undo.` After a `/redo`, `/undo` works again. +- `/undo` and `/redo` work only when the session is idle. During a turn, they answer `Wait until the turn ends, then try again.` +- While a revert stands, `/undo` and `/redo` answer `A revert stands. Run unrevert, or send a prompt, first.` +- The undo point is kept in memory. After a restart, or when the session opens again, `/undo` works again after the next turn. +- If you sent a message while the turn ran, `/undo` removes the conversation from that message on. The files still go back to the start of the turn. +- Without `git`, the plugin prints `snapshots: git was not found. Snapshots are off.` once, and adds no commands. `diff()` throws. The turns run as usual. + +## What you have now + +- `/undo` and `/redo` for the file changes of the last turn, also the changes of shell commands. +- A revert to any earlier message, with its files, and `unrevert()` to bring them back. +- A list of changed files and a unified diff for each model step. +- Snapshots that never touch the git repository of your project. diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 1272c358e4..89e2ac9d58 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -33,6 +33,10 @@ "./plugins": { "types": "./dist/esm/first-party/index.d.ts", "import": "./dist/esm/first-party/index.js" + }, + "./plugins/coding": { + "types": "./dist/esm/first-party/coding/index.d.ts", + "import": "./dist/esm/first-party/coding/index.js" } }, "scripts": { @@ -50,11 +54,23 @@ "@tanstack/ai": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", + "@vscode/ripgrep": "^1.18.0", + "turndown": "^7.2.4", "zod": "^4.2.0" }, "peerDependencies": { "@tanstack/ai": "workspace:^", - "@tanstack/ai-persistence": "workspace:^" + "@tanstack/ai-persistence": "workspace:^", + "@vscode/ripgrep": "^1.15.0", + "turndown": "^7.2.0" + }, + "peerDependenciesMeta": { + "@vscode/ripgrep": { + "optional": true + }, + "turndown": { + "optional": true + } }, "nx": { "targets": { diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts new file mode 100644 index 0000000000..0ccb42321e --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -0,0 +1,360 @@ +import * as childProcess from 'node:child_process' +import * as fs from 'node:fs/promises' +import * as nodePath from 'node:path' +import { isRecord } from '../../utils' +import type { WorkspaceHooks } from '../workspace-hooks' + +/** + * Where the workspace tools read and write files and run commands. Every + * path is absolute. {@link hostBackend} is this machine. Another backend, for + * example a sandbox, lets the same tools work there. + */ +export interface WorkspaceBackend { + /** + * The shell that runs the commands: `'sh'` (a POSIX shell) or `'cmd'` + * (Windows `cmd.exe`). The tools quote arguments for it. It also sets the + * path style: POSIX paths for `'sh'`, Windows paths for `'cmd'`. Without + * it, arguments are quoted for `sh`, and paths have the style of this + * machine. A Linux sandbox sets `'sh'`, so its root `/workspace` stays a + * POSIX path on a Windows host. + */ + shell?: 'sh' | 'cmd' + /** The bytes of a file. Throws when there is no file. */ + readFile: (path: string) => Promise + /** Create or replace a file. Missing parent folders are created. */ + writeFile: (path: string, data: Uint8Array | string) => Promise + /** + * Remove one file. Throws when there is no file, or for a folder. Optional: + * without it, the `patch` tool cannot delete or move files. + */ + remove?: (path: string) => Promise + /** A file or a folder, or `undefined` when nothing is at `path`. */ + stat: ( + path: string, + ) => Promise< + { type: 'file' | 'dir'; size: number; mtimeMs: number } | undefined + > + /** + * The real path of `path`, with every link resolved, or `undefined` when + * nothing is at `path`. Throws for a link to a target that is not there. + * The tools use it to refuse a link that leads out of the workspace. + * Optional: without it, the tools do not check links. + */ + realpath?: (path: string) => Promise + /** + * The entries of a folder. A symbolic link (or a Windows junction) has + * the type `'link'`, whatever it points to. The tools do not follow it. + */ + readdir: ( + path: string, + ) => Promise> + /** + * Run a shell command and wait for it to end. `env` is added to the + * environment. A command that fails resolves with its exit code. + */ + exec: ( + command: string, + options?: { + cwd?: string + env?: Record + timeoutMs?: number + signal?: AbortSignal + }, + ) => Promise<{ exitCode: number; stdout: string; stderr: string }> + /** + * Start a shell command in the background. `output()` is stdout and + * stderr so far. Optional: without it, background commands give an error. + */ + spawn?: ( + command: string, + options?: { cwd?: string; env?: Record }, + ) => { + wait: () => Promise<{ exitCode: number }> + kill: () => void + output: () => string + } +} + +/** What each workspace tool gets from `workspaceTools()`. */ +export interface ToolEnv { + backend: WorkspaceBackend + /** The workspace folder, an absolute path. */ + root: string + /** The hooks that plugins added. Read them when the tool runs. */ + hooks: () => ReadonlyArray + /** Run `fn` when the work before it on `path` is done, one at a time. */ + lock: (path: string, fn: () => Promise) => Promise + /** + * The absolute path of `path`, a file or a folder. A path outside the + * workspace is refused, or needs the user's yes for its folder. The check + * uses the real path, so a link that leads out counts as outside. + */ + reach: ( + path: string, + tool: string, + kind: 'file' | 'folder', + ) => Promise + /** A path for the model: from the workspace, or the full path outside it. */ + shown: (full: string) => string + /** + * A test for the files in `folder`, a path from `reach`. It takes a path + * from `folder` with `/`, and is true when the permission rules protect + * the file: a `read_file` of it, by its path or its real path, would ask + * or be denied. `undefined` without `permissions()`. + */ + protectedIn?: ( + folder: string, + ) => Promise<((path: string) => boolean) | undefined> +} + +/** + * The path functions for the paths of `backend`: POSIX for a `'sh'` + * backend, Windows for a `'cmd'` backend. Without `shell`, the functions of + * this machine. + */ +export function pathsOf(backend: WorkspaceBackend) { + switch (backend.shell) { + case 'sh': + return nodePath.posix + case 'cmd': + return nodePath.win32 + case undefined: + return nodePath + } +} + +/** + * The environment of a command: the environment of this process, plus + * `env`. Without `env` off Windows: `undefined`, so the command gets the + * environment of this process. On Windows it sets + * `NoDefaultCurrentDirectoryInExePath`. Then `cmd.exe` takes a program + * name like `rg` from the PATH only, and not from the current folder, where + * a cloned repo can put an `rg.cmd`. + */ +export function commandEnv( + env: Record | undefined, + platform: NodeJS.Platform, +) { + if (platform !== 'win32') return env && { ...process.env, ...env } + return { ...process.env, ...env, NoDefaultCurrentDirectoryInExePath: '1' } +} + +/** True for the error of a path that is not there. */ +const isMissing = (error: unknown) => + error instanceof Error && 'code' in error && error.code === 'ENOENT' + +/** `exec` stops a command that writes more than 10 MiB of output. */ +const MAX_EXEC_OUTPUT = 10 * 1024 * 1024 + +/** + * Start `command` in the system shell. `kill` stops the shell and every + * command it started: on Windows, a killed `cmd.exe` leaves its commands + * running. + */ +function startShell( + command: string, + options: { cwd?: string; env?: Record }, +) { + const isWindows = process.platform === 'win32' + // Off Windows, the command gets its own process group, so `kill` can + // stop the commands it started too. + const child = childProcess.spawn(command, { + cwd: options.cwd, + env: commandEnv(options.env, process.platform), + shell: true, + detached: !isWindows, + }) + const kill = () => { + if (child.pid === undefined) return + if (isWindows) { + // The full path: Node looks for a bare name in the current folder + // first, and that folder can be a cloned repo. + const system = nodePath.join( + process.env.SystemRoot ?? 'C:\\Windows', + 'System32', + ) + childProcess.execFile( + nodePath.join(system, 'taskkill.exe'), + ['/pid', String(child.pid), '/T', '/F'], + () => undefined, + ) + return + } + const group = -child.pid + try { + process.kill(group, 'SIGTERM') + } catch { + // The command already ended. + return + } + // A command that ignores SIGTERM gets SIGKILL after 2 seconds. + setTimeout(() => { + try { + process.kill(group, 'SIGKILL') + } catch { + // The command ended. + } + }, 2000).unref() + } + return { child, kill } +} + +/** + * The workspace on this machine, with `node:fs` and `node:child_process`. + * Commands run in the system shell. `exec` gives exit code 124, like GNU + * `timeout`, when the timeout stops a command or the output is over 10 MiB. + * It gives 1 when the signal stops a command, or the command did not + * start. A stopped command is stopped with every command it started. + */ +export const hostBackend = { + shell: process.platform === 'win32' ? 'cmd' : 'sh', + readFile: (path) => fs.readFile(path), + writeFile: async (path, data) => { + await fs.mkdir(nodePath.dirname(path), { recursive: true }) + await fs.writeFile(path, data) + }, + // Without `recursive`, `fs.rm` refuses a folder. + remove: (path) => fs.rm(path), + stat: async (path) => { + const info = await fs.stat(path).catch((error: unknown) => { + if (isMissing(error)) return undefined + throw error + }) + if (!info) return undefined + return { + type: info.isDirectory() ? 'dir' : 'file', + size: info.size, + mtimeMs: info.mtimeMs, + } + }, + realpath: async (path) => { + try { + return await fs.realpath(path) + } catch (error) { + // A link to a missing target is there, so it is not `undefined`. + const isThere = await fs.lstat(path).then( + () => true, + () => false, + ) + if (isMissing(error) && !isThere) return undefined + throw error + } + }, + readdir: async (path) => { + const entries = await fs.readdir(path, { withFileTypes: true }) + return entries.map((entry) => ({ + name: entry.name, + type: entry.isSymbolicLink() + ? ('link' as const) + : entry.isDirectory() + ? ('dir' as const) + : ('file' as const), + })) + }, + exec: (command, options = {}) => + new Promise((done) => { + const { child, kill } = startShell(command, options) + const stdout: Array = [] + const stderr: Array = [] + let size = 0 + // The exit code of a command that was stopped. The first reason wins. + let stoppedWith: number | undefined + const stop = (exitCode: number) => { + stoppedWith ??= exitCode + kill() + } + const collect = (chunks: Array) => (chunk: Buffer) => { + // Over the limit, the command stops. Drop what it still writes. + if (size > MAX_EXEC_OUTPUT) return + chunks.push(chunk) + size += chunk.length + if (size > MAX_EXEC_OUTPUT) stop(124) + } + child.stdout.on('data', collect(stdout)) + child.stderr.on('data', collect(stderr)) + // A timeout of 0 means no timeout. + const timer = options.timeoutMs + ? setTimeout(() => stop(124), options.timeoutMs) + : undefined + const onAbort = () => stop(1) + options.signal?.addEventListener('abort', onAbort) + if (options.signal?.aborted) onAbort() + let startError = '' + // Runs on `close`, and on `error` when the command did not start. + const finish = (code: number | null) => { + clearTimeout(timer) + options.signal?.removeEventListener('abort', onAbort) + done({ + exitCode: stoppedWith ?? code ?? 1, + stdout: Buffer.concat(stdout).toString(), + stderr: Buffer.concat(stderr).toString() + startError, + }) + } + child.on('close', finish) + child.on('error', (error) => { + startError = String(error) + finish(1) + }) + }), + spawn: (command, options = {}) => { + const { child, kill } = startShell(command, options) + // ponytail: keeps all output in memory. Cap it if long jobs need that. + let output = '' + const collect = (chunk: Buffer) => { + output += chunk.toString() + } + child.stdout.on('data', collect) + child.stderr.on('data', collect) + const exited = new Promise<{ exitCode: number }>((done) => { + child.on('close', (code) => done({ exitCode: code ?? 1 })) + child.on('error', (error) => { + output += String(error) + done({ exitCode: 1 }) + }) + }) + return { wait: () => exited, kill, output: () => output } + }, +} satisfies WorkspaceBackend + +const MAX_OUTPUT = 20_000 + +/** Cut text after 20,000 characters, and say how many were cut. */ +export function clip(text: string) { + return text.length > MAX_OUTPUT + ? `${text.slice(0, MAX_OUTPUT)}\n[${text.length - MAX_OUTPUT} more characters]` + : text +} + +/** The string argument `key` of a tool call. Throws when it is not a string. */ +export function stringArg(args: unknown, key: string) { + const value = optionalString(args, key) + if (value !== undefined) return value + throw new Error(`Argument "${key}" must be a string.`) +} + +/** The string argument `key` of a tool call, or `undefined`. */ +export function optionalString(args: unknown, key: string) { + const value = isRecord(args) ? args[key] : undefined + return typeof value === 'string' ? value : undefined +} + +// Keep a byte order mark, so an edit writes the file back as it was. +const decoder = new TextDecoder('utf-8', { ignoreBOM: true }) + +/** The text of the file at `path`, as UTF-8. */ +export async function readText(backend: WorkspaceBackend, path: string) { + return decoder.decode(await backend.readFile(path)) +} + +/** + * Run the `afterWrite` hooks for `path`, one after another. Resolves to the + * texts that the hooks return, for the tool result. + */ +export async function afterWrite(env: ToolEnv, path: string) { + const notes: Array = [] + for (const hook of env.hooks()) { + const note = await hook.afterWrite?.(path) + if (note) notes.push(note) + } + return notes +} diff --git a/packages/ai-harness/src/first-party/coding/bash.ts b/packages/ai-harness/src/first-party/coding/bash.ts new file mode 100644 index 0000000000..00709e4015 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/bash.ts @@ -0,0 +1,259 @@ +import { toolDefinition } from '@tanstack/ai' +import { isRecord } from '../../utils' +import { boundText } from '../bound-output' +import { pathsOf, stringArg } from './backend' +import type { ToolEnv } from './backend' + +/** The longest timeout that one call can ask for: 10 minutes. */ +const MAX_TIMEOUT_MS = 10 * 60 * 1000 +/** + * How much output the model gets: the last lines. The limits are below the + * defaults of `boundToolOutput` (2000 lines, 50 KiB), so that plugin does + * not cut the result again and remove the exit code. + */ +const OUTPUT_LIMITS = { + maxLines: 1000, + maxBytes: 20 * 1024, + keep: 'tail', +} as const +/** Added to the environment of every command, so tools can see an agent. */ +const AGENT_ENV = { AGENT: '1' } +/** Shell syntax that can hide a command inside another command. */ +const HIDDEN_COMMAND = /\$\(|`|<\(|>\(|<(`, a heredoc (`<<`), an + * unquoted `(`, or a quote that does not close, the whole command comes + * back as one part. The permission rules never allow a part that still has + * shell syntax in it. + * + * @example + * ```ts + * splitCommand('ls && rm -rf x') // ['ls', 'rm -rf x'] + * splitCommand('echo $(rm x)') // ['echo $(rm x)'] + * ``` + */ +export function splitCommand(command: string) { + if (HIDDEN_COMMAND.test(command)) return [command] + const parts: Array = [] + let part = '' + // The open quote. In `'...'` every character is as is. In `"..."` and + // `$'...'` a backslash escapes the next character. + let quote: "'" | '"' | "$'" | undefined + for (let i = 0; i < command.length; i++) { + const char = command.charAt(i) + if (quote === "'") { + if (char === "'") quote = undefined + part += char + continue + } + if (char === '\\') { + part += command.slice(i, i + 2) + i++ + continue + } + if (quote !== undefined) { + const closes = char === (quote === '"' ? '"' : "'") + if (closes) quote = undefined + part += char + continue + } + const operator = OPERATORS.find((item) => command.startsWith(item, i)) + if (operator !== undefined) { + parts.push(part) + part = '' + i += operator.length - 1 + continue + } + // A subshell or a function body, for example `ls() ( rm -rf x ); ls`. + if (char === '(') return [command] + if (command.startsWith("$'", i)) { + quote = "$'" + part += "$'" + i++ + continue + } + if (char === "'" || char === '"') quote = char + part += char + } + if (quote !== undefined) return [command] + parts.push(part) + // Remove only spaces and tabs, the blanks of the shell. A `\r` stays, so + // the permission rules see it. + return parts + .map((item) => item.replace(/^[ \t]+|[ \t]+$/g, '')) + .filter((item) => item !== '') +} + +/** + * The command parts of a `bash` call, for `PermissionResources`. Throws when + * `command` is not a string, and that refuses the call. + * + * @example + * ```ts + * PermissionResources.item({ bash: { commands: bashResources } }) + * ``` + */ +export function bashResources(input: unknown) { + return splitCommand(stringArg(input, 'command')) +} + +/** The number of background jobs started, by the signal that stops them. */ +const jobCounts = new WeakMap() + +/** + * `bash`: run a shell command in the workspace folder. Every command gets + * `AGENT=1` in its environment. The model gets the exit code and the last + * part of the output. + * + * With `background: true`, the call returns at once with a job id, and + * `note` tells the model when the job ends. Background jobs need + * `env.backend.spawn`. They are killed when `signal` aborts. + * + * A command that runs in the foreground supports `detach` of the tool + * context: the host can move it to the background while it runs. + */ +export function bashTools( + env: ToolEnv, + options: { + /** Stop the command after this many milliseconds. Default: 120,000. */ + timeoutMs?: number + /** Tells the model that a background job ended. Without it, no one is told. */ + note?: (text: string) => Promise + /** + * When the model gets only the end of the output, the full output is + * saved in this folder, and the model gets the path. Relative to the + * workspace, or absolute. + */ + spillDir?: string + /** Kills the background jobs that still run, for example on plugin cleanup. */ + signal?: AbortSignal + /** Told when a background job starts and ends. A durable session logs it. */ + jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } + } = {}, +) { + const { + timeoutMs: defaultTimeoutMs = 120_000, + note, + spillDir, + signal, + } = options + // A restart cannot reach these processes. `options.jobs` lets a durable + // session note the jobs that a crash stopped. + const jobs = new Set<{ kill: () => void }>() + // The count belongs to the signal. A session signal outlives a reload, so a + // job after a reload does not reuse the id of a job that still runs. + const count = (signal && jobCounts.get(signal)) || { started: 0 } + if (signal) jobCounts.set(signal, count) + signal?.addEventListener('abort', () => { + for (const job of jobs) job.kill() + }) + + /** Save the full output in `dir`. Gives back the note for the model. */ + const spill = async (dir: string, output: string) => { + const { join, resolve } = pathsOf(env.backend) + const path = join(resolve(env.root, dir), `bash-${crypto.randomUUID()}.txt`) + try { + await env.backend.writeFile(path, output) + return `[Full output saved to ${env.shown(path)}.]` + } catch (error) { + // The cut output is still useful, so a failed save does not fail the call. + return `[The full output was not saved: ${String(error)}]` + } + } + + /** The exit code and the end of `output`, as the model gets them. */ + const report = async (exitCode: number, output: string) => { + const bounded = boundText(output, OUTPUT_LIMITS) + const lines = [`exit code: ${exitCode}`, bounded.text] + if (bounded.truncated && spillDir !== undefined) { + lines.push(await spill(spillDir, output)) + } + return lines.join('\n') + } + + /** Start `command` in the background. `note` tells the model when it ends. */ + const startJob = (command: string, cwd: string) => { + if (!env.backend.spawn) { + throw new Error('This workspace cannot run commands in the background.') + } + const job = env.backend.spawn(command, { cwd, env: AGENT_ENV }) + count.started += 1 + const jobId = `bash-${count.started}` + jobs.add(job) + options.jobs?.started(jobId) + const tell = async () => { + const { exitCode } = await job.wait() + jobs.delete(job) + // A closed session keeps the job as running in the log, so the next + // host tells the model that it stopped. + if (signal?.aborted) return + try { + if (note === undefined) return + const result = await report(exitCode, job.output()) + await note(`Background job ${jobId} ended.\n${result}`) + } finally { + options.jobs?.ended(jobId) + } + } + // No one waits for the job, so a failed note must not crash the process. + tell().catch(() => undefined) + return { jobId, status: 'started' } + } + + return [ + toolDefinition({ + name: 'bash', + description: + 'Run a shell command in the workspace folder. You get the exit code and the last 1000 lines (20 KiB) of the output. For a command that runs long, like a dev server, set background to true: the call returns at once, and you get a note when the command ends.', + inputSchema: { + type: 'object', + properties: { + command: { type: 'string' }, + timeoutMs: { + type: 'number', + description: + 'Stop the command after this many milliseconds. At most 600000.', + }, + background: { + type: 'boolean', + description: 'Run the command in the background.', + }, + }, + required: ['command'], + }, + replay: 'never', + }).server(async (args: unknown, context) => { + const command = stringArg(args, 'command') + // The shell runs in the working folder, so it must be in the workspace. + const cwd = await env.reach('.', 'bash', 'folder') + if (isRecord(args) && args.background === true) { + return startJob(command, cwd) + } + const requested = isRecord(args) ? args.timeoutMs : undefined + // A timeout of 0 would mean no timeout, so it gets the default. + const isTimeout = typeof requested === 'number' && requested > 0 + const work = env.backend + .exec(command, { + cwd, + env: AGENT_ENV, + timeoutMs: isTimeout + ? Math.min(requested, MAX_TIMEOUT_MS) + : defaultTimeoutMs, + }) + .then(({ exitCode, stdout, stderr }) => + report(exitCode, `${stdout}${stderr ? `\nstderr:\n${stderr}` : ''}`), + ) + // The host can move the command to the background. The timeout still + // applies to it. + const detach = context?.detach + return detach ? Promise.race([work, detach(work)]) : work + }), + ] +} diff --git a/packages/ai-harness/src/first-party/coding/edit.ts b/packages/ai-harness/src/first-party/coding/edit.ts new file mode 100644 index 0000000000..1dba9f5180 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/edit.ts @@ -0,0 +1,190 @@ +import { toolDefinition } from '@tanstack/ai' +import { isRecord } from '../../utils' +import { afterWrite, readText, stringArg } from './backend' +import type { ToolEnv } from './backend' + +const BOM = '' + +/** + * How one match level changes text before the search. `pieces` cuts the text + * into pieces, and `change` gives the new text of each piece. + */ +interface Level { + pieces: RegExp + change: (piece: RegExpExecArray) => string +} + +// The match levels after the exact text, from strict to loose. +const LEVELS: ReadonlyArray = [ + // Unicode NFC form. A piece is a character and the marks after it. + // ponytail: decomposed Hangul jamo stay apart. Cut by grapheme if needed. + { + pieces: /\P{M}\p{M}*|\p{M}+/gu, + change: (piece) => piece[0].normalize('NFC'), + }, + // Spaces and tabs at the end of each line removed. + { + pieces: /([^\S\n]+(?=\n|$))|[\s\S]/g, + change: (piece) => (piece[1] ? '' : piece[0]), + }, +] + +/** + * `text` changed by `level`, and where each character of the result starts in + * `text`. `at[i]` is `undefined` inside a changed piece. The last entry of + * `at` is the end of `text`. + */ +function changed(text: string, level: Level) { + let result = '' + const at: Array = [] + const pieces = text.matchAll(level.pieces) + for (const piece of pieces) { + const next = level.change(piece) + // An unchanged piece maps one to one. A changed piece maps its start only. + const same = next === piece[0] + for (let i = 0; i < next.length; i++) { + at.push(same || i === 0 ? piece.index + i : undefined) + } + result += next + } + at.push(text.length) + return { text: result, at } +} + +/** Where `find` starts in `text`, without overlaps. None for an empty `find`. */ +function indexesOf(text: string, find: string) { + const found: Array = [] + if (find === '') return found + let index = text.indexOf(find) + while (index !== -1) { + found.push(index) + index = text.indexOf(find, index + find.length) + } + return found +} + +/** `text` with the line endings of `content`: CRLF when it has one, else LF. */ +export function toEndingsOf(text: string, content: string) { + return text.replace(/\r?\n/g, content.includes('\r\n') ? '\r\n' : '\n') +} + +/** + * Find `find` in `content`. Three levels are tried in order, and the first + * level with a match wins: + * + * 1. The exact text. + * 2. Both sides in Unicode NFC form. + * 3. Spaces and tabs at the end of each line ignored. A match then also covers + * the trailing spaces of its last line. + * + * `find` first gets the line endings of `content`, so LF text matches a CRLF + * file. + * + * Returns each match as `{ start, end }` offsets into the original `content`, + * in order and without overlaps. An empty array means no match. More than one + * match means the text is not unique. + * + * @example + * ```ts + * findMatches('a \nb\n', 'a\nb') // [{ start: 0, end: 5 }] + * ``` + */ +export function findMatches(content: string, find: string) { + const wanted = toEndingsOf(find, content) + const exact = indexesOf(content, wanted) + if (exact.length > 0) { + return exact.map((start) => ({ start, end: start + wanted.length })) + } + for (const level of LEVELS) { + const { text, at } = changed(content, level) + const needle = changed(wanted, level).text + const starts = indexesOf(text, needle) + const matches = starts.flatMap((index) => { + const start = at[index] + const end = at[index + needle.length] + // A match must start and end where a piece starts. + return start === undefined || end === undefined ? [] : [{ start, end }] + }) + if (matches.length > 0) return matches + } + return [] +} + +/** + * `write_file` and `edit_file`. After each write, the `afterWrite` hooks run + * for the file. `edit_file` finds the old text with {@link findMatches}. It + * writes the new text with the line endings of the file, and keeps a byte + * order mark. + */ +export function editTools(env: ToolEnv) { + return [ + toolDefinition({ + name: 'write_file', + description: 'Create or replace a file in the workspace.', + inputSchema: { + type: 'object', + properties: { path: { type: 'string' }, content: { type: 'string' } }, + required: ['path', 'content'], + }, + }).server(async (args: unknown) => { + const full = await env.reach( + stringArg(args, 'path'), + 'write_file', + 'file', + ) + const content = stringArg(args, 'content') + return env.lock(full, async () => { + await env.backend.writeFile(full, content) + const notes = await afterWrite(env, full) + return [`Wrote ${env.shown(full)}.`, ...notes].join('\n') + }) + }), + toolDefinition({ + name: 'edit_file', + description: + 'Replace exact text in a file. `old` must appear once, unless `replaceAll` is true.', + inputSchema: { + type: 'object', + properties: { + path: { type: 'string' }, + old: { type: 'string' }, + new: { type: 'string' }, + replaceAll: { type: 'boolean' }, + }, + required: ['path', 'old', 'new'], + }, + }).server(async (args: unknown) => { + const full = await env.reach(stringArg(args, 'path'), 'edit_file', 'file') + const oldText = stringArg(args, 'old') + const newText = stringArg(args, 'new') + const replaceAll = isRecord(args) && args.replaceAll === true + return env.lock(full, async () => { + const before = await readText(env.backend, full) + const matches = findMatches(before, oldText) + const count = matches.length + if (count === 0) throw new Error('The old text is not in the file.') + if (count > 1 && !replaceAll) { + throw new Error( + `The old text appears ${count} times. Add context, or set replaceAll.`, + ) + } + const replacement = toEndingsOf(newText, before) + let after = '' + let from = 0 + for (const { start, end } of matches) { + after += before.slice(from, start) + replacement + from = end + } + after += before.slice(from) + // Keep the byte order mark, also when the new text drops it. + const keepBom = before.startsWith(BOM) && !after.startsWith(BOM) + await env.backend.writeFile(full, keepBom ? BOM + after : after) + const notes = await afterWrite(env, full) + return [ + `Edited ${env.shown(full)} (${count} change${count > 1 ? 's' : ''}).`, + ...notes, + ].join('\n') + }) + }), + ] +} diff --git a/packages/ai-harness/src/first-party/coding/formatter.ts b/packages/ai-harness/src/first-party/coding/formatter.ts new file mode 100644 index 0000000000..c7b3aafc75 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/formatter.ts @@ -0,0 +1,272 @@ +import { createPluginEvent } from '../../extensions' +import { definePlugin } from '../../plugins' +import { isRecord } from '../../utils' +import { PermissionPrompt, PermissionRules } from '../permissions' +import { WorkspaceHooks } from '../workspace-hooks' +import { hostBackend, pathsOf, readText } from './backend' +import { quoteArg, shellPlatform } from './search' +import type { WorkspaceBackend } from './backend' + +/** What a formatter can check in the project root folder. */ +export interface FormatterProject { + /** The names of the files and folders in the root folder. */ + files: ReadonlyArray + /** The text of the file `name` in the root folder, or `''` without one. */ + read: (name: string) => Promise +} + +/** A formatter CLI for some file types. */ +export interface Formatter { + /** A short name, for the warning when it fails. */ + name: string + /** The file extensions it formats, with the dot, like `'.ts'`. */ + extensions: ReadonlyArray + /** + * The shell command that formats `file` in place. `file` is the absolute + * path, already quoted for the shell. Put it in the command as it is. + */ + command: (file: string) => string + /** Does the project use this formatter? Default: yes. */ + when?: (project: FormatterProject) => boolean | Promise +} + +export interface FormatterOptions { + /** The workspace folder. Commands run in it. */ + root: string + /** Where the files are and commands run. Default: {@link hostBackend}. */ + backend?: WorkspaceBackend + /** Your formatters. They come before the built-in formatters. */ + formatters?: ReadonlyArray + /** `false` turns the built-in formatters off. Default: `true`. */ + builtins?: boolean + /** The time limit of each run. Default: 20 seconds. */ + timeoutMs?: number +} + +/** A formatter run that failed. The written file stays as it was. */ +export interface FormatFailure { + /** The absolute path of the file. */ + path: string + message: string +} + +/** Sent when a formatter fails. The tool call does not fail. */ +export const FormatFailed = createPluginEvent( + 'tanstack/formatter:failed', +) + +const SCRIPT = ['.js', '.jsx', '.mjs', '.cjs', '.ts', '.tsx', '.mts', '.cts'] +/** The file types that prettier and oxfmt format. */ +const WEB = [ + ...SCRIPT, + ...['.json', '.jsonc', '.css', '.scss', '.less', '.html', '.vue'], + ...['.md', '.mdx', '.yaml', '.yml', '.graphql'], +] + +/** The field `key` of the project's `package.json`, or `undefined`. */ +async function packageField(project: FormatterProject, key: string) { + try { + const value: unknown = JSON.parse(await project.read('package.json')) + return isRecord(value) ? value[key] : undefined + } catch { + return undefined + } +} + +/** Does a name in the root folder start with one of the `prefixes`? */ +const hasFile = (project: FormatterProject, ...prefixes: Array) => + project.files.some((name) => + prefixes.some((prefix) => name.startsWith(prefix)), + ) + +// ponytail: six built-in formatters. Add more with `formatters`. +const BUILTINS: ReadonlyArray = [ + { + name: 'prettier', + extensions: WEB, + command: (file) => `npx --no-install prettier --write ${file}`, + when: async (project) => + hasFile(project, '.prettierrc', 'prettier.config.') || + (await packageField(project, 'prettier')) !== undefined, + }, + { + name: 'biome', + extensions: [...SCRIPT, '.json', '.jsonc', '.css', '.graphql'], + command: (file) => `npx --no-install biome format --write ${file}`, + when: (project) => + project.files.includes('biome.json') || + project.files.includes('biome.jsonc'), + }, + { + name: 'oxfmt', + extensions: WEB, + command: (file) => `npx --no-install oxfmt ${file}`, + when: async (project) => { + const devDependencies = await packageField(project, 'devDependencies') + return ( + hasFile(project, '.oxfmtrc') || + (isRecord(devDependencies) && 'oxfmt' in devDependencies) + ) + }, + }, + { + name: 'ruff', + extensions: ['.py', '.pyi'], + command: (file) => `ruff format ${file}`, + when: async (project) => + project.files.includes('ruff.toml') || + project.files.includes('.ruff.toml') || + (await project.read('pyproject.toml')).includes('[tool.ruff'), + }, + { + name: 'gofmt', + extensions: ['.go'], + command: (file) => `gofmt -w ${file}`, + when: (project) => project.files.includes('go.mod'), + }, + { + name: 'rustfmt', + extensions: ['.rs'], + command: (file) => `rustfmt ${file}`, + when: (project) => project.files.includes('Cargo.toml'), + }, +] + +const messageOf = (error: unknown) => + error instanceof Error ? error.message : String(error) + +/** + * The `afterWrite` hook of {@link formatter}. It runs the first formatter + * that the project uses for the file type. The project is checked at the + * first write, once. A failure goes to `onFailure` and never throws. + * `allow` says if the formatter `name` can run on `path`. When it says no, + * the hook resolves to a note for the tool result. + */ +export function formatOnWrite( + options: FormatterOptions, + onFailure: (failure: FormatFailure) => void, + allow?: (name: string, path: string) => Promise, +) { + const backend = options.backend ?? hostBackend + const { extname, join, resolve } = pathsOf(backend) + const root = resolve(options.root) + const timeoutMs = options.timeoutMs ?? 20_000 + const platform = shellPlatform(backend) + const candidates = [ + ...(options.formatters ?? []), + ...(options.builtins === false ? [] : BUILTINS), + ] + + /** The formatters that the project uses, in order. */ + const detect = async () => { + const entries = await backend.readdir(root) + const project: FormatterProject = { + files: entries.map((entry) => entry.name), + read: (name) => readText(backend, join(root, name)).catch(() => ''), + } + const used = await Promise.all( + candidates.map( + async (candidate) => (await candidate.when?.(project)) ?? true, + ), + ) + return candidates.filter((_, index) => used[index]) + } + let detected: ReturnType | undefined + + const format = async (path: string) => { + const extension = extname(path).toLowerCase() + const formatters = await (detected ??= detect()) + const match = formatters.find((candidate) => + candidate.extensions.includes(extension), + ) + if (!match) return + if (allow && !(await allow(match.name, path))) { + return `The formatter ${match.name} did not run: it is not allowed.` + } + const command = match.command(quoteArg(path, platform)) + const result = await backend.exec(command, { cwd: root, timeoutMs }) + if (result.exitCode === 0) return undefined + const reason = + result.exitCode === 124 + ? `it took longer than ${timeoutMs / 1000} seconds` + : result.stderr.trim().split('\n')[0] || `exit code ${result.exitCode}` + throw new Error(`${match.name} failed: ${reason}`) + } + + return (path: string) => + format(path).catch((error: unknown) => + onFailure({ path, message: messageOf(error) }), + ) +} + +/** + * Format each file that `write_file`, `edit_file`, and `patch` write, with + * the formatter that the project uses for its file type. Use it with + * `workspaceTools`. + * + * At the first write in a session, the plugin checks the config files in + * `root`, and keeps the result: + * + * - prettier: `.prettierrc*`, `prettier.config.*`, or `prettier` in + * `package.json`. + * - biome: `biome.json` or `biome.jsonc`. + * - oxfmt: `.oxfmtrc*`, or `oxfmt` in the `devDependencies`. + * - ruff (`.py`): `ruff.toml`, or `[tool.ruff]` in `pyproject.toml`. + * - gofmt (`.go`): `go.mod`. rustfmt (`.rs`): `Cargo.toml`. + * + * Your `formatters` come first. `builtins: false` turns the list above off. + * Each run is a command through `backend.exec` in `root`, with a time limit + * (`timeoutMs`, default 20 seconds). A formatter that fails, times out, or + * is not installed does not fail the tool call: the file stays as written, + * and the plugin sends a {@link FormatFailed} event. + * + * A formatter runs code from the project, like its config files. So with + * `permissions()`, each run is a command named `formatter:`, for + * example `formatter:prettier`. It asks like `bash`, also in `acceptEdits` + * mode. A no skips the formatter, keeps the write, and the tool result says + * so. Without `permissions()`, the formatter runs without a question. + * + * @example + * ```ts + * plugins: () => [ + * workspaceTools({ root }), + * formatter({ + * root, + * formatters: [ + * { name: 'taplo', extensions: ['.toml'], command: (file) => `taplo fmt ${file}` }, + * ], + * }), + * ] + * ``` + */ +export function formatter(options: FormatterOptions) { + return definePlugin({ + name: 'tanstack/formatter', + optionalRequires: [PermissionPrompt], + setup: (ctx) => ({ + contribute: [ + WorkspaceHooks.item({ + afterWrite: formatOnWrite( + options, + (failure) => ctx.emit(FormatFailed, failure), + // Read at each write, so permissions() can come after this plugin. + async (name, path) => { + const ask = ctx.getOptional(PermissionPrompt) + return ask + ? ask( + `formatter:${name}`, + `Allow formatter:${name} on ${path}? It runs code from the project.`, + ) + : true + }, + ), + }), + PermissionRules.item({ + tool: 'formatter:*', + decision: 'ask', + kind: 'execute', + }), + ], + }), + }) +} diff --git a/packages/ai-harness/src/first-party/coding/index.ts b/packages/ai-harness/src/first-party/coding/index.ts new file mode 100644 index 0000000000..a8adcfdcfe --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/index.ts @@ -0,0 +1,12 @@ +// First-party plugins for a coding agent. This entry is Node only: the +// plugins use the file system and the shell. +export { workspaceTools } from './workspace' +export type { WorkspaceToolsOptions } from './workspace' +export { hostBackend } from './backend' +export type { WorkspaceBackend } from './backend' +export { WorkspaceHooks } from '../workspace-hooks' +export type { SearchProvider, WebToolsOptions } from './web' +export { snapshots } from './snapshots' +export type { SnapshotStep, SnapshotsOptions } from './snapshots' +export { FormatFailed, formatter } from './formatter' +export type { Formatter, FormatterOptions, FormatterProject } from './formatter' diff --git a/packages/ai-harness/src/first-party/coding/patch.ts b/packages/ai-harness/src/first-party/coding/patch.ts new file mode 100644 index 0000000000..c617459f14 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/patch.ts @@ -0,0 +1,450 @@ +import { toolDefinition } from '@tanstack/ai' +import { afterWrite, readText, stringArg } from './backend' +import { findMatches, toEndingsOf } from './edit' +import type { ToolEnv } from './backend' + +const BOM = '' + +/** One chunk of an `*** Update File:` section. */ +export interface PatchChunk { + /** The text after `@@`: a line above the change, to find the place. */ + context?: string + /** The lines that the chunk replaces: the context and `-` lines. */ + oldLines: Array + /** The lines that replace them: the context and `+` lines. */ + newLines: Array + /** True after `*** End of File`: the old lines end the file. */ + endOfFile: boolean +} + +/** One file change of a patch. The paths are as the patch gives them. */ +export type PatchOperation = + | { type: 'add'; path: string; content: string } + | { type: 'delete'; path: string } + | { + type: 'update' + path: string + movePath?: string + chunks: Array + } + +/** The file changes of a patch, in order. */ +export type ParsedPatch = Array + +const DESCRIPTION = `Change files with a patch: + +*** Begin Patch +*** Add File: src/new.ts ++each line of a new file starts with + +*** Update File: src/app.ts +@@ function main() { + a context line starts with a space +-a removed line starts with - ++an added line starts with + +*** Delete File: src/old.ts +*** End Patch + +In an update, each chunk starts with @@. The text after @@ is optional. It is a line above the change, to find the place. Give 3 context lines above and below each change, so the old lines are unique. Chunks apply in order. "*** Move to: " right after "*** Update File:" renames the file. "*** End of File" after a chunk means that the chunk ends the file.` + +function lineError(index: number, message: string) { + return new Error(`Patch line ${index + 1}: ${message}`) +} + +/** The path after a header. Throws when it is empty. */ +function pathAt(index: number, rest = '') { + const path = rest.trim() + if (path === '') throw lineError(index, 'the path is missing.') + return path +} + +/** Throws when `operation` is an update with no lines and no move. */ +function checkNotEmpty(operation: PatchOperation | undefined, index: number) { + if (operation?.type !== 'update' || operation.movePath !== undefined) return + const hasLines = operation.chunks.some( + (chunk) => chunk.oldLines.length + chunk.newLines.length > 0, + ) + if (!hasLines) { + throw lineError( + index, + `"*** Update File: ${operation.path}" has no changes.`, + ) + } +} + +function newChunk(context: string) { + const chunk: PatchChunk = { + context: context || undefined, + oldLines: [], + newLines: [], + endOfFile: false, + } + return chunk +} + +const isFilled = (line: string) => line.trim() !== '' + +/** + * Read a patch in the `*** Begin Patch` format into its file changes. Throws + * an error with the line number when the text does not follow the format. + * + * @example + * ```ts + * parsePatch('*** Begin Patch\n*** Delete File: old.txt\n*** End Patch') + * // [{ type: 'delete', path: 'old.txt' }] + * ``` + */ +export function parsePatch(text: string) { + const lines = text.split(/\r?\n/) + const first = Math.max(lines.findIndex(isFilled), 0) + const last = lines.findLastIndex(isFilled) + if (lines[first]?.trim() !== '*** Begin Patch') { + throw lineError(first, 'the patch must start with "*** Begin Patch".') + } + if (last === first || lines[last]?.trim() !== '*** End Patch') { + throw lineError(last, 'the patch must end with "*** End Patch".') + } + const operations: ParsedPatch = [] + let current: PatchOperation | undefined + let currentLine = first + // The chunk that new lines go to. + let chunk: PatchChunk | undefined + for (let index = first + 1; index < last; index++) { + const line = lines[index] ?? '' + const head = line.trimEnd() + const header = /^\*\*\* (Add|Delete|Update) File:(.*)$/.exec(head) + if (header) { + checkNotEmpty(current, currentLine) + const [, kind, rest] = header + const path = pathAt(index, rest) + if (kind === 'Add') current = { type: 'add', path, content: '' } + else if (kind === 'Delete') current = { type: 'delete', path } + else current = { type: 'update', path, chunks: [] } + operations.push(current) + currentLine = index + chunk = undefined + continue + } + const move = /^\*\*\* Move to:(.*)$/.exec(head) + if (move) { + if ( + current?.type !== 'update' || + current.chunks.length > 0 || + current.movePath !== undefined + ) { + throw lineError( + index, + '"*** Move to:" must come right after "*** Update File:".', + ) + } + current.movePath = pathAt(index, move[1]) + continue + } + if (head === '*** End of File') { + if (!chunk) { + throw lineError(index, '"*** End of File" must come after a chunk.') + } + chunk.endOfFile = true + chunk = undefined + continue + } + switch (current?.type) { + case undefined: + case 'delete': + if (head !== '') { + throw lineError( + index, + 'expected "*** Add File:", "*** Delete File:", or "*** Update File:".', + ) + } + break + case 'add': + if (!line.startsWith('+')) { + throw lineError( + index, + 'each line of an added file must start with "+".', + ) + } + current.content += `${line.slice(1)}\n` + break + case 'update': { + if (line.startsWith('@@')) { + chunk = newChunk(line.slice(2).trim()) + current.chunks.push(chunk) + break + } + // An empty line is an empty context line. Editors often drop the + // space at the end of a line. + const sign = line.charAt(0) || ' ' + if (sign !== ' ' && sign !== '-' && sign !== '+') { + throw lineError( + index, + 'each chunk line must start with " ", "-", or "+".', + ) + } + // The first chunk can start without `@@`. + if (!chunk) { + chunk = newChunk('') + current.chunks.push(chunk) + } + const body = line.slice(1) + if (sign !== '+') chunk.oldLines.push(body) + if (sign !== '-') chunk.newLines.push(body) + break + } + } + } + checkNotEmpty(current, currentLine) + if (operations.length === 0) { + throw lineError(first, 'the patch has no file changes.') + } + return operations +} + +/** + * Every path that a patch changes, also the targets of `*** Move to:`, as + * the patch gives them. The permission rules check them. Throws when the + * patch does not parse. + * + * @example + * ```ts + * patchPaths(text) // ['src/a.ts', 'src/old.ts', 'src/new.ts'] + * ``` + */ +export function patchPaths(text: string) { + const operations = parsePatch(text) + return operations.flatMap((operation) => + operation.type === 'update' && operation.movePath !== undefined + ? [operation.path, operation.movePath] + : [operation.path], + ) +} + +/** The lines as text, each with a line ending. */ +const asText = (lines: ReadonlyArray) => + lines.map((line) => `${line}\n`).join('') + +/** + * `content` with `chunks` applied in order. A chunk is searched after the + * chunk before it, and from the line of its `@@` text. The old lines must be + * whole lines, and must be there once. Keeps a byte order mark, the line + * endings, and a missing line ending at the end. `name` is the file, for the + * errors. + */ +function applyChunks( + content: string, + chunks: ReadonlyArray, + name: string, +) { + const hasBom = content.startsWith(BOM) + const body = hasBom ? content.slice(1) : content + const eol = toEndingsOf('\n', body) + // Work on text that ends with a line ending, so the last line is like the + // others. + const ended = body === '' || body.endsWith('\n') + let text = ended ? body : body + eol + // Where the next chunk starts to search. + let from = 0 + const entries = chunks.entries() + for (const [index, chunk] of entries) { + const label = `Chunk ${index + 1} of ${name}` + const newText = toEndingsOf(asText(chunk.newLines), text) + // With `@@` text, search from its line. Only added lines go below it. + let searchFrom = from + let below = text.length + if (chunk.context !== undefined) { + const [found] = findMatches(text.slice(from), chunk.context) + if (!found) { + throw new Error(`${label}: "@@ ${chunk.context}" is not in the file.`) + } + searchFrom = text.lastIndexOf('\n', from + found.start - 1) + 1 + below = text.indexOf('\n', from + found.end) + 1 + } + if (chunk.oldLines.length === 0) { + const spot = chunk.endOfFile ? text.length : below + text = text.slice(0, spot) + newText + text.slice(spot) + from = spot + newText.length + continue + } + const oldText = toEndingsOf(asText(chunk.oldLines), text) + // ponytail: a match inside a line can hide a whole-line match that + // overlaps it. The chunk then is not found, and the model adds context. + const isWhole = (match: { start: number; end: number }) => { + const startsLine = match.start === 0 || text[match.start - 1] === '\n' + const endsRight = !chunk.endOfFile || match.end === text.length + return startsLine && endsRight + } + const matches = findMatches(text.slice(searchFrom), oldText) + .map((match) => ({ + start: match.start + searchFrom, + end: match.end + searchFrom, + })) + .filter(isWhole) + const [match, ...others] = matches + if (!match) throw new Error(`${label}: the old lines are not in the file.`) + if (others.length > 0) { + throw new Error( + `${label}: the old lines appear ${matches.length} times. Add context lines.`, + ) + } + text = text.slice(0, match.start) + newText + text.slice(match.end) + from = match.start + newText.length + } + const result = ended ? text : text.replace(/\r?\n$/, '') + return hasBom ? BOM + result : result +} + +/** + * One operation with its absolute paths. `to` is the move target, or `full` + * when the file does not move. + */ +interface Step { + operation: PatchOperation + full: string + to: string +} + +/** + * Apply each step. Every file is read and checked first. The files are + * written and removed only when every operation fits. Then the `afterWrite` + * hooks run for the written files. + */ +async function applyAll(env: ToolEnv, steps: ReadonlyArray) { + const { backend } = env + // The new state of each file: its text, or `null` to remove it. A later + // operation on a file sees it. A removed file is always a file that was + // there before the patch. + const planned = new Map() + const changes: Array = [] + for (const { operation, full, to } of steps) { + const shown = env.shown(full) + switch (operation.type) { + case 'add': + planned.set(full, operation.content) + changes.push(`added ${shown}`) + break + case 'delete': { + if (!backend.remove) { + throw new Error( + `Cannot delete ${shown}: the workspace cannot remove files.`, + ) + } + if (planned.has(full)) { + throw new Error( + `Cannot delete ${shown}: the patch changes it before.`, + ) + } + const isFile = (await backend.stat(full))?.type === 'file' + if (!isFile) throw new Error(`Cannot delete ${shown}: no such file.`) + planned.set(full, null) + changes.push(`deleted ${shown}`) + break + } + case 'update': { + const isMove = to !== full + const toShown = env.shown(to) + if (isMove) { + if (!backend.remove) { + throw new Error( + `Cannot move ${shown}: the workspace cannot remove files.`, + ) + } + if (planned.has(full)) { + throw new Error( + `Cannot move ${shown}: the patch changes it before.`, + ) + } + // The target must be free: no file, or a file that the patch + // removes before. + const isFree = + planned.get(to) === null || + (!planned.has(to) && (await backend.stat(to)) === undefined) + if (!isFree) { + throw new Error( + `Cannot move ${shown} to ${toShown}: ${toShown} already exists.`, + ) + } + } + const known = planned.get(full) + if (known === null) { + throw new Error( + `Cannot update ${shown}: the patch removes it before.`, + ) + } + const before = known ?? (await readText(backend, full)) + planned.set(to, applyChunks(before, operation.chunks, shown)) + if (isMove) planned.set(full, null) + changes.push( + isMove ? `moved ${shown} to ${toShown}` : `updated ${shown}`, + ) + break + } + } + } + const files = [...planned] + const written = files.flatMap(([full, text]) => + text === null ? [] : [{ full, text }], + ) + const removed = files.flatMap(([full, text]) => (text === null ? [full] : [])) + for (const { full, text } of written) await backend.writeFile(full, text) + // The checks above refuse a removal when the backend has no `remove`. + for (const full of removed) await backend.remove?.(full) + const notes: Array = [] + for (const { full } of written) notes.push(...(await afterWrite(env, full))) + return ['Applied the patch:', ...changes, ...notes].join('\n') +} + +/** + * The `patch` tool: it changes files with a patch in the `*** Begin Patch` + * format (see {@link parsePatch}). The old lines of a chunk are found with + * {@link findMatches}. All files are checked before the first write, so a + * patch that does not fit changes nothing. The `afterWrite` hooks run for + * each written file. The result lists the changed files. + * + * - `*** Add File:` replaces a file that is there, like `write_file`. + * - `*** Delete File:` needs a file there. + * - `*** Move to:` writes the new path, then removes the old one. It refuses + * a target that is already there, unless the patch removes it first. + * - A delete or a move needs `remove` on the backend. Without it, the tool + * refuses them, before it writes anything. + * - A file that the patch removed or moved away cannot change again in the + * same patch. A file cannot be removed or moved after the patch changed it. + */ +export function patchTools(env: ToolEnv) { + return [ + toolDefinition({ + name: 'patch', + description: DESCRIPTION, + inputSchema: { + type: 'object', + properties: { patch: { type: 'string' } }, + required: ['patch'], + }, + }).server(async (args: unknown) => { + const operations = parsePatch(stringArg(args, 'patch')) + // Resolve each path first. A path outside the workspace can ask the + // user. + const steps: Array = [] + for (const operation of operations) { + const full = await env.reach(operation.path, 'patch', 'file') + const movePath = + operation.type === 'update' ? operation.movePath : undefined + const to = + movePath === undefined + ? full + : await env.reach(movePath, 'patch', 'file') + steps.push({ operation, full, to }) + } + // Hold the lock of every file, taken in one fixed order, so two + // patches cannot wait for each other. + const paths = [ + ...new Set(steps.flatMap((step) => [step.full, step.to])), + ].sort() + let run = () => applyAll(env, steps) + for (const path of paths) { + const inner = run + run = () => env.lock(path, inner) + } + return run() + }), + ] +} diff --git a/packages/ai-harness/src/first-party/coding/read.ts b/packages/ai-harness/src/first-party/coding/read.ts new file mode 100644 index 0000000000..7aaa315901 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/read.ts @@ -0,0 +1,201 @@ +import { toolDefinition } from '@tanstack/ai' +import { toBase64 } from '../../media' +import { mimeTypeOf } from '../../media-ref' +import { isRecord } from '../../utils' +import { pathsOf, stringArg } from './backend' +import type { ContentPart } from '@tanstack/ai' +import type { ToolEnv } from './backend' + +// The page limits stay inside the defaults of `boundToolOutput` (2000 lines, +// 50 KiB), so that plugin does not cut a page again. +const MAX_LINES = 2000 +const MAX_PAGE_BYTES = 50 * 1024 +const MAX_LINE_CHARS = 2000 +// A NUL byte in this many first bytes marks a binary file. +const BINARY_CHECK_BYTES = 8 * 1024 +// ponytail: media goes to the model as inline base64, and it stays in the +// transcript. Store it as harness media if big files make turns too large. +const MAX_MEDIA_BYTES = 5 * 1024 * 1024 +const MAX_SUGGESTIONS = 3 + +// The first bytes of the files that go to the model as media parts. The +// bytes decide, not the extension: a Git LFS pointer named `x.png` is text. +const MEDIA = [ + { type: 'image', mimeType: 'image/png', signature: /^\x89PNG/ }, + { type: 'image', mimeType: 'image/jpeg', signature: /^\xff\xd8\xff/ }, + { type: 'image', mimeType: 'image/gif', signature: /^GIF8/ }, + { type: 'image', mimeType: 'image/webp', signature: /^RIFF[\s\S]{4}WEBP/ }, + { type: 'document', mimeType: 'application/pdf', signature: /^%PDF-/ }, +] as const + +// The default decoder drops a byte order mark, so the model does not see +// it. `edit_file` keeps the mark in the file. +const decoder = new TextDecoder() +const encoder = new TextEncoder() + +/** The edit distance of `a` and `b`: the fewest one-character changes. */ +function distance(a: string, b: string) { + // `row[j]` is the distance of the first i characters of `a` and the first + // j characters of `b`. Each pass over `a` makes the next row. + let row = Array.from({ length: b.length + 1 }, (_, j) => j) + for (let i = 0; i < a.length; i++) { + let diagonal = i + let left = i + 1 + row = row.map((above, j) => { + if (j === 0) return left + const cost = a[i] === b[j - 1] ? 0 : 1 + left = Math.min(above + 1, left + 1, diagonal + cost) + diagonal = above + return left + }) + } + return row.at(-1) ?? 0 +} + +/** + * The error for a file that does not exist. It names up to 3 files in the + * same folder with a similar name: the same name in other letter case, the + * same name with another extension, or a name 2 or fewer changes away. + */ +async function notFound(env: ToolEnv, full: string) { + const { dirname, join, parse } = pathsOf(env.backend) + const folder = dirname(full) + const wanted = parse(full).base.toLowerCase() + const wantedStem = parse(wanted).name + // A missing folder has no names to suggest. + const entries = await env.backend.readdir(folder).catch(() => []) + const similar = entries + .filter((entry) => entry.type === 'file') + .map((entry) => { + const name = entry.name.toLowerCase() + const isSameStem = parse(name).name === wantedStem + return { name: entry.name, isSameStem, distance: distance(name, wanted) } + }) + .filter((item) => item.isSameStem || item.distance <= 2) + .sort((a, b) => a.distance - b.distance) + .slice(0, MAX_SUGGESTIONS) + .map((item) => env.shown(join(folder, item.name))) + const message = `File not found: ${env.shown(full)}.` + return new Error( + similar.length > 0 + ? `${message} Did you mean ${similar.join(', ')}?` + : message, + ) +} + +/** `line` with its number. A line over 2000 characters is cut. */ +function numbered(lineNumber: number, line: string) { + const isLong = line.length > MAX_LINE_CHARS + const text = isLong + ? `${line.slice(0, MAX_LINE_CHARS)}... [line cut: ${line.length - MAX_LINE_CHARS} more characters]` + : line + return `${lineNumber}\t${text}` +} + +/** + * The numbered lines from line `offset`, at most `limit` lines (2000 when + * `limit` is not set) and 50 KiB. When the page shows fewer lines than were + * asked for and the file goes on, a note says how to read the next page. + */ +function page( + text: string, + options: { path: string; offset: number; limit: number | undefined }, +) { + const { path, offset, limit } = options + const lines = text.split('\n') + // A final newline ends the last line. It does not start a new one. + if (text.endsWith('\n')) lines.pop() + const window = lines.slice( + offset - 1, + offset - 1 + Math.min(limit ?? MAX_LINES, MAX_LINES), + ) + const shown: Array = [] + let bytes = 0 + for (const [index, line] of window.entries()) { + const row = numbered(offset + index, line) + // One line is at most about 8 KB, so the first line always fits. + bytes += encoder.encode(row).length + 1 + if (bytes > MAX_PAGE_BYTES) break + shown.push(row) + } + if (shown.length === 0) { + return `[No lines from line ${offset}. ${path} has ${lines.length} lines.]` + } + const last = offset + shown.length - 1 + const isCut = shown.length < (limit ?? Infinity) && last < lines.length + if (!isCut) return shown.join('\n') + return `${shown.join('\n')}\n\n[Showing lines ${offset}-${last} of ${lines.length}. Read more with offset ${last + 1}.]` +} + +/** + * `read_file`: numbered lines of a text file, a page at a time. An image + * (PNG, JPEG, GIF, WebP) or a PDF comes back as a content part. A binary + * file gets a short note, not its bytes. The text that `afterRead` hooks + * return is added after the lines of a text file. + */ +export function readTools(env: ToolEnv) { + return [ + toolDefinition({ + name: 'read_file', + description: + 'Read a file in the workspace. Text comes back as lines numbered from 1, at most 2000 lines or 50 KiB a page. Use offset and limit to read other lines. Images (PNG, JPEG, GIF, WebP) and PDFs up to 5 MB come back as media.', + inputSchema: { + type: 'object', + properties: { + path: { type: 'string' }, + offset: { type: 'number', description: 'First line, from 1' }, + limit: { type: 'number', description: 'Most lines to read' }, + }, + required: ['path'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const full = await env.reach(stringArg(args, 'path'), 'read_file', 'file') + const offset = + isRecord(args) && typeof args.offset === 'number' + ? Math.max(1, args.offset) + : 1 + const limit = + isRecord(args) && typeof args.limit === 'number' + ? args.limit + : undefined + return env.lock(full, async () => { + const path = env.shown(full) + if (!(await env.backend.stat(full))) throw await notFound(env, full) + const bytes = await env.backend.readFile(full) + const head = String.fromCharCode(...bytes.subarray(0, 12)) + const media = MEDIA.find((item) => item.signature.test(head)) + if (media) { + if (bytes.length > MAX_MEDIA_BYTES) { + const size = (bytes.length / 1024 / 1024).toFixed(1) + throw new Error( + `${path} is ${size} MB. read_file sends images and PDFs up to 5 MB.`, + ) + } + const part: ContentPart = { + type: media.type, + source: { + type: 'data', + value: toBase64(bytes), + mimeType: media.mimeType, + }, + } + return [part] + } + if (bytes.subarray(0, BINARY_CHECK_BYTES).includes(0)) { + const type = mimeTypeOf(full) + const typeNote = type === undefined ? '' : `, ${type}` + return `${path} is a binary file (${bytes.length} bytes${typeNote}). read_file shows text, images, and PDFs only.` + } + const text = decoder.decode(bytes) + const parts = [page(text, { path, offset, limit })] + const hooks = env.hooks() + for (const hook of hooks) { + const extra = await hook.afterRead?.(full) + if (extra) parts.push(extra) + } + return parts.join('\n\n') + }) + }), + ] +} diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts new file mode 100644 index 0000000000..4a79103781 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -0,0 +1,428 @@ +import { toolDefinition } from '@tanstack/ai' +import { globToRegExp } from '../glob' +import { + clip, + hostBackend, + optionalString, + pathsOf, + readText, + stringArg, +} from './backend' +import type { ToolEnv, WorkspaceBackend } from './backend' + +/** Folders that no search looks in. */ +const SKIP = new Set(['.git', 'node_modules']) +/** The time limit of each search command. */ +const TIMEOUT_MS = 30_000 +const MAX_FILES = 1000 +const MAX_MATCHES = 200 +/** A matching line is cut after this many characters. */ +const MAX_LINE = 2000 +/** grep skips bigger files, the same as `rg --max-filesize 1M`. */ +const MAX_SIZE = 1024 * 1024 +// ponytail: the JS walk stops after 5000 files, and the tools say so. Raise +// it if folders without rg and without git get that big. +const MAX_WALK = 5000 + +const FILES_NOTE = + '[Not all files are shown. Use a pattern or a folder to see fewer.]' +const MATCHES_NOTE = + '[Not all matches are shown. Use a narrower pattern, glob, or folder.]' + +/** `git` applies the `.gitignore` files. `-z` ends each path with NUL. */ +const GIT_FILES = 'git ls-files -z --cached --others --exclude-standard' +/** rg flags for grep. One match more than shown tells that there are more. */ +const RG_GREP = `--line-number --no-heading --color never --max-columns ${MAX_LINE} --max-columns-preview --max-filesize 1M --max-count ${MAX_MATCHES + 1}` + +/** Characters that `cmd.exe` reads as special. */ +const CMD_SPECIAL = /[()\][%!^"`<>&|;, *?]/g + +/** + * `value` as one argument of a shell command. The shell gives the value to + * the program as it is. `platform` is the platform of the shell. + * + * - POSIX `sh`: the value goes in single quotes. No character is special + * in them. A single quote in the value becomes `'\''`. + * - Windows `cmd.exe`: first the value is quoted for the program: a `"` + * gets a backslash, and the backslashes before a `"` are doubled. Then a + * `^` goes before each character that `cmd.exe` reads as special, the + * quotes too. So `cmd.exe` expands no `%VAR%` and runs no `&` or `|`. + * This needs delayed expansion off, which is the default. A line break + * cannot be passed safely, so it throws. + * + * @example + * ```ts + * quoteArg(`it's`, 'linux') // 'it'\''s' + * quoteArg('a & b', 'win32') // ^"a^ ^&^ b^" + * ``` + */ +export function quoteArg(value: string, platform: NodeJS.Platform) { + if (platform !== 'win32') return `'${value.replaceAll("'", `'\\''`)}'` + if (/[\r\n]/.test(value)) { + throw new Error('cmd.exe cannot take an argument with a line break.') + } + const quoted = `"${value.replace(/(\\*)"/g, '$1$1\\"').replace(/(\\*)$/, '$1$1')}"` + return quoted.replace(CMD_SPECIAL, '^$&') +} + +/** + * The platform of the shell that `backend.exec` uses, for {@link quoteArg}. + * A backend without `shell` has a POSIX `sh`, like a Linux sandbox. + */ +export function shellPlatform(backend: WorkspaceBackend) { + return backend.shell === 'cmd' ? 'win32' : 'linux' +} + +/** The rg binary of `@vscode/ripgrep`, when that optional peer is there. */ +async function bundledRg() { + try { + // Lazy and optional. The import throws when the binary package is missing. + const { rgPath } = await import('@vscode/ripgrep') + const info = await hostBackend.stat(rgPath) + return info?.type === 'file' ? rgPath : undefined + } catch { + return undefined + } +} + +/** The command that starts rg on `backend`, or `undefined` without rg. */ +async function findRg(backend: WorkspaceBackend) { + // The bundled binary is a host file. A sandbox cannot run it. + if (backend === hostBackend) { + const bundled = await bundledRg() + if (bundled !== undefined) return quoteArg(bundled, process.platform) + } + const probe = await backend.exec('rg --version', { timeoutMs: TIMEOUT_MS }) + return probe.exitCode === 0 ? 'rg' : undefined +} + +/** Run an rg command in `cwd`, and give its output. */ +async function runRg(backend: WorkspaceBackend, command: string, cwd: string) { + const { exitCode, stdout, stderr } = await backend.exec(command, { + cwd, + timeoutMs: TIMEOUT_MS, + }) + // 1 means that rg found nothing. 2 with output means that rg could not + // read some files, but searched the others. + const isDone = + exitCode === 0 || exitCode === 1 || (exitCode === 2 && stdout !== '') + if (isDone) return stdout + const reason = + exitCode === 124 + ? `it took longer than ${TIMEOUT_MS / 1000} seconds, or gave too much output` + : stderr.trim() || `exit code ${exitCode}` + throw new Error(`The search failed: ${reason}`) +} + +/** A line that grep found. `path` is from the searched folder, with `/`. */ +interface Hit { + path: string + line: number + text: string +} + +/** `path` without the `./` that rg puts first. */ +const fromDot = (path: string) => path.replace(/^\.\//, '') + +/** Is `path` in a folder that no search looks in? */ +const isSkipped = (path: string) => + path.split('/').some((part) => SKIP.has(part)) + +/** The paths without empty ones, skipped ones, or repeats, sorted. */ +function tidy(paths: Array) { + const kept = paths + .map(fromDot) + .filter((path) => path !== '' && !isSkipped(path)) + return [...new Set(kept)].sort() +} + +/** Sort by path only, so the lines of one file stay in order. */ +const byPath = (a: Hit, b: Hit) => + Number(a.path > b.path) - Number(a.path < b.path) + +/** A test for the paths that match `glob`. Without a glob, all match. */ +function globMatcher(glob: string | undefined) { + const regex = glob ? globToRegExp(glob) : undefined + return (path: string) => !regex || regex.test(path) +} + +/** The hits in rg output. Each line is the path, NUL, `line:text`. */ +function rgHits(out: string) { + const rows = out.split('\n') + const hits: Array = [] + for (const row of rows) { + const nul = row.indexOf('\0') + const colon = row.indexOf(':', nul) + if (nul === -1 || colon === -1) continue + hits.push({ + path: fromDot(row.slice(0, nul)), + line: Number(row.slice(nul + 1, colon)), + text: row.slice(colon + 1), + }) + } + return hits +} + +/** + * Add the files under `dir` to `out`, as paths that start with `prefix`. + * Stops after MAX_WALK files. It skips links, so it never leaves `dir`. + */ +async function walk( + backend: WorkspaceBackend, + dir: string, + prefix: string, + out: Array, +) { + const entries = await backend.readdir(dir) + for (const entry of entries) { + if (out.length > MAX_WALK) return + if (SKIP.has(entry.name) || entry.type === 'link') continue + const path = prefix + entry.name + if (entry.type === 'dir') { + const next = pathsOf(backend).join(dir, entry.name) + await walk(backend, next, `${path}/`, out) + } else out.push(path) + } +} + +/** + * `files`, paths from `base` with `/`, without the links and the files in a + * linked folder. git lists links, and git for Windows lists the files in a + * junction. Each folder is read once. A name is compared without letter + * case, so a case-insensitive file system cannot hide a link. + */ +async function withoutLinks( + backend: WorkspaceBackend, + base: string, + files: Array, +) { + const linksIn = new Map | undefined>>() + /** The links in `folder`, from `base`. `undefined` when it cannot be read. */ + const linksOf = (folder: string) => { + let links = linksIn.get(folder) + if (links === undefined) { + links = backend.readdir(pathsOf(backend).resolve(base, folder)).then( + (entries) => + new Set( + entries + .filter((entry) => entry.type === 'link') + .map((entry) => entry.name.toLowerCase()), + ), + () => undefined, + ) + linksIn.set(folder, links) + } + return links + } + const kept: Array = [] + for (const file of files) { + const parts = file.split('/') + let folder = '.' + let isLinked = false + for (const part of parts) { + const links = await linksOf(folder) + // A folder that cannot be read fails closed: its files are skipped. + if (links === undefined || links.has(part.toLowerCase())) { + isLinked = true + break + } + folder = `${folder}/${part}` + } + if (!isLinked) kept.push(file) + } + return kept +} + +/** grep in JS: read each file through the backend and test each line. */ +async function grepFiles( + backend: WorkspaceBackend, + base: string, + files: Array, + regex: RegExp, +) { + const hits: Array = [] + for (const path of files) { + const full = pathsOf(backend).resolve(base, path) + const info = await backend.stat(full) + if (info?.type !== 'file' || info.size > MAX_SIZE) continue + const text = await readText(backend, full) + // rg skips binary files too. + if (text.includes('\0')) continue + const lines = text.split('\n') + for (const [index, line] of lines.entries()) { + if (!regex.test(line)) continue + hits.push({ path, line: index + 1, text: line }) + if (hits.length > MAX_MATCHES) return hits + } + } + return hits +} + +/** + * `list_files` and `grep`. Both use the first engine that works: + * + * 1. The rg binary of `@vscode/ripgrep`, only with `hostBackend`. + * 2. `rg` on the PATH of the backend. + * 3. `git ls-files`, when the folder is in a git work tree. + * 4. A walk through `backend.readdir`. + * + * rg and git skip the files that `.gitignore` names. No engine looks in + * `.git` or `node_modules`, or gives links: rg does not follow them, and + * the git list and the walk leave them out. The results have the same + * shape whichever engine ran. + * + * With `permissions()`, grep skips the files that the rules protect (see + * `ToolEnv.protectedIn`), and a note says how many. `list_files` shows + * them: a name is not the contents. + */ +export function searchTools(env: ToolEnv) { + const { backend } = env + const { resolve } = pathsOf(backend) + const platform = shellPlatform(backend) + // Show hidden files, like git. End each path with NUL, and use `/`. + const rgCommon = [ + '--no-config --hidden --null --path-separator /', + ...[...SKIP].map((dir) => `--glob ${quoteArg(`!${dir}`, platform)}`), + ].join(' ') + let rg: Promise | undefined + /** The rg command. The first search finds it, and the tools keep it. */ + const findRgOnce = () => (rg ??= findRg(backend)) + + /** Every file under `base`, sorted, as paths from `base` with `/`. */ + async function filesUnder(base: string) { + const command = await findRgOnce() + if (command !== undefined) { + // The `.` path stops rg from reading stdin. + const out = await runRg(backend, `${command} --files ${rgCommon} .`, base) + return { files: tidy(out.split('\0')), partial: false } + } + const git = await backend.exec(GIT_FILES, { + cwd: base, + timeoutMs: TIMEOUT_MS, + }) + if (git.exitCode === 0) { + const listed = tidy(git.stdout.split('\0')) + const files = await withoutLinks(backend, base, listed) + return { files, partial: false } + } + const walked: Array = [] + await walk(backend, base, '', walked) + return { files: tidy(walked), partial: walked.length > MAX_WALK } + } + + /** + * The hits for `pattern` under `base`, in the files that `isWanted` keeps. + * `listed` is the result of `filesUnder(base)`, when the caller has it. + */ + async function grepUnder( + base: string, + pattern: string, + isWanted: (path: string) => boolean, + listed?: Awaited>, + ) { + // A pattern that is not a valid JS regular expression throws here, + // whichever engine runs. + const regex = new RegExp(pattern) + const command = await findRgOnce() + if (command === undefined) { + const { files, partial } = listed ?? (await filesUnder(base)) + const hits = await grepFiles(backend, base, files.filter(isWanted), regex) + return { hits, partial } + } + const out = await runRg( + backend, + `${command} ${RG_GREP} ${rgCommon} -e ${quoteArg(pattern, platform)} .`, + base, + ) + const hits = rgHits(out) + .filter((hit) => isWanted(hit.path)) + .sort(byPath) + return { hits, partial: false } + } + + return [ + toolDefinition({ + name: 'list_files', + description: + 'List files in a folder (default: the workspace), optionally matching a glob like `src/**/*.ts`. The glob is from that folder. Files that .gitignore names are left out.', + inputSchema: { + type: 'object', + properties: { + pattern: { type: 'string' }, + path: { + type: 'string', + description: 'The folder. Default: the workspace', + }, + }, + }, + replay: 'safe', + }).server(async (args: unknown) => { + const base = await env.reach( + optionalString(args, 'path') ?? '.', + 'list_files', + 'folder', + ) + const isWanted = globMatcher(optionalString(args, 'pattern')) + const { files, partial } = await filesUnder(base) + const listed = files.filter(isWanted) + const shown = listed + .slice(0, MAX_FILES) + .map((file) => env.shown(resolve(base, file))) + const text = shown.join('\n') || 'No files.' + const isCut = partial || listed.length > MAX_FILES + return clip(isCut ? `${text}\n${FILES_NOTE}` : text) + }), + toolDefinition({ + name: 'grep', + description: + 'Search file contents in a folder (default: the workspace) with a regular expression. Returns `file:line: text`. Files that .gitignore names are left out.', + inputSchema: { + type: 'object', + properties: { + pattern: { type: 'string' }, + glob: { type: 'string' }, + path: { + type: 'string', + description: 'The folder. Default: the workspace', + }, + }, + required: ['pattern'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const pattern = stringArg(args, 'pattern') + if (/[\r\n]/.test(pattern)) { + throw new Error('The pattern must be one line. grep tests each line.') + } + const base = await env.reach( + optionalString(args, 'path') ?? '.', + 'grep', + 'folder', + ) + const inGlob = globMatcher(optionalString(args, 'glob')) + // grep must not show the lines of a file that read_file could not + // read without a question. + const isProtected = await env.protectedIn?.(base) + // The note counts the protected files in the list, not in the hits, + // so it does not tell whether a protected file has the pattern. + const listed = isProtected ? await filesUnder(base) : undefined + const skipped = + listed?.files.filter((path) => inGlob(path) && isProtected?.(path)) + .length ?? 0 + const isWanted = (path: string) => + inGlob(path) && !(isProtected?.(path) ?? false) + const { hits, partial } = await grepUnder(base, pattern, isWanted, listed) + const shown = hits.slice(0, MAX_MATCHES).map((hit) => { + const text = hit.text.slice(0, MAX_LINE).trim() + return `${env.shown(resolve(base, hit.path))}:${hit.line}: ${text}` + }) + const lines = [shown.join('\n') || 'No matches.'] + if (partial || hits.length > MAX_MATCHES) lines.push(MATCHES_NOTE) + if (skipped > 0) { + const files = skipped === 1 ? '1 file' : `${skipped} files` + lines.push(`[Skipped ${files} that the permission rules protect.]`) + } + return clip(lines.join('\n')) + }), + ] +} diff --git a/packages/ai-harness/src/first-party/coding/snapshots.ts b/packages/ai-harness/src/first-party/coding/snapshots.ts new file mode 100644 index 0000000000..389b973923 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/snapshots.ts @@ -0,0 +1,398 @@ +import { createHash } from 'node:crypto' +import { LogRecordsCapability, getLogRecords } from '@tanstack/ai' +import { defineCommand } from '../../commands' +import { RevertFiles, RevertStanding, definePlugin } from '../../plugins' +import { hostBackend } from './backend' +import { isRecord } from '../../utils' +import { quoteArg, shellPlatform } from './search' +import type { ModelMessage } from '@tanstack/ai' +import type { WorkspaceBackend } from './backend' + +export interface SnapshotsOptions { + /** The workspace folder, an absolute path. */ + root: string + /** + * A folder for the shadow git repositories, an absolute path. Each + * workspace gets its own repository in it. The project's own `.git` is + * never used. + */ + dataDir: string + /** Where the files are and git runs. Default: {@link hostBackend}. */ + backend?: WorkspaceBackend +} + +/** The files that one model step changed, between two snapshots. */ +export interface SnapshotStep { + /** The tree id at the start of the step. */ + from: string + /** The tree id at the end of the step. */ + to: string + /** The changed paths, relative to the workspace, with `/`. */ + files: Array + /** The tool calls of the step. `session.revert` reads them. */ + toolCallIds: Array +} + +const isStep = (value: unknown): value is SnapshotStep => + isRecord(value) && + typeof value.from === 'string' && + typeof value.to === 'string' && + Array.isArray(value.files) && + Array.isArray(value.toolCallIds) + +/** + * What `unrevert` needs: the tree when the revert started, and the files + * that the revert put back. + */ +const isReverted = ( + value: unknown, +): value is { now: string; files: Array } => + isRecord(value) && + typeof value.now === 'string' && + Array.isArray(value.files) && + value.files.every((file) => typeof file === 'string') + +/** The log record type of a {@link SnapshotStep}. */ +const STEP_RECORD = 'tanstack/snapshots:step' +const TIMEOUT_MS = 60_000 +const BUSY = 'Wait until the turn ends, then try again.' +const REVERTED = 'A revert stands. Run unrevert, or send a prompt, first.' + +/** + * A shadow git repository for `options.root`. Git commands of one + * repository run one at a time, so sessions on one workspace do not clash + * on its index. + */ +function shadowRepo(options: SnapshotsOptions) { + const backend = options.backend ?? hostBackend + const platform = shellPlatform(backend) + const quote = (value: string) => quoteArg(value, platform) + const hash = createHash('sha1').update(options.root).digest('hex') + const gitDir = `${options.dataDir}/${hash}` + // The env wins over a GIT_DIR or GIT_INDEX_FILE from a git hook, so git + // never writes to the project's repository. + const env = { GIT_DIR: gitDir, GIT_INDEX_FILE: `${gitDir}/index` } + // `core.autocrlf=false` keeps the line endings of every file. + const git = `git -c core.autocrlf=false --git-dir ${quote(gitDir)} --work-tree ${quote(options.root)}` + + /** Run a command in the workspace. Throws with git's message when it fails. */ + const exec = async (command: string) => { + const result = await backend.exec(command, { + cwd: options.root, + env, + timeoutMs: TIMEOUT_MS, + }) + if (result.exitCode !== 0) { + throw new Error(`${command} failed: ${result.stderr.trim()}`) + } + return result.stdout + } + const run = (args: string) => exec(`${git} ${args}`) + + /** + * Save the workspace files, and give the id of their tree. One shell runs + * both git commands: each shell start is slow on Windows. + */ + const track = async () => + (await exec(`${git} add -A && ${git} write-tree`)).trim() + const changed = async (from: string, to: string) => { + const names = await run( + `diff --name-only -z --no-renames ${quote(from)} ${quote(to)}`, + ) + return names.split('\0').filter((name) => name !== '') + } + + let queue: Promise = Promise.resolve() + const serial = (fn: () => Promise) => { + const result = queue.then(fn) + queue = result.then( + () => undefined, + () => undefined, + ) + return result + } + + return { + /** Create the repository. `false` when the backend has no git. */ + init: () => + serial(async () => { + const probe = await backend.exec('git --version', { + timeoutMs: TIMEOUT_MS, + }) + if (probe.exitCode !== 0) return false + // A bare repository: it has no work tree of its own. + await exec(`git init --bare -q ${quote(gitDir)}`) + // Keep the bytes of every file: no line ending change, no filter. + await backend.writeFile( + `${gitDir}/info/attributes`, + '* -text -ident -filter\n', + ) + return true + }), + track: () => serial(track), + changed: (from: string, to: string) => serial(() => changed(from, to)), + diff: (from: string, to?: string) => + serial(async () => { + const target = to ?? (await track()) + const patch = await run( + `diff --no-color --no-ext-diff --no-renames ${quote(from)} ${quote(target)}`, + ) + return { files: await changed(from, target), patch } + }), + /** + * Make the files that differ between `from` and `tree` match `tree`, and + * give their count. Of these files, the ones that are not in `tree` are + * removed. With `only`, only the files in it change. All other files stay + * as they are. + */ + restore: (tree: string, from: string, only?: ReadonlySet) => + serial(async () => { + // Pairs of a status letter and a path. `D`: the path is not in `tree`. + const status = ( + await run( + `diff --name-status -z --no-renames ${quote(from)} ${quote(tree)}`, + ) + ).split('\0') + const removed: Array = [] + const restored: Array = [] + for (let i = 1; i < status.length; i += 2) { + const path = status[i] + if (!path || (only && !only.has(path))) continue + if (status[i - 1] === 'D') removed.push(path) + else restored.push(path) + } + // The paths go to git in a file, so no list is too long for the + // command line. `--literal-pathspecs`: `[x].txt` is a name, not a glob. + const paths = async (name: string, list: Array) => { + const file = `${gitDir}/${name}` + await backend.writeFile(file, list.join('\0')) + return `--pathspec-from-file ${quote(file)} --pathspec-file-nul` + } + const commands: Array = [] + if (removed.length > 0) { + // `git rm` removes only the files in the index, so first the index + // gets the files on disk. + commands.push( + `${git} add -A`, + `${git} --literal-pathspecs rm -q -f --ignore-unmatch ${await paths('restore-removed', removed)}`, + ) + } + if (restored.length > 0) { + commands.push( + `${git} --literal-pathspecs checkout ${quote(tree)} ${await paths('restore-restored', restored)}`, + ) + } + // One shell runs all of them: each shell start is slow on Windows. + if (commands.length > 0) await exec(commands.join(' && ')) + return removed.length + restored.length + }), + } +} + +const messageOf = (error: unknown) => + error instanceof Error ? error.message : String(error) + +/** + * Snapshots of the workspace files at the start and the end of each model + * step, in a shadow git repository in `dataDir`. Adds two commands: + * + * - `/undo`: put the files that the last turn changed back as they were + * before it, and remove that turn from the transcript. Other files stay + * as they are. + * - `/redo`: bring back what the last `/undo` removed. A new turn clears it. + * + * `session.revert(messageId)` also puts back the files that the tool calls + * after that message changed, and `session.unrevert()` brings them back. + * + * The changed files of each step go to the session log as a + * `tanstack/snapshots:step` record on a durable host, else to the plugin + * state. `diff(from, to?)` gives the changed files and a unified diff + * between two tree ids. Without `to`, it compares with the files now. + * + * Needs the git CLI on the backend. Without git, the plugin does nothing. + * + * @example + * ```ts + * const root = process.cwd() + * const history = snapshots({ root, dataDir: '/var/lib/agent/snapshots' }) + * defineHarness({ adapter, plugins: () => [workspaceTools({ root }), history] }) + * // `step` is a `tanstack/snapshots:step` record. + * const { files, patch } = await history.diff(step.from, step.to) + * ``` + */ +export function snapshots(options: SnapshotsOptions) { + const repo = shadowRepo(options) + let ready: Promise | undefined + const isReady = () => + (ready ??= repo.init().then( + (hasGit) => { + if (!hasGit) + console.warn('snapshots: git was not found. Snapshots are off.') + return hasGit + }, + (error: unknown) => { + console.warn(`snapshots: ${messageOf(error)}. Snapshots are off.`) + return false + }, + )) + + return definePlugin({ + name: 'tanstack/snapshots', + diff: async (from: string, to?: string) => { + if (!(await isReady())) + throw new Error('snapshots: git is not available.') + return repo.diff(from, to) + }, + provides: [RevertFiles], + setup: async (ctx) => { + if (!(await isReady())) { + // Without git, a revert changes only the transcript. + ctx.provide(RevertFiles, { + revert: async () => undefined, + unrevert: async () => {}, + }) + return + } + const reverted = ctx.getOptional(RevertStanding) + const noSteps: Array = [] + const steps = ctx.state(noSteps) + ctx.provide(RevertFiles, { + revert: async (hidden, records) => { + const calls = new Set( + hidden.flatMap((message) => + (message.toolCalls ?? []).map((call) => call.id), + ), + ) + // A durable host keeps the steps in the log, else in the state. + const all = [...(await steps.get()), ...(await records(STEP_RECORD))] + const later = all + .filter(isStep) + .filter((step) => step.toolCallIds.some((id) => calls.has(id))) + const first = later[0] + if (!first) return undefined + // Only the files that the steps changed go back. The edits of the + // user in other files stay. `now` keeps all files for `unrevert`. + const files = [...new Set(later.flatMap((step) => step.files))] + const now = await repo.track() + await repo.restore(first.from, now, new Set(files)) + return { now, files } + }, + unrevert: async (saved) => { + if (!isReverted(saved)) return + await repo.restore( + saved.now, + await repo.track(), + new Set(saved.files), + ) + }, + }) + /** The tree at the start of the step that runs. */ + let stepStart: string | undefined + /** The tree at the start of the last turn: what `/undo` restores. */ + let turnStart: string | undefined + /** What the last `/undo` removed: the trees of the turn, and its messages. */ + let redo: + | { start: string; end: string; messages: Array } + | undefined + + // A failed snapshot must not fail the turn. It only says so. + const guard = async (fn: () => Promise) => { + try { + return await fn() + } catch (error) { + console.warn(`snapshots: ${messageOf(error)}`) + return undefined + } + } + + return { + middleware: [ + { + name: 'tanstack/snapshots', + optionalRequires: [LogRecordsCapability], + onIteration: async (_ctx, { iteration }) => { + // A later step starts with the files the step before it left, + // so it reuses that snapshot. The model call changes no file. + if (iteration > 0 && stepStart !== undefined) return + stepStart = await guard(repo.track) + if (iteration > 0) return + turnStart = stepStart + redo = undefined + }, + onToolPhaseComplete: async (ctx, { toolCalls }) => { + const from = stepStart + if (from === undefined) return + const step = await guard(async () => { + const to = await repo.track() + const files = to === from ? [] : await repo.changed(from, to) + const toolCallIds = toolCalls.map((call) => call.id) + return { from, to, files, toolCallIds } + }) + // A failed snapshot leaves no start: the next step takes one. + stepStart = step?.to + if (!step || step.files.length === 0) return + const log = getLogRecords(ctx, { optional: true }) + if (log) await log.append([{ type: STEP_RECORD, ...step }]) + else await steps.update((saved) => [...saved, step]) + }, + }, + ], + commands: { + // ponytail: one level. A stack of turn starts gives more undos. + undo: defineCommand({ + description: + 'Undo the last turn: restore the files and remove it from the transcript', + run: async (_input, { session }) => { + if (session.snapshot().status !== 'idle') return BUSY + if (reverted?.()) return REVERTED + const messages = await session.transcript() + // ponytail: the turn starts at the last user message. A steer + // that joined the turn is a user message too, so only the + // part after it goes. + const start = messages.findLastIndex( + (message) => message.role === 'user', + ) + // When the session is idle, the step start is the tree at the + // end of the last turn. + const turnEnd = stepStart + if ( + turnStart === undefined || + turnEnd === undefined || + start === -1 + ) { + return 'Nothing to undo.' + } + // Only the files that the turn changed go back. A later edit of + // the user in one of them is lost. Other files stay as they are. + const count = await repo.restore(turnStart, turnEnd) + redo = { + start: turnStart, + end: turnEnd, + messages: messages.slice(start), + } + turnStart = undefined + await session.replaceTranscript(messages.slice(0, start)) + return `Undid the last turn. Files restored: ${count}.` + }, + }), + redo: defineCommand({ + description: 'Bring back the files and the turn of the last /undo', + run: async (_input, { session }) => { + if (session.snapshot().status !== 'idle') return BUSY + if (reverted?.()) return REVERTED + if (!redo) return 'Nothing to redo.' + const { start, end, messages } = redo + await repo.restore(end, start) + // The same turn can be undone again. `stepStart` is still `end`. + turnStart = start + redo = undefined + const transcript = await session.transcript() + await session.replaceTranscript([...transcript, ...messages]) + return 'Redid the last turn.' + }, + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/coding/web.ts b/packages/ai-harness/src/first-party/coding/web.ts new file mode 100644 index 0000000000..cb45d92cc4 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/web.ts @@ -0,0 +1,472 @@ +import { lookup } from 'node:dns/promises' +import * as http from 'node:http' +import * as https from 'node:https' +import { BlockList, isIP } from 'node:net' +import { toolDefinition } from '@tanstack/ai' +import { isRecord } from '../../utils' +import { optionalString, stringArg } from './backend' +import type { IncomingMessage } from 'node:http' +import type { LookupFunction } from 'node:net' + +const DEFAULT_TIMEOUT_MS = 30_000 +const MAX_TIMEOUT_MS = 120_000 +/** `webfetch` stops reading a body after this many bytes. */ +const MAX_BYTES = 5 * 1024 * 1024 +const MAX_REDIRECTS = 5 +/** A NUL byte in this many first bytes marks a binary body. */ +const BINARY_CHECK_BYTES = 8 * 1024 +const DEFAULT_RESULTS = 8 +const MAX_RESULTS = 20 + +const FORMATS = ['markdown', 'text', 'html'] as const +type Format = (typeof FORMATS)[number] + +const HTML_TYPES = /^(text\/html|application\/xhtml\+xml)$/ +const TEXT_TYPES = /^text\/|[/+](json|xml)$|javascript$/ + +/** The part of `turndown` that `webfetch` uses. It ships no types. */ +type TurndownClass = new (options: { + headingStyle: 'atx' + codeBlockStyle: 'fenced' +}) => { + remove: (tags: Array) => unknown + turndown: (html: string) => string +} + +const ENTITIES: Record = { + amp: '&', + lt: '<', + gt: '>', + quot: '"', + apos: "'", + nbsp: ' ', +} + +/** The addresses that are not on the public internet. */ +const PRIVATE = new BlockList() +// This network. 0.0.0.0 reaches this machine. +PRIVATE.addSubnet('0.0.0.0', 8) +PRIVATE.addSubnet('10.0.0.0', 8) +// Shared address space. Some clouds keep metadata here. +PRIVATE.addSubnet('100.64.0.0', 10) +PRIVATE.addSubnet('127.0.0.0', 8) +// Link-local, with the cloud metadata address 169.254.169.254. +PRIVATE.addSubnet('169.254.0.0', 16) +PRIVATE.addSubnet('172.16.0.0', 12) +PRIVATE.addSubnet('192.168.0.0', 16) +// Multicast and reserved. +PRIVATE.addSubnet('224.0.0.0', 3) +// `::`, `::1`, and IPv4 in the old `::a.b.c.d` form. The IPv4 rules above +// also match the `::ffff:a.b.c.d` form. +PRIVATE.addSubnet('::', 96, 'ipv6') +PRIVATE.addSubnet('fc00::', 7, 'ipv6') +PRIVATE.addSubnet('fe80::', 10, 'ipv6') +PRIVATE.addSubnet('ff00::', 8, 'ipv6') + +/** A search engine for the `websearch` tool. */ +export interface SearchProvider { + /** At most `options.limit` results for `query`. */ + search: ( + query: string, + options: { limit: number; signal?: AbortSignal }, + ) => Promise> +} + +export interface WebToolsOptions { + /** + * The `fetch` that `webfetch` uses. Default: a request with `node:http` + * and `node:https` that checks the address of each connection. A custom + * `fetch` looks up host names itself. Then a DNS server that changes its + * answer after the check (DNS rebinding) can reach a private host. + */ + fetch?: typeof fetch + /** + * Let `webfetch` read localhost and private network hosts. Default: + * `false`. Turn it on only when the model may reach your network. + */ + allowPrivateHosts?: boolean + /** Adds the `websearch` tool. Without it, there is no `websearch` tool. */ + search?: SearchProvider +} + +/** The `turndown` class, or `undefined` when that optional peer is missing. */ +async function loadTurndown() { + try { + // @ts-expect-error -- turndown ships no types. TurndownClass is the part we use. + const turndown: { default: TurndownClass } = await import('turndown') + return turndown.default + } catch { + return undefined + } +} + +/** Decode `&`, `<`, `'`, `A`, and the other common entities. */ +function decodeEntities(text: string) { + return text.replace( + /&(#x[\da-f]+|#\d+|[a-z]+);/gi, + (entity, name: string) => { + if (!name.startsWith('#')) return ENTITIES[name.toLowerCase()] ?? entity + const isHex = name[1] === 'x' || name[1] === 'X' + const code = Number.parseInt(name.slice(isHex ? 2 : 1), isHex ? 16 : 10) + return code <= 0x10ffff ? String.fromCodePoint(code) : entity + }, + ) +} + +/** The text of an HTML page: no tags, scripts, or styles. One block a line. */ +function htmlToText(html: string) { + const text = html + .replace(/<(script|style|noscript|template)\b[\s\S]*?<\/\1\s*>/gi, '') + .replace(//g, '') + .replace(/<\/?(br|p|div|li|tr|h[1-6])\b[^>]*>/gi, '\n') + .replace(/<[^>]*>/g, '') + return decodeEntities(text) + .replace(/[^\S\n]+/g, ' ') + .replace(/ *\n\s*/g, '\n') + .trim() +} + +/** `html` in the format the model asked for. */ +async function convert(html: string, format: Format | undefined) { + const Turndown = await loadTurndown() + const wanted = format ?? (Turndown ? 'markdown' : 'text') + switch (wanted) { + case 'html': + return html + case 'text': + return htmlToText(html) + case 'markdown': { + if (!Turndown) { + throw new Error( + 'Install turndown to get markdown (npm install turndown), or ask for format "text".', + ) + } + const service = new Turndown({ + headingStyle: 'atx', + codeBlockStyle: 'fenced', + }) + service.remove(['script', 'style', 'noscript', 'template']) + return service.turndown(html) + } + } +} + +/** The `format` argument, or `undefined` when the model did not set it. */ +function formatArg(args: unknown) { + const value = optionalString(args, 'format') + if (value === undefined) return undefined + const format = FORMATS.find((item) => item === value) + if (format === undefined) { + throw new Error('Argument "format" must be markdown, text, or html.') + } + return format +} + +/** True when `address`, an IP address, is not on the public internet. */ +function isPrivateAddress(address: string) { + return PRIVATE.check(address, isIP(address) === 6 ? 'ipv6' : 'ipv4') +} + +/** The error for a host that `webfetch` does not read. */ +function refused(host: string) { + return new Error(`webfetch does not read private or local hosts: ${host}`) +} + +/** + * Throws unless `url` is http or https. Unless `allowPrivateHosts` is set, + * it also throws for a private or local host: the host as written, and + * every address that DNS gives for it. The default request checks the + * addresses again when it connects. + */ +async function checkUrl(url: URL, allowPrivateHosts: boolean) { + const isWeb = url.protocol === 'http:' || url.protocol === 'https:' + if (!isWeb) { + throw new Error( + `webfetch reads http and https URLs only, not ${url.protocol}`, + ) + } + if (allowPrivateHosts) return + // The URL keeps IPv6 hosts in brackets. A name can end with a dot. + const host = url.hostname.replace(/^\[|\]$/g, '').replace(/\.$/, '') + const isLocalName = host === 'localhost' || host.endsWith('.localhost') + if (isLocalName) throw refused(url.hostname) + if (isIP(host) !== 0) { + if (isPrivateAddress(host)) throw refused(url.hostname) + return + } + const addresses = await lookup(host, { all: true }).catch(() => { + throw new Error(`Could not resolve ${host}.`) + }) + const isPrivate = addresses.some(({ address }) => isPrivateAddress(address)) + if (isPrivate) throw refused(url.hostname) +} + +/** + * A `lookup` for Node sockets. It refuses `hostname` when any of its + * addresses is private, unless `allowPrivateHosts` is set. The socket + * connects only to the addresses that were checked here, so a DNS answer + * that changes after `checkUrl` cannot reach a private host. + */ +function checkedLookup(allowPrivateHosts: boolean) { + const checked: LookupFunction = (hostname, options, callback) => { + void lookup(hostname, { ...options, all: true }).then( + (addresses) => { + const isPrivate = + !allowPrivateHosts && + addresses.some(({ address }) => isPrivateAddress(address)) + const [first] = addresses + if (isPrivate) callback(refused(hostname), []) + else if (options.all) callback(null, addresses) + else if (first) callback(null, first.address, first.family) + else callback(new Error(`Could not resolve ${hostname}.`), []) + }, + (error) => callback(error, []), + ) + } + return checked +} + +/** + * A Node response as a fetch `Response`. Only the headers that `webfetch` + * reads are copied. The body is read from the socket on demand. + */ +function toResponse(response: IncomingMessage) { + const status = response.statusCode ?? 500 + const headers = new Headers() + const { location, 'content-type': type } = response.headers + if (location !== undefined) headers.set('location', location) + if (type !== undefined) headers.set('content-type', type) + const hasNoBody = status === 204 || status === 205 || status === 304 + if (hasNoBody) { + response.resume() + return new Response(null, { status, headers }) + } + // The chunks are Buffers, so the stream gives Uint8Array chunks. + const chunks = response[Symbol.asyncIterator]() + const body = new ReadableStream({ + pull: async (controller) => { + const next = await chunks.next() + if (next.done) controller.close() + else controller.enqueue(next.value) + }, + cancel: () => { + response.destroy() + }, + }) + return new Response(body, { status, headers }) +} + +/** + * GET `url` with `node:http` or `node:https`. It never follows redirects. + * Each request gets its own connection (`agent: false`), so every + * connection goes through `checkedLookup`. + */ +async function request( + url: URL, + options: { signal: AbortSignal; allowPrivateHosts: boolean }, +) { + const settings = { + agent: false, + signal: options.signal, + lookup: checkedLookup(options.allowPrivateHosts), + } + const response = await new Promise((resolve, reject) => { + const sent = + url.protocol === 'https:' + ? https.get(url, settings, resolve) + : http.get(url, settings, resolve) + sent.on('error', reject) + }) + return toResponse(response) +} + +/** + * Fetch `start` with `send` and follow at most 5 redirects. Each URL is + * checked before it is sent, so a redirect cannot lead to a private host. + */ +async function follow( + start: URL, + options: { + send: (url: URL) => Promise + allowPrivateHosts: boolean + }, +) { + let url = start + for (let hop = 0; hop <= MAX_REDIRECTS; hop++) { + await checkUrl(url, options.allowPrivateHosts) + const response = await options.send(url) + const location = response.headers.get('location') + const isRedirect = + response.status >= 300 && response.status < 400 && location !== null + if (!isRedirect) return { response, url } + await response.body?.cancel() + url = new URL(location, url) + } + throw new Error(`${start.href} redirected more than ${MAX_REDIRECTS} times.`) +} + +/** The body bytes, at most 5 MB. `isCut` is true when there was more. */ +async function readBody(response: Response) { + const reader = response.body?.getReader() + const chunks: Array = [] + let size = 0 + let isCut = false + while (reader) { + const { done, value } = await reader.read() + if (done) break + chunks.push(value) + size += value.length + if (size > MAX_BYTES) { + isCut = true + await reader.cancel() + break + } + } + return { bytes: Buffer.concat(chunks).subarray(0, MAX_BYTES), isCut } +} + +/** The page as the model gets it: HTML converted, text and JSON as they are. */ +async function readPage( + response: Response, + url: URL, + format: Format | undefined, +) { + if (!response.ok) { + await response.body?.cancel() + throw new Error(`${url.href} answered with HTTP ${response.status}.`) + } + const header = response.headers.get('content-type') ?? '' + const type = (header.split(';', 1)[0] ?? '').trim().toLowerCase() + const isHtml = HTML_TYPES.test(type) + const binary = `${url.href} is ${type || 'binary'}. webfetch reads HTML, text, and JSON only.` + // A missing type is checked on the bytes below. + const isKnownBinary = type !== '' && !isHtml && !TEXT_TYPES.test(type) + if (isKnownBinary) { + await response.body?.cancel() + throw new Error(binary) + } + const { bytes, isCut } = await readBody(response) + if (type === '' && bytes.subarray(0, BINARY_CHECK_BYTES).includes(0)) { + throw new Error(binary) + } + // ponytail: reads every page as UTF-8. Decode the charset of the + // content type if other charsets matter. + const text = new TextDecoder().decode(bytes) + const page = isHtml ? await convert(text, format) : text + return isCut ? `${page}\n\n[Cut at 5 MB. The rest was not read.]` : page +} + +/** + * The web tools for a coding agent: + * + * - `webfetch` reads a URL. HTML comes back as Markdown when the optional + * `turndown` package is installed, else as text. Text and JSON come back + * as they are. Other content types are refused. It reads at most 5 MB, + * stops after 30 seconds (the model can ask for up to 120), and follows + * at most 5 redirects. + * - `websearch` searches with `options.search`. Without a provider, there + * is no `websearch` tool. + * + * `webfetch` fetches URLs that the model picks. So it reads only http and + * https, and it refuses localhost, loopback, link-local (with cloud + * metadata), and private network hosts, unless `allowPrivateHosts` is set. + * Without `options.fetch`, it checks the address of each connection, so a + * DNS answer that changes after the check cannot reach a private host. + * + * @example + * ```ts + * const tools = webTools({ + * search: { search: (query, { limit }) => mySearch(query, limit) }, + * }) + * ``` + */ +export function webTools(options: WebToolsOptions = {}) { + const { allowPrivateHosts = false, search } = options + + const webfetch = toolDefinition({ + name: 'webfetch', + description: + 'Read a web page by its http or https URL. HTML comes back as Markdown by default, or set format to text or html. Text and JSON come back as they are. At most 5 MB is read.', + inputSchema: { + type: 'object', + properties: { + url: { type: 'string' }, + format: { type: 'string', enum: ['markdown', 'text', 'html'] }, + timeoutMs: { + type: 'number', + description: 'Stop after this many milliseconds. At most 120000.', + }, + }, + required: ['url'], + }, + replay: 'safe', + }).server(async (args: unknown, context) => { + const url = new URL(stringArg(args, 'url')) + const format = formatArg(args) + const requested = isRecord(args) ? args.timeoutMs : undefined + // A timeout of 0 would stop at once, so it gets the default. + const isTimeout = typeof requested === 'number' && requested > 0 + const timeoutMs = isTimeout + ? Math.min(requested, MAX_TIMEOUT_MS) + : DEFAULT_TIMEOUT_MS + const timeout = AbortSignal.timeout(timeoutMs) + const runSignal = context?.abortSignal + const signal = runSignal ? AbortSignal.any([timeout, runSignal]) : timeout + const customFetch = options.fetch + const send = customFetch + ? (target: URL) => + customFetch(target.href, { redirect: 'manual', signal }) + : (target: URL) => request(target, { signal, allowPrivateHosts }) + try { + const fetched = await follow(url, { send, allowPrivateHosts }) + return await readPage(fetched.response, fetched.url, format) + } catch (error) { + if (timeout.aborted) { + throw new Error(`${url.href} timed out after ${timeoutMs} ms.`) + } + throw error + } + }) + + if (!search) return [webfetch] + + const websearch = toolDefinition({ + name: 'websearch', + description: + 'Search the web. You get a numbered list of results: title, URL, and a short snippet. Read a result with webfetch.', + inputSchema: { + type: 'object', + properties: { + query: { type: 'string' }, + limit: { + type: 'number', + description: `Most results to get. Default ${DEFAULT_RESULTS}, at most ${MAX_RESULTS}.`, + }, + }, + required: ['query'], + }, + replay: 'safe', + }).server(async (args: unknown, context) => { + const query = stringArg(args, 'query') + const requested = isRecord(args) ? args.limit : undefined + const isLimit = typeof requested === 'number' && requested >= 1 + const limit = isLimit + ? Math.min(Math.floor(requested), MAX_RESULTS) + : DEFAULT_RESULTS + const results = await search.search(query, { + limit, + signal: context?.abortSignal, + }) + if (results.length === 0) return `No results for "${query}".` + return results + .slice(0, limit) + .map((result, index) => { + const lines = [`${index + 1}. ${result.title}`, ` ${result.url}`] + if (result.snippet) lines.push(` ${result.snippet}`) + return lines.join('\n') + }) + .join('\n\n') + }) + + return [webfetch, websearch] +} diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts new file mode 100644 index 0000000000..8f2f62d3ed --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -0,0 +1,446 @@ +import { BackgroundJobs, SessionSignal, definePlugin } from '../../plugins' +import { + PERMISSION_MODES, + PermissionDecisionCapability, + PermissionResources, + PermissionRules, +} from '../permissions' +import { WorkspaceHooks } from '../workspace-hooks' +import { hostBackend, optionalString, pathsOf, stringArg } from './backend' +import { bashResources, bashTools } from './bash' +import { editTools } from './edit' +import { patchPaths, patchTools } from './patch' +import { readTools } from './read' +import { searchTools } from './search' +import { webTools } from './web' +import type { AnyTool } from '@tanstack/ai' +import type { PermissionDecision, ToolResources } from '../permissions' +import type { ToolEnv, WorkspaceBackend } from './backend' +import type { WebToolsOptions } from './web' + +/** True for `true`, `y`, or `yes` (any case). */ +function isYes(answer: unknown) { + return answer === true || /^y(es)?$/i.test(String(answer).trim()) +} + +export interface WorkspaceToolsOptions { + root: string + bashTimeoutMs?: number + /** + * A path outside `root`. `'deny'` (default) refuses it. `'ask'` asks the + * user first, and a yes allows that folder, and the folders in it, for the + * rest of the session. The `bypass` mode of `permissions()` allows it + * without a question. + */ + outside?: 'deny' | 'ask' + /** Where the files are and commands run. Default: {@link hostBackend}. */ + backend?: WorkspaceBackend + /** + * The tools that change files. `'auto'` (default): GPT models (an id that + * starts with `gpt-`, or `o` and a digit, like `o3`) get `patch` and no + * `edit_file`. Other models get `edit_file` and `write_file`, and no + * `patch`. `'edit'` or `'patch'` gives that style to every model. + */ + editStyle?: 'auto' | 'edit' | 'patch' + /** + * The web tools. Default: on, with `webfetch`. Pass `search` to add + * `websearch`. `false` adds no web tools. + */ + web?: false | WebToolsOptions + /** + * When `bash` gives the model only the end of the output, it saves the + * full output in this folder. Relative to `root`, or absolute. + */ + spillDir?: string +} + +/** What the tools need from the session for a path outside the workspace. */ +export interface OutsideAccess { + /** Ask the user a question, and resolve with the answer. */ + ask: (message: string) => Promise + /** The permission mode of the session, for example `'bypass'`. */ + mode: () => unknown + /** + * The working folder of the thread, from `root` (the `cwd` setting of + * `session.configure`). Relative paths start there. Read at each call. + */ + cwd?: () => string | undefined + /** + * What `permissions()` decides in the current mode for a call of a tool + * that touches `paths`, or `undefined` without `permissions()`. Read at + * each call. + */ + rules?: () => + | ((tool: string, paths: ReadonlyArray) => PermissionDecision) + | undefined +} + +/** Is `full` the folder `dir`, or a path in it? `sep` is the separator. */ +function within(dir: string, full: string, sep: string) { + return full === dir || full.startsWith(dir.endsWith(sep) ? dir : dir + sep) +} + +/** A lock per path: work on one path runs one at a time, in order. */ +function createLock() { + const queues = new Map>() + return (path: string, fn: () => Promise) => { + const result = (queues.get(path) ?? Promise.resolve()).then(fn) + const settled = result.then( + () => undefined, + () => undefined, + ) + queues.set(path, settled) + void settled.then(() => { + if (queues.get(path) === settled) queues.delete(path) + }) + return result + } +} + +/** + * Does `editStyle: 'auto'` give `model` the `patch` tool? GPT models do: + * `gpt-` ids and the `o` series, like `o3`. + */ +function prefersPatch(model: string) { + // A gateway id can name the provider first, like `openai/gpt-5`. + const id = model.slice(model.lastIndexOf('/') + 1) + return /^(gpt-|o\d)/i.test(id) +} + +/** + * The tools of {@link workspaceTools} for one session, and what each call + * touches for `PermissionResources`. `access` asks the user about a path + * outside the workspace when `options.outside` is `'ask'`. `hooks` are the + * hooks that plugins added. `note` tells the model that a background `bash` + * job ended, and `signal` kills the jobs that still run. + */ +export function createWorkspaceTools( + options: WorkspaceToolsOptions, + session: { + access?: OutsideAccess + hooks?: () => ReadonlyArray + note?: (text: string) => Promise + signal?: AbortSignal + jobs?: { started: (jobId: string) => void; ended: (jobId: string) => void } + } = {}, +) { + const backend = options.backend ?? hostBackend + const paths = pathsOf(backend) + const root = paths.resolve(options.root) + const { access } = session + // The real outside folders the user allowed in this session, and the + // questions that wait for an answer, by folder, so parallel calls ask once. + const allowed: Array = [] + const asking = new Map>() + /** The working folder: `root`, or the thread's `cwd` in it. */ + const here = () => { + const cwd = access?.cwd?.() + return cwd ? paths.resolve(root, cwd) : root + } + + /** Ask once about `folder`. A yes allows it for the session. */ + const allow = (folder: string, message: string) => { + let answer = asking.get(folder) + if (!answer) { + answer = Promise.resolve(access?.ask(message)).then((reply) => { + asking.delete(folder) + if (isYes(reply)) allowed.push(folder) + return isYes(reply) + }) + asking.set(folder, answer) + } + return answer + } + + /** + * The real path of `full`, with every link resolved. For a path that is + * not there yet: the real path of the nearest parent that is there, plus + * the rest. A link to a missing target throws. + */ + const realPath = async (full: string) => { + const rest: Array = [] + let head = full + for (;;) { + const real = backend.realpath ? await backend.realpath(head) : head + if (real !== undefined) return paths.join(real, ...rest) + const parent = paths.dirname(head) + // Not even the top folder is there. + if (parent === head) return full + rest.unshift(paths.basename(head)) + head = parent + } + } + + /** + * The real path of `full`, as the permission rules see it. In the real + * root, it is a path in `root`, so a root that is a link (like `/tmp` on + * macOS) does not move every path outside. + */ + const realForRules = async (full: string) => { + const [real, realRoot] = await Promise.all([realPath(full), realPath(root)]) + return within(realRoot, real, paths.sep) + ? paths.join(root, paths.relative(realRoot, real)) + : real + } + + const shown = (full: string) => + (within(root, full, paths.sep) ? paths.relative(root, full) || '.' : full) + .split(paths.sep) + .join('/') + + // What each call touches, as the tool resolves it, for the permission + // rules with a `resource`. + const file = (input: unknown) => [ + paths.resolve(here(), stringArg(input, 'path')), + ] + const folder = (input: unknown) => [ + paths.resolve(here(), optionalString(input, 'path') ?? '.'), + ] + const resources: Record = { + read_file: { paths: file }, + write_file: { paths: file }, + edit_file: { paths: file }, + list_files: { paths: folder }, + grep: { paths: folder }, + patch: { + paths: (input) => + patchPaths(stringArg(input, 'patch')).map((each) => + paths.resolve(here(), each), + ), + }, + bash: { commands: bashResources }, + } + + /** + * The rules see the path that the call names. A link can lead to a path + * with a stricter rule, like `notes` to `.env`. So the rules check the + * real path too, and the strictest decision wins. + */ + const checkRealPath = async (path: string, tool: string, full: string) => { + const decide = access?.rules?.() + // `bash` touches commands. Its rules do not match paths. + if (!access || !decide || !resources[tool]?.paths) return + const real = await realForRules(full) + if (paths.relative(full, real) === '') return + // `permissions()` checked `full` before the call. Act only when the + // real path makes the decision stricter. + const decision = decide(tool, [full, real]) + if (decision === decide(tool, [full])) return + if (decision === 'deny') { + throw new Error( + `The permission rules do not allow ${tool} on ${shown(real)}, the real path of ${path}.`, + ) + } + const reply = await access.ask( + `${path} is a link to ${shown(real)}. The permission rules ask before ${tool} uses ${shown(real)}. Allow it? (y/n)`, + ) + if (!isYes(reply)) { + throw new Error( + `The user did not allow ${tool} on ${shown(real)}, the real path of ${path}.`, + ) + } + } + + /** Throws when the user or the options do not allow `real`, outside. */ + const allowOutside = async ( + path: string, + tool: string, + kind: 'file' | 'folder', + full: string, + real: string, + ) => { + if (options.outside !== 'ask' || !access) { + throw new Error(`Path "${path}" is outside the workspace.`) + } + if (access.mode() === 'bypass') return + const dir = kind === 'folder' ? real : paths.dirname(real) + const yes = await allow( + dir, + `${tool} wants ${real}, outside the workspace (${root}). Allow ${dir} for this session? (y/n)`, + ) + if (!yes) { + throw new Error( + `The user did not allow ${full}. It is outside the workspace.`, + ) + } + } + + // The boundary check uses real paths, so a link inside the workspace that + // leads out of it counts as outside. + const reach: ToolEnv['reach'] = async (path, tool, kind) => { + const full = paths.resolve(here(), path) + const real = await realPath(full) + const folders = [await realPath(root), ...allowed] + const isAllowed = folders.some((dir) => within(dir, real, paths.sep)) + if (!isAllowed) await allowOutside(path, tool, kind, full, real) + await checkRealPath(path, tool, full) + return full + } + + const env: ToolEnv = { + backend, + root, + hooks: session.hooks ?? (() => []), + // ponytail: one lock per session. Sessions on one folder do not wait + // for each other. + lock: createLock(), + reach, + shown, + protectedIn: async (dir) => { + const decide = access?.rules?.() + if (!decide) return undefined + // rg and the other engines do not follow links in `dir`. So only `dir` + // itself can be a link, and its real path covers the files in it. + const real = await realForRules(dir) + return (path) => + decide('read_file', [ + paths.resolve(dir, path), + paths.resolve(real, path), + ]) !== 'allow' + }, + } + const tools = [ + ...readTools(env), + ...editTools(env), + ...patchTools(env), + ...searchTools(env), + ...bashTools(env, { + timeoutMs: options.bashTimeoutMs, + note: session.note, + spillDir: options.spillDir, + signal: session.signal, + jobs: session.jobs, + }), + ...(options.web === false ? [] : webTools(options.web)), + ] + + const prompt = + options.outside === 'ask' + ? `Your workspace is ${root}. Paths are relative to it. A path outside it asks the user first.` + : `Your workspace is ${root}. Paths are relative to it.` + return { tools, prompt, resources } +} + +/** + * `tools` with the edit tools of `editStyle` for `model`: without + * `edit_file` for the patch style, else without `patch`. + */ +function withEditStyle( + tools: ReadonlyArray, + editStyle: NonNullable, + model: string, +) { + const usesPatch = + editStyle === 'patch' || (editStyle === 'auto' && prefersPatch(model)) + const hidden = usesPatch ? 'edit_file' : 'patch' + return tools.filter((tool) => tool.name !== hidden) +} + +/** + * File, shell, and web tools for a coding agent, in `root`: `read_file`, + * `write_file`, `edit_file` or `patch` (see `editStyle`), `list_files`, + * `grep`, `bash`, `webfetch`, and `websearch` with a search provider. Edits, + * `bash`, and `webfetch` ask for approval through `permissions()`, and each + * call tells the permission rules which paths or commands it touches. A + * path outside `root` is refused, or with `outside: 'ask'` the user is + * asked first. A link that leads out of `root` counts as outside. Other + * plugins add {@link WorkspaceHooks} to run code after a read or a write. + * + * The tools use `backend` for files and commands. The default, + * `hostBackend`, runs on this machine with the host's authority. Use a + * sandbox backend for code you do not trust. + * + * @example + * ```ts + * workspaceTools({ root: process.cwd(), outside: 'ask' }) + * ``` + */ +export function workspaceTools(options: WorkspaceToolsOptions) { + return definePlugin({ + name: 'tanstack/workspace-tools', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => { + const hooks = ctx.collect(WorkspaceHooks) + // The working folder is the thread's `cwd` setting, read at each call. + const cwd = () => ctx.session.settings().cwd + const { tools, prompt, resources } = createWorkspaceTools(options, { + access: { + ask: (message) => ctx.session.ask({ message }), + mode: () => ctx.config.get('mode'), + cwd, + // Read at each call, so permissions() can come after this plugin. + rules: () => { + const decide = ctx.getOptional(PermissionDecisionCapability) + if (!decide) return undefined + const value = ctx.config.get('mode') + const mode = + PERMISSION_MODES.find((each) => each === value) ?? 'default' + return (tool, paths) => decide(tool, mode, { paths }) + }, + }, + hooks: () => hooks, + note: (text) => ctx.session.note(text, { wake: true }), + // Kills the background `bash` jobs when the session closes. A reload + // keeps them. Without a session, the plugin cleanup kills them. + signal: ctx.getOptional(SessionSignal) ?? ctx.resources.signal, + jobs: ctx.getOptional(BackgroundJobs), + }) + const paths = pathsOf(options.backend ?? hostBackend) + const folder = () => { + const dir = cwd() + return dir + ? ` The working folder of this thread is ${paths.resolve(options.root, dir)}. Relative paths start there.` + : '' + } + return { + tools, + prompts: [() => prompt + folder()], + prepareTools: (turn) => + withEditStyle(turn.tools, options.editStyle ?? 'auto', turn.model), + contribute: [ + PermissionRules.item({ + tool: 'read_file', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ + tool: 'list_files', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ + tool: 'grep', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ + tool: 'write_file', + decision: 'ask', + kind: 'edit', + }), + PermissionRules.item({ + tool: 'edit_file', + decision: 'ask', + kind: 'edit', + }), + PermissionRules.item({ + tool: 'patch', + decision: 'ask', + kind: 'edit', + }), + PermissionRules.item({ + tool: 'bash', + decision: 'ask', + kind: 'execute', + }), + // webfetch sends a request to a URL that the model picks, so it + // asks. A user rule can allow it. websearch stays allowed. + ...(options.web === false + ? [] + : [PermissionRules.item({ tool: 'webfetch', decision: 'ask' })]), + PermissionResources.item(resources), + ], + } + }, + }) +} diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts new file mode 100644 index 0000000000..a52e55e95d --- /dev/null +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -0,0 +1,183 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { basename, join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { + WorkspaceHooks, + hostBackend, + workspaceTools, +} from '../src/first-party/coding' +import { commandEnv } from '../src/first-party/coding/backend' +import { mockAdapter, text, toolCall } from './helpers' + +let dir = '' +beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-backend-')) +}) +afterEach(async () => { + await rm(dir, { recursive: true, force: true }) +}) + +const decoder = new TextDecoder() + +describe('hostBackend', () => { + it('writes into new folders, then reads, stats, and lists them', async () => { + const file = join(dir, 'a', 'b.txt') + await hostBackend.writeFile(file, 'hi') + expect(decoder.decode(await hostBackend.readFile(file))).toBe('hi') + expect(await hostBackend.stat(file)).toMatchObject({ + type: 'file', + size: 2, + }) + expect(await hostBackend.stat(join(dir, 'a'))).toMatchObject({ + type: 'dir', + }) + expect(await hostBackend.readdir(dir)).toEqual([{ name: 'a', type: 'dir' }]) + }) + + it('stats a missing path as undefined', async () => { + expect(await hostBackend.stat(join(dir, 'missing'))).toBeUndefined() + }) + + it('removes a file, and refuses a folder', async () => { + const file = join(dir, 'a', 'b.txt') + await hostBackend.writeFile(file, 'hi') + await expect(hostBackend.remove(join(dir, 'a'))).rejects.toThrow() + await hostBackend.remove(file) + expect(await hostBackend.stat(file)).toBeUndefined() + expect(await hostBackend.stat(join(dir, 'a'))).toMatchObject({ + type: 'dir', + }) + }) + + it('runs a command with added env, and resolves with its exit code', async () => { + expect( + await hostBackend.exec( + `node -e "console.log(process.env.HARNESS_VALUE); process.exit(2)"`, + { cwd: dir, env: { HARNESS_VALUE: 'set' } }, + ), + ).toEqual({ exitCode: 2, stdout: 'set\n', stderr: '' }) + }) + + // The command sleeps for one second. On Windows, `exec` waits for it even + // after the timeout kills the shell, so keep it short. It runs outside + // `dir`: on Windows, a running command keeps its folder from removal. + const sleep = `node -e "setTimeout(function () {}, 1000)"` + + it('gives exit code 124 when the timeout kills a command', async () => { + const result = await hostBackend.exec(sleep, { timeoutMs: 100 }) + expect(result.exitCode).toBe(124) + }) + + it('gives exit code 1 when the signal stops a command', async () => { + const controller = new AbortController() + const running = hostBackend.exec(sleep, { signal: controller.signal }) + controller.abort() + expect((await running).exitCode).toBe(1) + }) + + it('starts a background command, then gives its exit code and output', async () => { + // --no-warnings: a runner that sets NO_COLOR and FORCE_COLOR makes node + // print a warning to the output. + const job = hostBackend.spawn( + `node --no-warnings -e "console.log('done')"`, + { + cwd: dir, + }, + ) + expect(await job.wait()).toEqual({ exitCode: 0 }) + expect(job.output()).toBe('done\n') + }) + + it('kills a background command', async () => { + const job = hostBackend.spawn(`node -e "setTimeout(function () {}, 60000)"`) + job.kill() + expect((await job.wait()).exitCode).not.toBe(0) + }) + + // cmd.exe looks in the current folder before the PATH. A cloned repo can + // put an `rg.cmd` there, and grep runs `rg` with no question. + it.runIf(process.platform === 'win32')( + 'does not run a program from the working folder by its name', + async () => { + // A host can set the variable already. Test the host that does not. + const before = process.env.NoDefaultCurrentDirectoryInExePath + delete process.env.NoDefaultCurrentDirectoryInExePath + try { + await writeFile(join(dir, 'rg.cmd'), '@echo ran> "%~dp0ran.txt"\r\n') + await hostBackend.exec('rg --version', { cwd: dir }) + await hostBackend.spawn('rg --version', { cwd: dir }).wait() + expect(await hostBackend.stat(join(dir, 'ran.txt'))).toBeUndefined() + } finally { + if (before !== undefined) { + process.env.NoDefaultCurrentDirectoryInExePath = before + } + } + }, + ) + + it('sets NoDefaultCurrentDirectoryInExePath in the env of commands on Windows only', () => { + expect(commandEnv(undefined, 'win32')).toMatchObject({ + NoDefaultCurrentDirectoryInExePath: '1', + }) + expect( + commandEnv( + { AGENT: '1', NoDefaultCurrentDirectoryInExePath: '' }, + 'win32', + ), + ).toMatchObject({ AGENT: '1', NoDefaultCurrentDirectoryInExePath: '1' }) + // Off Windows, only `env` is added to the environment of this process. + expect(commandEnv(undefined, 'linux')).toBeUndefined() + expect(commandEnv({ AGENT: '1' }, 'linux')).toEqual({ + ...process.env, + AGENT: '1', + }) + }) +}) + +describe('workspace hooks', () => { + it('runs afterWrite after each write, and adds afterRead text to read_file', async () => { + const written: Array = [] + const hooks = definePlugin({ + name: 'test/hooks', + setup: () => ({ + contribute: [ + WorkspaceHooks.item({ + afterWrite: async (path) => { + written.push(decoder.decode(await hostBackend.readFile(path))) + }, + afterRead: async (path) => `read ${basename(path)}`, + }), + ], + }), + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('write_file', { path: 'a.txt', content: 'one' }, 'c1'), + () => + toolCall('edit_file', { path: 'a.txt', old: 'one', new: 'two' }, 'c2'), + () => toolCall('read_file', { path: 'a.txt' }, 'c3'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/hooks', + adapter, + // The hooks plugin comes after the tools: the tools read hooks + // when they run. + plugins: () => [workspaceTools({ root: dir }), hooks], + }), + { threadId: 't' }, + ) + + await session.prompt('work') + + expect(written).toEqual(['one', 'two']) + expect(JSON.stringify(calls[3].messages)).toContain( + '1\\ttwo\\n\\nread a.txt', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/coding-bash.test.ts b/packages/ai-harness/tests/coding-bash.test.ts new file mode 100644 index 0000000000..ed3a310f25 --- /dev/null +++ b/packages/ai-harness/tests/coding-bash.test.ts @@ -0,0 +1,283 @@ +import { tmpdir } from 'node:os' +import { dirname, join, posix, resolve } from 'node:path' +import { describe, expect, it } from 'vitest' +import { + bashResources, + bashTools, + splitCommand, +} from '../src/first-party/coding/bash' +import { decidePermission } from '../src/first-party/permissions' +import type { + ToolEnv, + WorkspaceBackend, +} from '../src/first-party/coding/backend' +import type { PermissionRule } from '../src/first-party/permissions' + +type ExecOptions = Parameters[1] +type BashOptions = Parameters[1] + +const root = resolve(tmpdir(), 'harness-bash') +const OK = { exitCode: 0, stdout: 'ok', stderr: '' } +// 3000 lines: the model gets only the last 1000. +const LONG = { + exitCode: 1, + stdout: 'old\n'.repeat(2000) + 'new\n'.repeat(1000), + stderr: '', +} +const LONG_RESULT = `exit code: 1\n[Output cut. Showing the last 1000 of 3000 lines, 4000 of 12000 bytes.]\n\n${'new\n'.repeat(1000)}` + +/** + * A workspace whose commands all give back `result`. It records each + * command and each written file. `backend` replaces parts of the backend. + */ +function workspace( + result: { exitCode: number; stdout: string; stderr: string }, + backend: Partial = {}, +) { + const execs: Array<{ command: string; options: ExecOptions }> = [] + const files = new Map() + const env: ToolEnv = { + backend: { + readFile: async (path) => { + throw new Error(`No file ${path}`) + }, + writeFile: async (path, data) => { + files.set(path, data) + }, + stat: async () => undefined, + readdir: async () => [], + exec: async (command, options) => { + execs.push({ command, options }) + return result + }, + ...backend, + }, + root, + hooks: () => [], + lock: (_path, fn) => fn(), + // Like the real one: a path in the workspace, from the root. + reach: async (path) => resolve(root, path), + shown: (full) => full, + } + return { env, execs, files } +} + +/** Run the `bash` tool once. */ +async function bash(env: ToolEnv, args: object, options?: BashOptions) { + const [tool] = bashTools(env, options) + if (!tool?.execute) throw new Error('No bash tool') + return tool.execute(args) +} + +/** A background job that ends when the test calls `end`, or when killed. */ +function fakeJob(output: string) { + let end!: (exitCode: number) => void + const exited = new Promise<{ exitCode: number }>((done) => { + end = (exitCode) => done({ exitCode }) + }) + const job = { + killed: false, + wait: () => exited, + kill: () => { + job.killed = true + end(143) + }, + output: () => output, + } + return { job, end } +} + +/** A `note` option that keeps each note. */ +function collectNotes() { + const notes: Array = [] + const note = async (text: string) => { + notes.push(text) + } + return { notes, note } +} + +/** Wait until the promises that are ready now have run. */ +const settle = () => new Promise((done) => setTimeout(done, 0)) + +describe('splitCommand', () => { + it.each([ + { command: 'ls && rm -rf x', parts: ['ls', 'rm -rf x'] }, + { command: 'a | b; c || d', parts: ['a', 'b', 'c', 'd'] }, + { command: 'ls\npwd', parts: ['ls', 'pwd'] }, + { command: 'echo "a && b"', parts: ['echo "a && b"'] }, + { command: "echo 'a; b' && ls", parts: ["echo 'a; b'", 'ls'] }, + { command: 'echo "a \\" b" && ls', parts: ['echo "a \\" b"', 'ls'] }, + { command: 'echo a\\;b', parts: ['echo a\\;b'] }, + { command: "echo $'a\\'b' && ls", parts: ["echo $'a\\'b'", 'ls'] }, + ])('splits `$command`', ({ command, parts }) => { + expect(splitCommand(command)).toEqual(parts) + }) + + it.each([ + { name: 'a command substitution', command: 'echo $(rm x) && ls' }, + { name: 'a backtick', command: 'echo `rm x` && ls' }, + { name: 'an input process substitution', command: 'diff <(ls a) b' }, + { name: 'an output process substitution', command: 'ls | tee >(rm x)' }, + { name: 'a heredoc', command: 'cat < { + expect(splitCommand(command)).toEqual([command]) + }) +}) + +describe('bash permissions', () => { + const rules: Array = [ + { tool: 'bash', decision: 'ask', kind: 'execute' }, + { tool: 'bash', resource: 'ls*', decision: 'allow' }, + ] + + it.each([ + { command: 'ls -la && ls src', expected: 'allow' }, + { command: 'ls && rm -rf x', expected: 'ask' }, + { command: 'ls $(rm -rf x)', expected: 'ask' }, + { command: 'ls() ( rm -rf x ); ls', expected: 'ask' }, + ])('$expected for `$command`, with an allow rule for ls*', (example) => { + const commands = bashResources({ command: example.command }) + expect( + decidePermission(rules, 'bash', 'default', { resources: { commands } }), + ).toBe(example.expected) + }) + + it('refuses a call without a command string', () => { + expect(() => bashResources({ command: 1 })).toThrow( + 'Argument "command" must be a string.', + ) + }) +}) + +describe('bash', () => { + it('runs the command in the workspace, with AGENT=1 and the default timeout', async () => { + const { env, execs } = workspace(OK) + expect(await bash(env, { command: 'make' })).toBe('exit code: 0\nok') + expect(execs).toEqual([ + { + command: 'make', + options: { cwd: root, env: { AGENT: '1' }, timeoutMs: 120_000 }, + }, + ]) + }) + + it.each([ + { name: 'the factory timeout without one', args: {}, expected: 30_000 }, + { + name: 'the timeout of the call', + args: { timeoutMs: 5_000 }, + expected: 5_000, + }, + { + name: 'at most 10 minutes', + args: { timeoutMs: 3_600_000 }, + expected: 600_000, + }, + { + name: 'the factory timeout for 0', + args: { timeoutMs: 0 }, + expected: 30_000, + }, + ])('gives exec $name', async ({ args, expected }) => { + const { env, execs } = workspace(OK) + await bash(env, { command: 'make', ...args }, { timeoutMs: 30_000 }) + expect(execs[0]?.options?.timeoutMs).toBe(expected) + }) + + it('keeps the last lines of long output, with a note, and saves no file', async () => { + const { env, files } = workspace(LONG) + expect(await bash(env, { command: 'test' })).toBe(LONG_RESULT) + expect(files.size).toBe(0) + }) + + it('saves the full output in spillDir, and gives its path', async () => { + const { env, files } = workspace(LONG) + const result = await bash( + env, + { command: 'test' }, + { spillDir: '.agent/bash' }, + ) + const [path = ''] = files.keys() + expect(files.size).toBe(1) + expect(dirname(path)).toBe(join(root, '.agent', 'bash')) + expect(files.get(path)).toBe(LONG.stdout) + expect(result).toBe(`${LONG_RESULT}\n[Full output saved to ${path}.]`) + }) + + it('saves it with a POSIX path in a sandbox, also on a Windows host', async () => { + const { env, files } = workspace(LONG, { shell: 'sh' }) + await bash( + { ...env, root: '/workspace' }, + { command: 'test' }, + { spillDir: '.agent/bash' }, + ) + const [path = ''] = files.keys() + expect(posix.dirname(path)).toBe('/workspace/.agent/bash') + }) + + it('still gives the output when the save fails', async () => { + const { env } = workspace(LONG, { + writeFile: async () => { + throw new Error('disk full') + }, + }) + expect( + await bash(env, { command: 'test' }, { spillDir: '.agent/bash' }), + ).toBe(`${LONG_RESULT}\n[The full output was not saved: Error: disk full]`) + }) +}) + +describe('bash in the background', () => { + it('returns at once, and notes the exit code and output when the job ends', async () => { + const { job, end } = fakeJob('built\n') + const spawned: Array<{ command: string; options: unknown }> = [] + const { env } = workspace(OK, { + spawn: (command, options) => { + spawned.push({ command, options }) + return job + }, + }) + const { notes, note } = collectNotes() + expect( + await bash(env, { command: 'npm run build', background: true }, { note }), + ).toEqual({ jobId: 'bash-1', status: 'started' }) + expect(spawned).toEqual([ + { + command: 'npm run build', + options: { cwd: root, env: { AGENT: '1' } }, + }, + ]) + await settle() + expect(notes).toEqual([]) + end(2) + await settle() + expect(notes).toEqual([ + 'Background job bash-1 ended.\nexit code: 2\nbuilt\n', + ]) + }) + + it('kills running jobs when the signal aborts, and sends no note', async () => { + const { job } = fakeJob('') + const { env } = workspace(OK, { spawn: () => job }) + const { notes, note } = collectNotes() + const controller = new AbortController() + await bash( + env, + { command: 'npm run dev', background: true }, + { note, signal: controller.signal }, + ) + controller.abort() + await settle() + expect(job.killed).toBe(true) + expect(notes).toEqual([]) + }) + + it('gives a tool error when the backend cannot start background commands', async () => { + const { env } = workspace(OK) + await expect( + bash(env, { command: 'npm run dev', background: true }), + ).rejects.toThrow('This workspace cannot run commands in the background.') + }) +}) diff --git a/packages/ai-harness/tests/coding-edit.test.ts b/packages/ai-harness/tests/coding-edit.test.ts new file mode 100644 index 0000000000..10619c5d80 --- /dev/null +++ b/packages/ai-harness/tests/coding-edit.test.ts @@ -0,0 +1,162 @@ +import { tmpdir } from 'node:os' +import { join, resolve } from 'node:path' +import { describe, expect, it } from 'vitest' +import { createWorkspaceTools } from '../src/first-party/coding/workspace' +import type { AnyTool } from '@tanstack/ai' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' + +/** A backend that keeps files in a map, by absolute path. */ +function memoryBackend() { + const files = new Map() + const backend: WorkspaceBackend = { + readFile: async (path) => { + const data = files.get(path) + if (!data) throw new Error(`No file ${path}`) + return data + }, + writeFile: async (path, data) => { + files.set( + path, + typeof data === 'string' ? new TextEncoder().encode(data) : data, + ) + }, + stat: async () => undefined, + readdir: async () => [], + exec: async () => ({ exitCode: 0, stdout: '', stderr: '' }), + } + return backend +} + +const root = resolve(tmpdir(), 'harness-edit') +const backend = memoryBackend() +const tools: Array = createWorkspaceTools({ root, backend }).tools +// Keep a byte order mark in the text, so the tests can see it. +const decoder = new TextDecoder('utf-8', { ignoreBOM: true }) + +async function run(name: string, args: object) { + const tool = tools.find((item) => item.name === name) + if (!tool?.execute) throw new Error(`No tool ${name}`) + const result: string = await tool.execute(args) + return result +} + +async function read(path: string) { + return decoder.decode(await backend.readFile(join(root, path))) +} + +/** Write `text` to `path`, run `edit_file` on it, and read the file back. */ +async function edit( + path: string, + text: string, + change: { old: string; new: string; replaceAll?: boolean }, +) { + await backend.writeFile(join(root, path), text) + const result = await run('edit_file', { path, ...change }) + return { result, text: await read(path) } +} + +describe('edit_file', () => { + it('replaces the exact text', async () => { + expect( + await edit('exact.txt', 'let a = 1\nlet b = 2\n', { + old: 'a = 1', + new: 'a = 3', + }), + ).toEqual({ + result: 'Edited exact.txt (1 change).', + text: 'let a = 3\nlet b = 2\n', + }) + }) + + it('matches text in another Unicode form, and keeps the rest of the file', async () => { + // The file is in NFD form. The old text is in NFC form. + const { text } = await edit('nfd.txt', 'café and naïve\n', { + old: 'café', + new: 'tea', + }) + expect(text).toBe('tea and naïve\n') + }) + + it('ignores spaces and tabs at the end of lines', async () => { + const { text } = await edit('trailing.txt', 'a \nb\t\nc \n', { + old: 'a\nb', + new: 'x', + }) + expect(text).toBe('x\nc \n') + }) + + it('prefers an exact match over a looser one', async () => { + // With trailing spaces ignored, `foo\n` is in the file twice. + const { text } = await edit('prefer.txt', 'foo\nfoo \n', { + old: 'foo\n', + new: 'bar\n', + }) + expect(text).toBe('bar\nfoo \n') + }) + + it('refuses text that is not unique at the level that finds it, unless replaceAll', async () => { + const file = 'foo \nbar\nfoo\t\nbar\n' + await expect( + edit('twice.txt', file, { old: 'foo\nbar', new: 'x' }), + ).rejects.toThrow('The old text appears 2 times.') + expect( + await edit('twice.txt', file, { + old: 'foo\nbar', + new: 'x', + replaceAll: true, + }), + ).toEqual({ result: 'Edited twice.txt (2 changes).', text: 'x\nx\n' }) + }) + + it('refuses text that is not in the file, and empty text', async () => { + await expect( + edit('missing.txt', 'abc\n', { old: 'zzz', new: 'q' }), + ).rejects.toThrow('The old text is not in the file.') + await expect( + edit('missing.txt', 'abc\n', { old: '', new: 'q' }), + ).rejects.toThrow('The old text is not in the file.') + }) + + it.each([ + { + name: 'a CRLF file with LF text', + file: 'one\r\ntwo\r\nthree\r\n', + old: 'one\ntwo', + new: '1\n2', + expected: '1\r\n2\r\nthree\r\n', + }, + { + name: 'an LF file with CRLF text', + file: 'one\ntwo\nthree\n', + old: 'one\r\ntwo', + new: '1\r\n2', + expected: '1\n2\nthree\n', + }, + ])('edits $name, and keeps the line endings of the file', async (example) => { + const { text } = await edit('endings.txt', example.file, { + old: example.old, + new: example.new, + }) + expect(text).toBe(example.expected) + }) + + it.each([ + { name: 'without the mark', old: 'first' }, + { name: 'that starts with the mark', old: 'first' }, + ])('keeps the byte order mark, with old text $name', async ({ old }) => { + const { text } = await edit('bom.txt', 'first\nsecond\n', { + old, + new: '1st', + }) + expect(text).toBe('1st\nsecond\n') + }) +}) + +describe('write_file', () => { + it('writes the content as given', async () => { + expect( + await run('write_file', { path: 'new/w.txt', content: 'hi\r\n' }), + ).toBe('Wrote new/w.txt.') + expect(await read('new/w.txt')).toBe('hi\r\n') + }) +}) diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts new file mode 100644 index 0000000000..01e0760fa2 --- /dev/null +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -0,0 +1,446 @@ +import { tmpdir } from 'node:os' +import * as nodePath from 'node:path' +import { describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { workspaceTools } from '../src/first-party/coding' +import { permissions } from '../src/first-party/permissions' +import { + FormatFailed, + formatOnWrite, + formatter, +} from '../src/first-party/coding/formatter' +import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessSession } from '../src' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' +import type { + FormatFailure, + Formatter, + FormatterOptions, +} from '../src/first-party/coding/formatter' + +const { join, posix, resolve } = nodePath +const root = resolve(tmpdir(), 'harness-format') +const encoder = new TextEncoder() +const decoder = new TextDecoder() + +type Exec = (command: string) => Awaited> + +const ok: Exec = () => ({ exitCode: 0, stdout: '', stderr: '' }) + +/** `path` from the root, as the fake POSIX shell gets it. */ +const quoted = (path: string) => `'${join(root, path)}'` + +/** + * A project in memory that starts as `files`, by path from its root. `exec` + * answers each command. `runs` lists each command with its options. With + * `shell: 'sh'`, the project is at `/workspace` with POSIX paths, like a + * Linux sandbox. Else it is at `root` on this machine. + */ +function project(files: Record, exec: Exec = ok, shell?: 'sh') { + const paths = shell === 'sh' ? posix : nodePath + const base = shell === 'sh' ? '/workspace' : root + const store = new Map() + const put = (path: string, content: string) => + store.set(paths.join(base, path), encoder.encode(content)) + for (const [path, content] of Object.entries(files)) put(path, content) + const runs: Array<{ command: string; cwd?: string; timeoutMs?: number }> = [] + const backend: WorkspaceBackend = { + shell, + readFile: async (path) => { + const data = store.get(path) + if (!data) throw new Error(`No file ${path}`) + return data + }, + writeFile: async (path, data) => { + store.set(path, typeof data === 'string' ? encoder.encode(data) : data) + }, + stat: async (path) => { + const data = store.get(path) + return data && { type: 'file', size: data.length, mtimeMs: 0 } + }, + readdir: async (dir) => + [...store.keys()] + .filter((path) => paths.dirname(path) === dir) + .map((path) => ({ name: paths.basename(path), type: 'file' as const })), + exec: async (command, options) => { + runs.push({ command, cwd: options?.cwd, timeoutMs: options?.timeoutMs }) + return exec(command) + }, + } + const read = (path: string) => + decoder.decode(store.get(paths.join(base, path))) + return { backend, runs, put, read } +} + +/** Run the hook once for `path` in a project of `files`. */ +async function formatIn( + files: Record, + path: string, + options: Omit = {}, + exec: Exec = ok, +) { + const { backend, runs } = project(files, exec) + const failures: Array = [] + const afterWrite = formatOnWrite({ root, backend, ...options }, (failure) => + failures.push(failure), + ) + await afterWrite(join(root, path)) + return { runs, commands: runs.map((run) => run.command), failures } +} + +const prettier = { '.prettierrc': '{}' } +const mine: Formatter = { + name: 'mine', + extensions: ['.ts'], + command: (file) => `mine ${file}`, +} + +describe('formatOnWrite', () => { + it.each([ + ['.prettierrc', prettier, 'a.ts', 'npx --no-install prettier --write'], + [ + 'prettier.config.js', + { 'prettier.config.js': '' }, + 'a.css', + 'npx --no-install prettier --write', + ], + [ + 'a prettier key in package.json', + { 'package.json': '{"prettier": {}}' }, + 'a.md', + 'npx --no-install prettier --write', + ], + [ + 'biome.json', + { 'biome.json': '{}' }, + 'a.tsx', + 'npx --no-install biome format --write', + ], + [ + '.oxfmtrc.json', + { '.oxfmtrc.json': '{}' }, + 'a.ts', + 'npx --no-install oxfmt', + ], + [ + 'oxfmt in devDependencies', + { 'package.json': '{"devDependencies": {"oxfmt": "^0.59.0"}}' }, + 'a.json', + 'npx --no-install oxfmt', + ], + ['ruff.toml', { 'ruff.toml': '' }, 'a.py', 'ruff format'], + [ + '[tool.ruff] in pyproject.toml', + { 'pyproject.toml': '[tool.ruff]\nline-length = 88\n' }, + 'a.py', + 'ruff format', + ], + ['go.mod', { 'go.mod': 'module a\n' }, 'main.go', 'gofmt -w'], + ['Cargo.toml', { 'Cargo.toml': '' }, 'src/main.rs', 'rustfmt'], + ])( + 'detects %s and runs its command in the root', + async (_, files, path, command) => { + const { runs } = await formatIn(files, path) + expect(runs).toEqual([ + { command: `${command} ${quoted(path)}`, cwd: root, timeoutMs: 20_000 }, + ]) + }, + ) + + it.each([ + ['no config file', {}, 'a.ts'], + [ + 'a pyproject.toml without ruff', + { 'pyproject.toml': '[project]\n' }, + 'a.py', + ], + ['a file type no formatter takes', { 'go.mod': 'module a\n' }, 'a.ts'], + ['a package.json that is not JSON', { 'package.json': '{' }, 'a.ts'], + ])('runs nothing for %s', async (_, files, path) => { + expect(await formatIn(files, path)).toEqual({ + runs: [], + commands: [], + failures: [], + }) + }) + + it('quotes the path, so a space and a & stay in one argument', async () => { + const { commands } = await formatIn(prettier, 'src/my file & co.ts') + expect(commands).toEqual([ + `npx --no-install prettier --write '${join(root, 'src', 'my file & co.ts')}'`, + ]) + }) + + it('runs a user formatter before the built-in one', async () => { + const { commands } = await formatIn(prettier, 'a.ts', { + formatters: [mine], + }) + expect(commands).toEqual([`mine ${quoted('a.ts')}`]) + }) + + it('reads the root files for when, and skips a formatter that says no', async () => { + const taplo: Formatter = { + name: 'taplo', + extensions: ['.toml'], + command: (file) => `taplo fmt ${file}`, + when: async (project) => + (await project.read('taplo.toml')).includes('[formatting]'), + } + const used = await formatIn({ 'taplo.toml': '[formatting]\n' }, 'a.toml', { + formatters: [taplo], + }) + const skipped = await formatIn({}, 'a.toml', { formatters: [taplo] }) + expect(used.commands).toEqual([`taplo fmt ${quoted('a.toml')}`]) + expect(skipped.commands).toEqual([]) + }) + + it('keeps POSIX paths in a sandbox, also on a Windows host', async () => { + const { backend, runs } = project(prettier, ok, 'sh') + const afterWrite = formatOnWrite( + { root: '/workspace', backend }, + () => undefined, + ) + await afterWrite('/workspace/src/a.ts') + expect(runs).toEqual([ + { + command: `npx --no-install prettier --write '/workspace/src/a.ts'`, + cwd: '/workspace', + timeoutMs: 20_000, + }, + ]) + }) + + it('runs no built-in formatter with builtins: false', async () => { + const { commands } = await formatIn(prettier, 'a.ts', { builtins: false }) + expect(commands).toEqual([]) + }) + + it('checks the config files once, at the first write', async () => { + const { backend, runs, put } = project({}) + const afterWrite = formatOnWrite({ root, backend }, () => undefined) + await afterWrite(join(root, 'a.ts')) + put('.prettierrc', '{}') + await afterWrite(join(root, 'a.ts')) + expect(runs).toEqual([]) + }) + + it.each([ + [ + 'a non-zero exit, with the first stderr line', + { exitCode: 2, stdout: '', stderr: '[error] a.ts: SyntaxError\nat 1:1' }, + 'prettier failed: [error] a.ts: SyntaxError', + ], + [ + 'a missing formatter, with the exit code', + { exitCode: 127, stdout: '', stderr: '' }, + 'prettier failed: exit code 127', + ], + [ + 'a timeout', + { exitCode: 124, stdout: '', stderr: '' }, + 'prettier failed: it took longer than 5 seconds', + ], + ])('reports %s, and does not throw', async (_, result, message) => { + const { runs, failures } = await formatIn( + prettier, + 'a.ts', + { timeoutMs: 5000 }, + () => result, + ) + expect(runs[0]?.timeoutMs).toBe(5000) + expect(failures).toEqual([{ path: join(root, 'a.ts'), message }]) + }) + + it('reports a backend that throws, and does not throw', async () => { + const { failures } = await formatIn(prettier, 'a.ts', {}, () => { + throw new Error('spawn sh ENOENT') + }) + expect(failures).toEqual([ + { path: join(root, 'a.ts'), message: 'spawn sh ENOENT' }, + ]) + }) +}) + +describe('formatter', () => { + it('formats after write_file and edit_file, and a failure keeps the file and the tool result', async () => { + const { backend, runs, read } = project(prettier, () => ({ + exitCode: 2, + stdout: '', + stderr: 'SyntaxError', + })) + const failures: Array = [] + const listener = definePlugin({ + name: 'test/listener', + setup: (ctx) => { + ctx.on(FormatFailed, (failure) => failures.push(failure)) + }, + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('write_file', { path: 'a.ts', content: 'let a=1' }, 'c1'), + () => + toolCall('edit_file', { path: 'a.ts', old: 'a=1', new: 'a=2' }, 'c2'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/formatter', + adapter, + plugins: () => [ + workspaceTools({ root, backend }), + formatter({ root, backend }), + listener, + ], + }), + { threadId: 't' }, + ) + + await session.prompt('work') + await host.close() + + const command = `npx --no-install prettier --write ${quoted('a.ts')}` + const failure = { + path: join(root, 'a.ts'), + message: 'prettier failed: SyntaxError', + } + expect(runs.map((run) => run.command)).toEqual([command, command]) + expect(failures).toEqual([failure, failure]) + expect(read('a.ts')).toBe('let a=2') + expect(JSON.stringify(calls[2].messages)).toContain( + 'Edited a.ts (1 change).', + ) + }) +}) + +describe('formatter with permissions()', () => { + /** + * A session with the workspace tools, the formatter, and `permissions()` + * in `mode`, in a project of `files`. Each reply writes a file. `runs` + * lists the formatter runs. + */ + async function openWith( + mode: string, + paths: Array, + persistence = memoryPersistence(), + files: Record = {}, + ) { + const { backend, runs } = project(files) + const { adapter, calls } = mockAdapter([ + ...paths.map( + (path, index) => () => + toolCall('write_file', { path, content: '' }, `c${index}`), + ), + () => text('done'), + ]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/formatter-permissions', + adapter, + plugins: () => [ + permissions({ root }), + workspaceTools({ root, backend }), + formatter({ root, backend }), + ], + }), + { threadId: 't' }, + ) + await session.command('mode', mode) + return { host, session, runs, 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 + } + + it('in acceptEdits, asks before prettier runs after an edit of prettier.config.js', async () => { + const { host, session, runs, calls } = await openWith('acceptEdits', [ + 'prettier.config.js', + ]) + + const turn = session.prompt('change the config') + const message = await answer(session, { answer: 'reject' }) + await turn + + expect(message).toContain('formatter:prettier') + expect(runs).toEqual([]) + // The write stays, and the result says the formatter did not run. + expect(JSON.stringify(calls[1].messages)).toContain( + 'Wrote prettier.config.js.\\nThe formatter prettier did not run: it is not allowed.', + ) + await host.close() + }) + + it('stops asking after an always answer, also in a new session', async () => { + const persistence = memoryPersistence() + const first = await openWith( + 'acceptEdits', + ['prettier.config.js', 'a.ts'], + persistence, + ) + const turn = first.session.prompt('write twice') + await answer(first.session, { answer: 'always' }) + await turn + expect(first.runs.map((run) => run.command)).toEqual([ + `npx --no-install prettier --write ${quoted('prettier.config.js')}`, + `npx --no-install prettier --write ${quoted('a.ts')}`, + ]) + expect(first.session.snapshot().pendingQuestions).toHaveLength(0) + await first.host.close() + + const second = await openWith('acceptEdits', ['b.ts'], persistence, { + 'prettier.config.js': '', + }) + await second.session.prompt('write again') + expect(second.runs.map((run) => run.command)).toEqual([ + `npx --no-install prettier --write ${quoted('b.ts')}`, + ]) + await second.host.close() + }) + + it('runs with no question in bypass', async () => { + const { host, session, runs } = await openWith('bypass', [ + 'prettier.config.js', + ]) + await session.prompt('change the config') + expect(runs.map((run) => run.command)).toEqual([ + `npx --no-install prettier --write ${quoted('prettier.config.js')}`, + ]) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) + + it('runs with no question and no note without permissions()', async () => { + const { backend, runs } = project({}) + const { adapter, calls } = mockAdapter([ + () => toolCall('write_file', { path: 'prettier.config.js', content: '' }), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/formatter-no-permissions', + adapter, + plugins: () => [ + workspaceTools({ root, backend }), + formatter({ root, backend }), + ], + }), + { threadId: 't' }, + ) + await session.prompt('change the config') + expect(runs.map((run) => run.command)).toEqual([ + `npx --no-install prettier --write ${quoted('prettier.config.js')}`, + ]) + expect(JSON.stringify(calls[1].messages)).not.toContain('did not run') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/coding-patch.test.ts b/packages/ai-harness/tests/coding-patch.test.ts new file mode 100644 index 0000000000..9ff5a176aa --- /dev/null +++ b/packages/ai-harness/tests/coding-patch.test.ts @@ -0,0 +1,483 @@ +import { tmpdir } from 'node:os' +import { join, relative, resolve, sep } from 'node:path' +import { describe, expect, it } from 'vitest' +import { + parsePatch, + patchPaths, + patchTools, +} from '../src/first-party/coding/patch' +import type { AnyTool } from '@tanstack/ai' +import type { + ToolEnv, + WorkspaceBackend, +} from '../src/first-party/coding/backend' + +const root = resolve(tmpdir(), 'harness-patch') +const encoder = new TextEncoder() +// Keep a byte order mark in the text, so the tests can see it. +const decoder = new TextDecoder('utf-8', { ignoreBOM: true }) + +/** A workspace path, with `/` between the parts. */ +const shown = (full: string) => relative(root, full).split(sep).join('/') + +/** A patch with `lines` between the begin and end lines. */ +function patch(...lines: Array) { + return ['*** Begin Patch', ...lines, '*** End Patch'].join('\n') +} + +/** + * The patch tool on files in memory that start as `files`. `written` lists + * the files that the `afterWrite` hook saw, in order. With `canRemove: + * false`, the backend has no `remove`. + */ +function workspace(files: Record, { canRemove = true } = {}) { + const store = new Map() + for (const [path, text] of Object.entries(files)) { + store.set(join(root, path), encoder.encode(text)) + } + const backend: WorkspaceBackend = { + readFile: async (path) => { + const data = store.get(path) + if (!data) throw new Error(`No file ${path}`) + return data + }, + writeFile: async (path, data) => { + store.set(path, typeof data === 'string' ? encoder.encode(data) : data) + }, + remove: canRemove + ? async (path) => { + if (!store.delete(path)) throw new Error(`No file ${path}`) + } + : undefined, + stat: async (path) => { + const data = store.get(path) + return data && { type: 'file', size: data.length, mtimeMs: 0 } + }, + readdir: async () => [], + exec: async () => ({ exitCode: 0, stdout: '', stderr: '' }), + } + const written: Array = [] + const env: ToolEnv = { + backend, + root, + hooks: () => [ + { + afterWrite: async (path) => { + written.push(shown(path)) + }, + }, + ], + lock: (_path, fn) => fn(), + reach: async (path) => resolve(root, path), + shown, + } + const tools: Array = patchTools(env) + const [tool] = tools + return { + written, + run: async (text: string) => { + if (!tool?.execute) throw new Error('No patch tool') + const result: string = await tool.execute({ patch: text }) + return result + }, + /** Every file, by workspace path. */ + files: () => + Object.fromEntries( + [...store].map( + ([full, data]) => [shown(full), decoder.decode(data)] as const, + ), + ), + } +} + +/** Apply `text` to a workspace with one file, `a.txt`, and read it back. */ +async function patchFile(content: string, text: string) { + const files = workspace({ 'a.txt': content }) + await files.run(text) + return files.files()['a.txt'] +} + +describe('parsePatch', () => { + it('reads added, deleted, and updated files, and a move', () => { + const text = patch( + '*** Add File: src/new.ts', + '+export const a = 1', + '+', + '*** Delete File: old.txt', + '*** Update File: src/app.ts', + '*** Move to: src/main.ts', + '@@ function main() {', + ' const x = 1', + '-const y = 2', + '+const y = 3', + ) + expect(parsePatch(text)).toEqual([ + { type: 'add', path: 'src/new.ts', content: 'export const a = 1\n\n' }, + { type: 'delete', path: 'old.txt' }, + { + type: 'update', + path: 'src/app.ts', + movePath: 'src/main.ts', + chunks: [ + { + context: 'function main() {', + oldLines: ['const x = 1', 'const y = 2'], + newLines: ['const x = 1', 'const y = 3'], + endOfFile: false, + }, + ], + }, + ]) + }) + + it('reads several chunks, a first chunk without @@, End of File, and CRLF text', () => { + const text = patch( + '*** Update File: a.txt', + '-one', + '+1', + '@@', + ' two', + '', + '-three', + '*** End of File', + ).replaceAll('\n', '\r\n') + expect(parsePatch(text)).toEqual([ + { + type: 'update', + path: 'a.txt', + chunks: [ + { + oldLines: ['one'], + newLines: ['1'], + endOfFile: false, + }, + { + oldLines: ['two', '', 'three'], + newLines: ['two', ''], + endOfFile: true, + }, + ], + }, + ]) + }) + + it.each([ + { + name: 'a missing begin line', + text: '\n*** Update File: a.txt\n*** End Patch', + error: 'Patch line 2: the patch must start with "*** Begin Patch".', + }, + { + name: 'a missing end line', + text: '*** Begin Patch\n*** Delete File: a.txt\n', + error: 'Patch line 2: the patch must end with "*** End Patch".', + }, + { + name: 'an unknown header', + text: patch('*** Rename File: a.txt'), + error: + 'Patch line 2: expected "*** Add File:", "*** Delete File:", or "*** Update File:".', + }, + { + name: 'an added line without +', + text: patch('*** Add File: a.txt', '+one', 'two'), + error: 'Patch line 4: each line of an added file must start with "+".', + }, + { + name: 'a chunk line without a prefix', + text: patch('*** Update File: a.txt', '@@', ' one', 'two'), + error: 'Patch line 5: each chunk line must start with " ", "-", or "+".', + }, + { + name: 'an update without changes', + text: patch('*** Update File: a.txt', '@@', '*** Delete File: b.txt'), + error: 'Patch line 2: "*** Update File: a.txt" has no changes.', + }, + { + name: 'a header without a path', + text: patch('*** Add File: '), + error: 'Patch line 2: the path is missing.', + }, + ])('refuses $name', ({ text, error }) => { + expect(() => parsePatch(text)).toThrow(error) + }) +}) + +describe('patchPaths', () => { + it('lists every path, with the move targets', () => { + const text = patch( + '*** Add File: new.txt', + '+x', + '*** Delete File: old.txt', + '*** Update File: src/a.ts', + '*** Move to: src/b.ts', + '-x', + ) + expect(patchPaths(text)).toEqual([ + 'new.txt', + 'old.txt', + 'src/a.ts', + 'src/b.ts', + ]) + }) + + it('throws when the patch does not parse', () => { + expect(() => patchPaths('*** Delete File: a.txt')).toThrow( + 'Patch line 1: the patch must start with "*** Begin Patch".', + ) + }) +}) + +describe('patch tool', () => { + it('adds and updates files, and lists them', async () => { + const files = workspace({ 'src/a.ts': 'const a = 1\nconst b = 2\n' }) + const result = await files.run( + patch( + '*** Add File: src/new.ts', + '+hello', + '*** Update File: src/a.ts', + '@@', + '-const b = 2', + '+const b = 3', + ), + ) + expect(result).toBe( + 'Applied the patch:\nadded src/new.ts\nupdated src/a.ts', + ) + expect(files.files()).toEqual({ + 'src/a.ts': 'const a = 1\nconst b = 3\n', + 'src/new.ts': 'hello\n', + }) + }) + + it('searches from the @@ line', async () => { + const file = + 'function a() {\n return 1\n}\nfunction b() {\n return 1\n}\n' + const text = patch( + '*** Update File: a.txt', + '@@ function b() {', + '- return 1', + '+ return 2', + ) + expect(await patchFile(file, text)).toBe( + 'function a() {\n return 1\n}\nfunction b() {\n return 2\n}\n', + ) + }) + + it('puts only added lines below the @@ line', async () => { + const text = patch('*** Update File: a.txt', '@@ two', '+two and a half') + expect(await patchFile('one\ntwo\nthree\n', text)).toBe( + 'one\ntwo\ntwo and a half\nthree\n', + ) + }) + + it('searches each chunk after the chunk before it', async () => { + // `x` is in the file twice, but once after the first chunk. + const text = patch( + '*** Update File: a.txt', + '@@', + '-a', + '+A', + '@@', + '-x', + '+X', + ) + expect(await patchFile('x\na\nx\n', text)).toBe('x\nA\nX\n') + }) + + it('matches the end of the file after End of File', async () => { + const text = patch( + '*** Update File: a.txt', + '@@', + '-end', + '+END', + '*** End of File', + ) + expect(await patchFile('end\nmiddle\nend\n', text)).toBe( + 'end\nmiddle\nEND\n', + ) + }) + + it('sees the changes of an earlier section on the same file', async () => { + const text = patch( + '*** Update File: a.txt', + '-one', + '+1', + '*** Update File: a.txt', + '-1', + '+uno', + ) + expect(await patchFile('one\n', text)).toBe('uno\n') + }) + + it.each([ + { + name: 'a CRLF file stays CRLF', + file: 'one\r\ntwo\r\nthree\r\n', + expected: 'one\r\n2\r\nthree\r\n', + }, + { + name: 'a byte order mark stays', + file: 'one\ntwo\nthree\n', + expected: 'one\n2\nthree\n', + }, + { + name: 'a file without a last line ending stays without it', + file: 'one\ntwo', + expected: 'one\n2', + }, + ])('$name', async ({ file, expected }) => { + const text = patch('*** Update File: a.txt', '@@', ' one', '-two', '+2') + expect(await patchFile(file, text)).toBe(expected) + }) + + it('matches lines in another Unicode form', async () => { + // The file is in NFD form. The patch is in NFC form. + const file = 'café\nnaïve\n' + const text = patch('*** Update File: a.txt', '@@', '-café', '+tea') + expect(await patchFile(file, text)).toBe('tea\nnaïve\n') + }) + + it.each([ + { + name: 'old lines that are not in the file', + file: 'x\nmid\nx\n', + lines: ['@@', '-zzz', '+q'], + error: 'Chunk 1 of a.txt: the old lines are not in the file.', + }, + { + name: 'old lines that are in the file twice', + file: 'x\nmid\nx\n', + lines: ['@@', '-x', '+y'], + error: + 'Chunk 1 of a.txt: the old lines appear 2 times. Add context lines.', + }, + { + name: 'an @@ line that is not in the file', + file: 'x\nmid\nx\n', + lines: ['@@ function nope() {', '-x', '+y'], + error: 'Chunk 1 of a.txt: "@@ function nope() {" is not in the file.', + }, + { + // The file has `x` only at the end of the line `ax`. + name: 'old lines that start inside a line', + file: 'ax\nb\n', + lines: ['@@', '-x', '+y'], + error: 'Chunk 1 of a.txt: the old lines are not in the file.', + }, + ])('refuses $name, and changes nothing', async ({ file, lines, error }) => { + const files = workspace({ 'a.txt': file }) + await expect( + files.run(patch('*** Update File: a.txt', ...lines)), + ).rejects.toThrow(error) + expect(files.files()).toEqual({ 'a.txt': file }) + }) + + it('changes no file when a later file does not fit', async () => { + const files = workspace({ 'a.txt': 'a\n', 'b.txt': 'b\n' }) + const text = patch( + '*** Update File: a.txt', + '-a', + '+A', + '*** Add File: c.txt', + '+c', + '*** Update File: b.txt', + '-zzz', + '+Z', + ) + await expect(files.run(text)).rejects.toThrow( + 'Chunk 1 of b.txt: the old lines are not in the file.', + ) + expect(files.files()).toEqual({ 'a.txt': 'a\n', 'b.txt': 'b\n' }) + expect(files.written).toEqual([]) + }) + + it('runs the afterWrite hook once for each written file', async () => { + const files = workspace({ 'a.txt': 'a\n', 'b.txt': 'b\n' }) + await files.run( + patch( + '*** Update File: a.txt', + '-a', + '+A', + '*** Update File: b.txt', + '-b', + '+B', + '*** Update File: a.txt', + '-A', + '+AA', + ), + ) + expect(files.written).toEqual(['a.txt', 'b.txt']) + }) + + it('deletes a file, and runs no hook for it', async () => { + const files = workspace({ 'a.txt': 'a\n', 'b.txt': 'b\n' }) + const result = await files.run( + patch('*** Add File: c.txt', '+c', '*** Delete File: a.txt'), + ) + expect(result).toBe('Applied the patch:\nadded c.txt\ndeleted a.txt') + expect(files.files()).toEqual({ 'b.txt': 'b\n', 'c.txt': 'c\n' }) + expect(files.written).toEqual(['c.txt']) + }) + + it('moves a file, and applies its chunks', async () => { + const files = workspace({ 'a.txt': 'one\ntwo\n' }) + const result = await files.run( + patch( + '*** Update File: a.txt', + '*** Move to: src/b.txt', + '@@', + ' one', + '-two', + '+2', + ), + ) + expect(result).toBe('Applied the patch:\nmoved a.txt to src/b.txt') + expect(files.files()).toEqual({ 'src/b.txt': 'one\n2\n' }) + expect(files.written).toEqual(['src/b.txt']) + }) + + it.each([ + { + name: 'a delete of a missing file', + lines: ['*** Delete File: missing.txt'], + error: 'Cannot delete missing.txt: no such file.', + }, + { + name: 'a move onto a file that is there', + lines: ['*** Update File: a.txt', '*** Move to: b.txt'], + error: 'Cannot move a.txt to b.txt: b.txt already exists.', + }, + ])('refuses $name, and changes nothing', async ({ lines, error }) => { + const before = { 'a.txt': 'a\n', 'b.txt': 'b\n', 'd.txt': 'd\n' } + const files = workspace(before) + await expect( + files.run( + patch('*** Add File: c.txt', '+c', '*** Delete File: d.txt', ...lines), + ), + ).rejects.toThrow(error) + expect(files.files()).toEqual(before) + }) + + it.each([ + { + name: 'a delete', + lines: ['*** Delete File: a.txt'], + error: 'Cannot delete a.txt: the workspace cannot remove files.', + }, + { + name: 'a move', + lines: ['*** Update File: a.txt', '*** Move to: b.txt'], + error: 'Cannot move a.txt: the workspace cannot remove files.', + }, + ])( + 'refuses $name when the backend cannot remove files, and changes nothing', + async ({ lines, error }) => { + const files = workspace({ 'a.txt': 'a\n' }, { canRemove: false }) + await expect( + files.run(patch('*** Add File: c.txt', '+c', ...lines)), + ).rejects.toThrow(error) + expect(files.files()).toEqual({ 'a.txt': 'a\n' }) + }, + ) +}) diff --git a/packages/ai-harness/tests/coding-read.test.ts b/packages/ai-harness/tests/coding-read.test.ts new file mode 100644 index 0000000000..bf3d93bacf --- /dev/null +++ b/packages/ai-harness/tests/coding-read.test.ts @@ -0,0 +1,308 @@ +import { tmpdir } from 'node:os' +import { join, posix, resolve, sep } from 'node:path' +import { beforeAll, describe, expect, it } from 'vitest' +import { createWorkspaceTools } from '../src/first-party/coding/workspace' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' +import type { WorkspaceHooks } from '../src/first-party/workspace-hooks' + +/** + * A backend that keeps files in a map, by absolute path. It runs no + * commands. `shell: 'sh'` gives it POSIX paths, like a Linux sandbox. + */ +function memoryBackend(shell?: 'sh') { + const separator = shell === 'sh' ? posix.sep : sep + const files = new Map() + const backend: WorkspaceBackend = { + shell, + readFile: async (path) => { + const data = files.get(path) + if (!data) throw new Error(`No file ${path}`) + return data + }, + writeFile: async (path, data) => { + files.set( + path, + typeof data === 'string' ? new TextEncoder().encode(data) : data, + ) + }, + stat: async (path) => { + const data = files.get(path) + if (!data) return undefined + return { type: 'file', size: data.length, mtimeMs: 0 } + }, + readdir: async (path) => { + const prefix = path + separator + const paths = [...files.keys()].filter((file) => file.startsWith(prefix)) + // Like the host, a folder that is not there throws. + if (paths.length === 0) throw new Error(`No folder ${path}`) + return paths + .map((file) => file.slice(prefix.length)) + .filter((name) => !name.includes(separator)) + .sort() + .map((name) => ({ name, type: 'file' as const })) + }, + exec: async () => ({ exitCode: 0, stdout: '', stderr: '' }), + } + return backend +} + +const root = resolve(tmpdir(), 'harness-read') +const backend = memoryBackend() + +const ascii = (text: string) => new TextEncoder().encode(text) +const put = (path: string, data: Uint8Array | string) => + backend.writeFile(join(root, path), data) + +/** + * `read_file` with `hooks`, in `workspace`. Default: the memory backend at + * `root`. + */ +function reader( + options: { + hooks?: Array + workspace?: { root: string; backend: WorkspaceBackend } + } = {}, +) { + const { hooks = [], workspace = { root, backend } } = options + const tool = createWorkspaceTools(workspace, { + hooks: () => hooks, + }).tools.find((item) => item.name === 'read_file') + return async ( + path: string, + options: { offset?: number; limit?: number } = {}, + ) => { + if (!tool?.execute) throw new Error('No read_file tool') + const result: unknown = await tool.execute({ path, ...options }) + return result + } +} + +const read = reader() + +/** The result of `read_file` for a text file. Throws for any other result. */ +async function readText( + path: string, + options: { offset?: number; limit?: number } = {}, +) { + const result = await read(path, options) + if (typeof result !== 'string') throw new Error('Not a text result') + return result +} + +const PNG = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]) +const pngPart = { + type: 'image', + source: { type: 'data', value: 'iVBORw0KGgo=', mimeType: 'image/png' }, +} +const BINARY = new Uint8Array([0x49, 0x44, 0x00, 0x01]) + +describe('read_file text', () => { + it('reads a file with a byte order mark without the mark', async () => { + await put('bom.txt', new Uint8Array([0xef, 0xbb, 0xbf, ...ascii('a\nb\n')])) + expect(await read('bom.txt')).toBe('1\ta\n2\tb') + }) + + it('cuts a line over 2000 characters, and says how much is cut', async () => { + await put('long.txt', `${'a'.repeat(2500)}\nshort`) + expect(await read('long.txt')).toBe( + `1\t${'a'.repeat(2000)}... [line cut: 500 more characters]\n2\tshort`, + ) + }) + + it('reads a Git LFS pointer named .png as text', async () => { + await put('pointer.png', 'version https://git-lfs.github.com/spec/v1\n') + expect(await read('pointer.png')).toBe( + '1\tversion https://git-lfs.github.com/spec/v1', + ) + }) +}) + +describe('read_file pages', () => { + beforeAll(async () => { + const lines = Array.from({ length: 2500 }, (_, index) => `L${index + 1}`) + await put('many.txt', lines.join('\n')) + }) + + it.each([{}, { limit: 5000 }])( + 'shows at most 2000 lines, then says where the next page starts (%o)', + async (options) => { + const lines = (await readText('many.txt', options)).split('\n') + expect(lines.slice(0, 2)).toEqual(['1\tL1', '2\tL2']) + expect(lines.slice(-3)).toEqual([ + '2000\tL2000', + '', + '[Showing lines 1-2000 of 2500. Read more with offset 2001.]', + ]) + }, + ) + + it('reads the next page from its offset to the end, with no note', async () => { + const lines = (await readText('many.txt', { offset: 2001 })).split('\n') + expect(lines).toHaveLength(500) + expect(lines[0]).toBe('2001\tL2001') + expect(lines.at(-1)).toBe('2500\tL2500') + }) + + it('keeps a page under 50 KiB', async () => { + // Each line is about 1 KB, so 51 lines fit in 50 KiB. + await put( + 'wide.txt', + Array.from({ length: 100 }, () => 'x'.repeat(1000)).join('\n'), + ) + const lines = (await readText('wide.txt')).split('\n') + expect(lines.at(-3)).toBe(`51\t${'x'.repeat(1000)}`) + expect(lines.at(-1)).toBe( + '[Showing lines 1-51 of 100. Read more with offset 52.]', + ) + }) + + it('says how many lines there are when offset is past the end', async () => { + expect(await read('many.txt', { offset: 3000 })).toBe( + '[No lines from line 3000. many.txt has 2500 lines.]', + ) + }) +}) + +describe('read_file media and binary files', () => { + it.each([ + { file: 'a.png', bytes: PNG, part: pngPart }, + { + file: 'a.jpg', + bytes: new Uint8Array([0xff, 0xd8, 0xff, 0xe0]), + part: { + type: 'image', + source: { type: 'data', value: '/9j/4A==', mimeType: 'image/jpeg' }, + }, + }, + { + file: 'a.gif', + bytes: ascii('GIF89a'), + part: { + type: 'image', + source: { type: 'data', value: 'R0lGODlh', mimeType: 'image/gif' }, + }, + }, + { + file: 'a.webp', + bytes: new Uint8Array([ + ...ascii('RIFF'), + 0, + 0, + 0, + 0, + ...ascii('WEBPVP8 '), + ]), + part: { + type: 'image', + source: { + type: 'data', + value: 'UklGRgAAAABXRUJQVlA4IA==', + mimeType: 'image/webp', + }, + }, + }, + { + file: 'a.pdf', + bytes: ascii('%PDF-1.7'), + part: { + type: 'document', + source: { + type: 'data', + value: 'JVBERi0xLjc=', + mimeType: 'application/pdf', + }, + }, + }, + ])('returns $file as a $part.type part', async ({ file, bytes, part }) => { + await put(file, bytes) + expect(await read(file)).toEqual([part]) + }) + + it('refuses an image over 5 MB', async () => { + const big = new Uint8Array(6 * 1024 * 1024) + big.set(PNG) + await put('big.png', big) + await expect(read('big.png')).rejects.toThrow( + 'big.png is 6.0 MB. read_file sends images and PDFs up to 5 MB.', + ) + }) + + it.each([ + { + file: 'data.bin', + expected: + 'data.bin is a binary file (4 bytes). read_file shows text, images, and PDFs only.', + }, + { + file: 'clip.mp3', + expected: + 'clip.mp3 is a binary file (4 bytes, audio/mpeg). read_file shows text, images, and PDFs only.', + }, + ])('says $file is binary, with its size', async ({ file, expected }) => { + await put(file, BINARY) + expect(await read(file)).toBe(expected) + }) +}) + +describe('read_file on a missing file', () => { + beforeAll(async () => { + const names = [ + 'src/Button.tsx', + 'src/button.css', + 'src/helpers.ts', + 'src/index.ts', + 'parts/a1.ts', + 'parts/a2.ts', + 'parts/a3.ts', + 'parts/a4.ts', + ] + await Promise.all(names.map((name) => put(name, 'x'))) + }) + + it.each([ + { + path: 'src/button.ts', + message: + 'File not found: src/button.ts. Did you mean src/Button.tsx, src/button.css?', + }, + { + path: 'src/helpres.ts', + message: 'File not found: src/helpres.ts. Did you mean src/helpers.ts?', + }, + { + path: 'parts/a.ts', + message: + 'File not found: parts/a.ts. Did you mean parts/a1.ts, parts/a2.ts, parts/a3.ts?', + }, + { path: 'nowhere/x.ts', message: 'File not found: nowhere/x.ts.' }, + ])('names similar files for $path', async ({ path, message }) => { + await expect(read(path)).rejects.toThrow(message) + }) + + it('names them with POSIX paths in a sandbox, also on a Windows host', async () => { + const sandbox = memoryBackend('sh') + await sandbox.writeFile('/workspace/src/Button.tsx', 'x') + const readInSandbox = reader({ + workspace: { root: '/workspace', backend: sandbox }, + }) + await expect(readInSandbox('src/button.tsx')).rejects.toThrow( + 'File not found: src/button.tsx. Did you mean src/Button.tsx?', + ) + }) +}) + +describe('read_file hooks', () => { + it('adds afterRead text to a text file, not to media or binary files', async () => { + const readWithHook = reader({ + hooks: [{ afterRead: async () => 'hook note' }], + }) + await put('hooked.txt', 'hello') + await put('hooked.png', PNG) + await put('hooked.bin', BINARY) + expect(await readWithHook('hooked.txt')).toBe('1\thello\n\nhook note') + expect(await readWithHook('hooked.png')).toEqual([pngPart]) + expect(await readWithHook('hooked.bin')).toBe( + 'hooked.bin is a binary file (4 bytes). read_file shows text, images, and PDFs only.', + ) + }) +}) diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts new file mode 100644 index 0000000000..3cb99ab0bd --- /dev/null +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -0,0 +1,502 @@ +import { mkdtemp, rm, symlink } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join, resolve, sep } from 'node:path' +import { + afterAll, + afterEach, + beforeAll, + describe, + expect, + it, + vi, +} from 'vitest' +import { hostBackend } from '../src/first-party/coding/backend' +import { quoteArg } from '../src/first-party/coding/search' +import { createWorkspaceTools } from '../src/first-party/coding/workspace' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' +import type { OutsideAccess } from '../src/first-party/coding/workspace' + +// The bundled rg of `@vscode/ripgrep`. It points at a file that exists, so +// no test needs the real binary. No test runs it. +const ripgrep = vi.hoisted(() => ({ rgPath: process.execPath })) +vi.mock('@vscode/ripgrep', () => ({ + get rgPath() { + return ripgrep.rgPath + }, +})) + +afterEach(() => { + vi.restoreAllMocks() + ripgrep.rgPath = process.execPath +}) + +interface ExecResult { + exitCode: number + stdout: string + stderr: string +} + +const ok = (stdout: string): ExecResult => ({ exitCode: 0, stdout, stderr: '' }) +const NOT_FOUND: ExecResult = { exitCode: 127, stdout: '', stderr: 'not found' } +const GIT = 'git ls-files -z --cached --others --exclude-standard' +const NOT_GIT: ExecResult = { + exitCode: 128, + stdout: '', + stderr: 'fatal: not a git repository', +} +const FILES_NOTE = + '[Not all files are shown. Use a pattern or a folder to see fewer.]' +const MATCHES_NOTE = + '[Not all matches are shown. Use a narrower pattern, glob, or folder.]' + +const root = resolve(tmpdir(), 'coding-search') +const hasGit = (await hostBackend.exec('git --version')).exitCode === 0 +// A real process can take seconds to start on a busy machine. +const SPAWN_TIMEOUT = 60_000 + +/** + * A backend with `files` in memory, by path from `root`. `script` answers + * each command. A command that it does not answer is not found. + */ +function fakeBackend( + files: Record, + script: (command: string) => ExecResult | undefined = () => undefined, +) { + const encoder = new TextEncoder() + const byPath = new Map( + Object.entries(files).map(([path, text]) => [ + join(root, path), + encoder.encode(text), + ]), + ) + const commands: Array = [] + const backend: WorkspaceBackend = { + readFile: async (path) => { + const data = byPath.get(path) + if (!data) throw new Error(`No file ${path}`) + return data + }, + writeFile: async () => { + throw new Error('Read only') + }, + stat: async (path) => { + const data = byPath.get(path) + if (!data) return undefined + return { type: 'file', size: data.length, mtimeMs: 0 } + }, + readdir: async (path) => { + const prefix = path + sep + const entries = new Map() + for (const file of byPath.keys()) { + if (!file.startsWith(prefix)) continue + const [name = '', ...rest] = file.slice(prefix.length).split(sep) + entries.set(name, rest.length > 0 ? 'dir' : 'file') + } + // Sorted by name, like the host. + return [...entries] + .sort(([a], [b]) => (a < b ? -1 : 1)) + .map(([name, type]) => ({ name, type })) + }, + exec: async (command) => { + commands.push(command) + return script(command) ?? NOT_FOUND + }, + } + return { backend, commands } +} + +/** + * `list_files` and `grep` on `backend`, with the workspace at `folder`. + * `access` is what the session gives the tools. + */ +function searchOn( + backend: WorkspaceBackend, + folder = root, + access?: OutsideAccess, +) { + const { tools } = createWorkspaceTools({ root: folder, backend }, { access }) + const run = async (name: string, args: object) => { + const tool = tools.find((each) => each.name === name) + if (!tool?.execute) throw new Error(`No tool ${name}`) + const result: unknown = await tool.execute(args) + if (typeof result !== 'string') throw new Error(`${name} gave no text`) + return result + } + return { + listFiles: (args: object = {}) => run('list_files', args), + grep: (args: object) => run('grep', args), + } +} + +/** Run `script` for the commands of the host backend, not a real shell. */ +function scriptHostExec(script: (command: string) => ExecResult | undefined) { + return vi + .spyOn(hostBackend, 'exec') + .mockImplementation(async (command) => script(command) ?? NOT_FOUND) +} + +describe('list_files and grep engines', () => { + const long = `two${'x'.repeat(3000)}` + const files = { + 'src/b.ts': 'const two = 2\r\n', + 'src/a.ts': 'one\ntwo\n', + 'long.txt': `${long}\n`, + '.github/ci.yml': 'two: yes\n', + 'image.png': 'two\0', + 'node_modules/pkg/i.js': 'two\n', + '.git/HEAD': 'two\n', + } + // What rg prints for `files`: not sorted, with `./`, NUL after each path. + // A long line is cut after 2000 bytes, with a note after it. + const rgFiles = + './src/b.ts\0./src/a.ts\0./long.txt\0./image.png\0./.github/ci.yml\0' + const rgMatches = [ + './src/b.ts\x001:const two = 2\r', + './.github/ci.yml\x001:two: yes', + `./long.txt\x001:${long.slice(0, 2000)} [... omitted end of long line]`, + './src/a.ts\x002:two', + '', + ].join('\n') + const gitFiles = + '.github/ci.yml\0image.png\0long.txt\0node_modules/pkg/i.js\0src/a.ts\0src/b.ts\0' + + const engines = [ + { + name: 'rg on the PATH', + script: (command: string) => { + if (command === 'rg --version') return ok('ripgrep 14.1.0') + if (command.startsWith('rg --files ')) return ok(rgFiles) + const isGrepForTwo = + command.startsWith('rg --line-number ') && + command.endsWith(` -e 'two' .`) + return isGrepForTwo ? ok(rgMatches) : undefined + }, + commands: [ + 'rg --version', + expect.stringMatching(/^rg --files /), + expect.stringMatching(/^rg --line-number /), + ], + }, + { + name: 'git ls-files', + script: (command: string) => (command === GIT ? ok(gitFiles) : undefined), + commands: ['rg --version', GIT, GIT], + }, + { + name: 'the JS walk', + script: (command: string) => (command === GIT ? NOT_GIT : undefined), + commands: ['rg --version', GIT, GIT], + }, + ] + + it.each(engines)( + '$name gives the same files and matches', + async ({ script, commands }) => { + const fake = fakeBackend(files, script) + const { listFiles, grep } = searchOn(fake.backend) + + expect(await listFiles()).toBe( + '.github/ci.yml\nimage.png\nlong.txt\nsrc/a.ts\nsrc/b.ts', + ) + expect(await grep({ pattern: 'two' })).toBe( + [ + '.github/ci.yml:1: two: yes', + `long.txt:1: two${'x'.repeat(1997)}`, + 'src/a.ts:2: two', + 'src/b.ts:1: const two = 2', + ].join('\n'), + ) + // The engine order, and rg looked for once. + expect(fake.commands).toEqual(commands) + + expect(await listFiles({ pattern: 'src/**' })).toBe('src/a.ts\nsrc/b.ts') + expect(await grep({ pattern: 'two', glob: 'src/**' })).toBe( + 'src/a.ts:2: two\nsrc/b.ts:1: const two = 2', + ) + }, + ) + + it('runs the bundled rg with the host backend, and does not look on the PATH', async () => { + const exec = scriptHostExec(() => ok('./a.ts\0')) + expect(await searchOn(hostBackend, tmpdir()).listFiles()).toBe('a.ts') + const commands = exec.mock.calls.map(([command]) => command) + expect(commands).toHaveLength(1) + const bundled = quoteArg(process.execPath, process.platform) + expect(commands[0]?.startsWith(`${bundled} --files `)).toBe(true) + }) + + it('looks for rg on the PATH when the bundled binary is missing', async () => { + ripgrep.rgPath = join(tmpdir(), 'no-such-rg-binary') + const exec = scriptHostExec((command) => { + if (command === 'rg --version') return ok('ripgrep 14.1.0') + return command.startsWith('rg --files ') ? ok('./a.ts\0') : undefined + }) + expect(await searchOn(hostBackend, tmpdir()).listFiles()).toBe('a.ts') + expect(exec.mock.calls.map(([command]) => command)).toEqual([ + 'rg --version', + expect.stringMatching(/^rg --files /), + ]) + }) + + it.each([ + { + exitCode: 2, + stderr: 'regex parse error', + message: 'The search failed: regex parse error', + }, + { + exitCode: 124, + stderr: '', + message: 'The search failed: it took longer than 30 seconds', + }, + ])( + 'throws when rg ends with exit code $exitCode and no output', + async ({ exitCode, stderr, message }) => { + const fake = fakeBackend({}, (command) => + command === 'rg --version' + ? ok('ripgrep 14.1.0') + : { exitCode, stdout: '', stderr }, + ) + await expect( + searchOn(fake.backend).grep({ pattern: 'two' }), + ).rejects.toThrow(message) + }, + ) + + it('refuses a pattern with a line break', async () => { + const { grep } = searchOn(fakeBackend({}).backend) + await expect(grep({ pattern: 'a\nb' })).rejects.toThrow( + 'The pattern must be one line.', + ) + }) + + it('rg: grep drops the lines of a protected file, and lists the files to count them', async () => { + const fake = fakeBackend( + { '.env': 'two=1\n', 'a.ts': 'two\n' }, + (command) => { + if (command === 'rg --version') return ok('ripgrep 14.1.0') + if (command.startsWith('rg --files ')) return ok('./.env\0./a.ts\0') + return command.startsWith('rg --line-number ') + ? ok('./.env\x001:two=1\n./a.ts\x001:two\n') + : undefined + }, + ) + // Rules that ask about every .env file, like permissions() does. + const { grep } = searchOn(fake.backend, root, { + ask: async () => 'n', + mode: () => 'default', + rules: () => (_tool, paths) => + paths.some((path) => path.endsWith('.env')) ? 'ask' : 'allow', + }) + + expect(await grep({ pattern: 'two' })).toBe( + 'a.ts:1: two\n[Skipped 1 file that the permission rules protect.]', + ) + expect(fake.commands).toEqual([ + 'rg --version', + expect.stringMatching(/^rg --files /), + expect.stringMatching(/^rg --line-number /), + ]) + }) +}) + +describe('quoting for the shell of the backend', () => { + /** rg on the PATH answers every rg command with no output. */ + const rgOnPath = (command: string) => + command.startsWith('rg ') ? ok('') : undefined + + /** The rg command that a grep for `a b` runs on `backend`. */ + async function grepCommand(backend: WorkspaceBackend) { + await searchOn(backend).grep({ pattern: 'a b' }) + return ( + exec.mock.calls + .map(([command]) => command) + .find((command) => command.startsWith('rg --line-number')) ?? '' + ) + } + const exec = vi.fn(async (command: string) => rgOnPath(command) ?? NOT_FOUND) + + afterEach(() => { + exec.mockClear() + }) + + it.each([ + { shell: undefined, quoted: `'a b'` }, + { shell: 'sh' as const, quoted: `'a b'` }, + { shell: 'cmd' as const, quoted: '^"a^ b^"' }, + ])('quotes the pattern for shell $shell', async ({ shell, quoted }) => { + const backend = { ...fakeBackend({}).backend, shell, exec } + expect(await grepCommand(backend)).toContain(` -e ${quoted} .`) + }) + + it('quotes for the shell of this host in a copy of hostBackend', async () => { + const quoted = process.platform === 'win32' ? '^"a^ b^"' : `'a b'` + expect(await grepCommand({ ...hostBackend, exec })).toContain( + ` -e ${quoted} .`, + ) + }) +}) + +describe('limits', () => { + /** `count` files named `f0000.txt` and up, each with `text`. */ + const numbered = (count: number, text = '') => + Object.fromEntries( + Array.from({ length: count }, (_, index) => [ + `f${String(index).padStart(4, '0')}.txt`, + text, + ]), + ) + + it('lists 1000 files, then a note', async () => { + const lines = ( + await searchOn(fakeBackend(numbered(1001)).backend).listFiles() + ).split('\n') + expect(lines).toHaveLength(1001) + expect(lines[999]).toBe('f0999.txt') + expect(lines[1000]).toBe(FILES_NOTE) + }) + + it('shows 200 matches, then a note', async () => { + const fake = fakeBackend({ 'a.txt': 'two\n'.repeat(201) }) + const lines = (await searchOn(fake.backend).grep({ pattern: 'two' })).split( + '\n', + ) + expect(lines).toHaveLength(201) + expect(lines[199]).toBe('a.txt:200: two') + expect(lines[200]).toBe(MATCHES_NOTE) + }) + + it('says so when the JS walk stops before the end', async () => { + const fake = fakeBackend({ ...numbered(5001), 'z.md': '' }) + expect(await searchOn(fake.backend).listFiles({ pattern: '*.md' })).toBe( + `No files.\n${FILES_NOTE}`, + ) + }) +}) + +describe('on the host', () => { + /** The host backend, but rg is never found. */ + const withoutRg: WorkspaceBackend = { + ...hostBackend, + exec: (command, options) => + command.startsWith('rg ') + ? Promise.resolve(NOT_FOUND) + : hostBackend.exec(command, options), + } + /** The host backend, but no command runs, so the JS walk runs. */ + const walkOnly: WorkspaceBackend = { + ...hostBackend, + exec: async () => NOT_FOUND, + } + const put = (dir: string, path: string, text: string) => + hostBackend.writeFile(join(dir, path), text) + + it.skipIf(!hasGit)( + 'leaves out what .gitignore names, through git ls-files', + { timeout: SPAWN_TIMEOUT }, + async () => { + const dir = await mkdtemp(join(tmpdir(), 'coding-search-git-')) + try { + await put(dir, '.gitignore', 'ignored/\n*.log\n') + await put(dir, 'src/a.ts', 'two\n') + await put(dir, 'ignored/b.ts', 'two\n') + await put(dir, 'debug.log', 'two\n') + await put(dir, 'node_modules/pkg/i.js', 'two\n') + const init = await hostBackend.exec('git init -q', { cwd: dir }) + expect(init.exitCode).toBe(0) + + const { listFiles, grep } = searchOn(withoutRg, dir) + expect(await listFiles()).toBe('.gitignore\nsrc/a.ts') + expect(await grep({ pattern: 'two' })).toBe('src/a.ts:1: two') + } finally { + await rm(dir, { recursive: true, force: true }) + } + }, + ) + + describe('links to outside the workspace', () => { + let dir = '' + let outside = '' + beforeAll(async () => { + dir = await mkdtemp(join(tmpdir(), 'coding-search-walk-')) + outside = await mkdtemp(join(tmpdir(), 'coding-search-outside-')) + await put(dir, 'a.txt', 'two\n') + await put(outside, 'secret.txt', 'two\n') + // A junction on Windows needs no admin rights. Elsewhere it is a link. + await symlink(outside, join(dir, 'linked'), 'junction') + // A link to a file can need admin rights on Windows. Without it, the + // folder link still covers the tests. + await symlink( + join(outside, 'secret.txt'), + join(dir, 'secret-link.txt'), + 'file', + ).catch(() => undefined) + }) + afterAll(async () => { + await rm(outside, { recursive: true, force: true }) + await rm(dir, { recursive: true, force: true }) + }) + + it('the JS walk leaves the links out', async () => { + const { listFiles, grep } = searchOn(walkOnly, dir) + expect(await listFiles()).toBe('a.txt') + expect(await grep({ pattern: 'two' })).toBe('a.txt:1: two') + }) + + it.skipIf(!hasGit)( + 'the git list leaves the links out, and the files in a linked folder', + { timeout: SPAWN_TIMEOUT }, + async () => { + const init = await hostBackend.exec('git init -q', { cwd: dir }) + expect(init.exitCode).toBe(0) + const { listFiles, grep } = searchOn(withoutRg, dir) + expect(await listFiles()).toBe('a.txt') + expect(await grep({ pattern: 'two' })).toBe('a.txt:1: two') + }, + ) + }) +}) + +describe('quoteArg', () => { + const hostile = [ + { + value: '"; rm -rf x', + posix: `'"; rm -rf x'`, + win32: '^"\\^"^;^ rm^ -rf^ x^"', + }, + { value: '$(x)', posix: `'$(x)'`, win32: '^"$^(x^)^"' }, + { value: '`x`', posix: "'`x`'", win32: '^"^`x^`^"' }, + { value: '%PATH%', posix: `'%PATH%'`, win32: '^"^%PATH^%^"' }, + { value: 'a & b', posix: `'a & b'`, win32: '^"a^ ^&^ b^"' }, + { value: 'a | b', posix: `'a | b'`, win32: '^"a^ ^|^ b^"' }, + { value: `it's`, posix: `'it'\\''s'`, win32: `^"it's^"` }, + { value: 'C:\\dir\\', posix: `'C:\\dir\\'`, win32: '^"C:\\dir\\\\^"' }, + ] + + it.each(hostile)( + 'quotes $value for sh and cmd.exe', + ({ value, posix, win32 }) => { + expect(quoteArg(value, 'linux')).toBe(posix) + expect(quoteArg(value, 'win32')).toBe(win32) + }, + ) + + it('refuses a line break for cmd.exe', () => { + expect(() => quoteArg('a\nb', 'win32')).toThrow('line break') + }) + + it( + 'passes hostile values through the real shell of this host unchanged', + { timeout: SPAWN_TIMEOUT }, + async () => { + const values = [...hostile.map((row) => row.value), ''] + const quote = (value: string) => quoteArg(value, process.platform) + const script = 'console.log(JSON.stringify(process.argv.slice(1)))' + const command = [process.execPath, '-e', script, ...values] + .map(quote) + .join(' ') + const result = await hostBackend.exec(command) + expect(JSON.parse(result.stdout)).toEqual(values) + }, + ) +}) diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts new file mode 100644 index 0000000000..9ec08635ca --- /dev/null +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -0,0 +1,736 @@ +import { + mkdir, + mkdtemp, + readFile, + readdir, + rm, + 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 { toolDefinition } from '@tanstack/ai' +import { memoryLogStore, memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { hostBackend } from '../src/first-party/coding/backend' +import { quoteArg } from '../src/first-party/coding/search' +import { snapshots } from '../src/first-party/coding/snapshots' +import { + after, + gate, + messageTexts, + mockAdapter, + text, + toolCall, +} from './helpers' +import type { HarnessPersistence, HarnessSession } from '../src' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' +import type { Reply } from './helpers' + +const THREAD = 't1' +const hasGit = (await hostBackend.exec('git --version')).exitCode === 0 + +let base = '' +let root = '' +let dataDir = '' + +beforeEach(async () => { + // Each test gets its own folder. The workspace and the data folder are in it. + base = await mkdtemp(join(tmpdir(), 'tanstack-snapshots-')) + root = join(base, 'work') + dataDir = join(base, 'data') + await mkdir(root) +}) + +afterEach(async () => { + vi.restoreAllMocks() + await rm(base, { recursive: true, force: true }) +}) + +/** Write `text` to `path` in the workspace. Parent folders are created. */ +async function put(path: string, text: string) { + const full = join(root, path) + await mkdir(dirname(full), { recursive: true }) + await writeFile(full, text) +} + +const read = (path: string) => readFile(join(root, path), 'utf8') + +/** One turn: the model calls `change`, then answers `answer`. */ +const editTurn = (answer: string): Array => [ + () => toolCall('change', {}), + () => text(answer), +] + +/** + * Open a session with `snapshots()` and a `change` tool. Each call of the + * tool runs the next function of `changes`. + */ +async function open(options: { + replies: Array + changes?: Array<() => Promise> + persistence?: HarnessPersistence + backend?: WorkspaceBackend +}) { + const changes = options.changes ?? [] + const change = toolDefinition({ + name: 'change', + description: 'Change the files', + }).server(async () => { + await changes.shift()?.() + return 'changed' + }) + const plugin = snapshots({ root, dataDir, backend: options.backend }) + const { adapter, calls } = mockAdapter(options.replies) + const host = createHarnessHost({ + persistence: options.persistence ?? memoryPersistence(), + }) + const session = await host.open( + defineHarness({ + name: 'test/snapshots', + adapter, + tools: [change], + plugins: () => [plugin], + }), + { threadId: THREAD }, + ) + return { host, session, plugin, calls } +} + +/** The steps that the plugin state keeps. */ +function savedSteps(session: HarnessSession) { + const steps = session.snapshot().plugins['tanstack/snapshots'] + if (!Array.isArray(steps)) throw new Error('No steps in the plugin state.') + return steps +} + +/** Every file under `folder`, by path, with its bytes. */ +async function filesUnder(folder: string) { + const entries = await readdir(folder, { + recursive: true, + withFileTypes: true, + }) + const files = new Map() + for (const entry of entries) { + if (!entry.isFile()) continue + const full = join(entry.parentPath, entry.name) + files.set(full, await readFile(full)) + } + return files +} + +const TREE = /^[0-9a-f]{40}$/ + +/** A backend that keeps no files. Each command goes to `exec`. */ +function fakeBackend( + exec: WorkspaceBackend['exec'], + shell?: WorkspaceBackend['shell'], +) { + const fail = () => Promise.reject(new Error('Not used.')) + const backend: WorkspaceBackend = { + shell, + readFile: fail, + writeFile: async () => undefined, + stat: fail, + readdir: fail, + exec, + } + return backend +} + +describe.skipIf(!hasGit)('snapshots', () => { + // Each git call starts a shell and a process, which is slow on Windows. + it('starts few git commands for each step', async () => { + await put('a.txt', 'one') + const { host, session } = await open({ + replies: [ + () => toolCall('change', {}), + () => toolCall('change', {}), + () => text('Done.'), + ], + changes: [async () => put('a.txt', 'two'), async () => undefined], + }) + const exec = vi.spyOn(hostBackend, 'exec') + + await session.prompt('Edit the file') + + const commands = exec.mock.calls.map(([command]) => command) + // `git --version` and `git init` once, the turn start, and one snapshot + // after each tool step. Step starts reuse the snapshot before them. + expect( + commands.filter((command) => /write-tree/.test(command)), + ).toHaveLength(3) + // The second step changed nothing, so it needs no diff. + expect(commands.filter((command) => / diff /.test(command))).toHaveLength(1) + expect(commands.length).toBeLessThanOrEqual(6) + await host.close() + }) + + it('keeps the changed files of each step in the plugin state', async () => { + await put('a.txt', 'one') + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [ + async () => { + await put('a.txt', 'two') + await put('src/b.txt', 'new') + }, + ], + }) + + await session.prompt('Edit the files') + + const steps = savedSteps(session) + expect(steps).toMatchObject([{ files: ['a.txt', 'src/b.txt'] }]) + expect(steps[0].from).toMatch(TREE) + expect(steps[0].to).toMatch(TREE) + expect(steps[0].to).not.toBe(steps[0].from) + await host.close() + }) + + it('writes each step as a log record on a durable host', async () => { + await put('a.txt', 'one') + const { runs, metadata } = memoryPersistence().stores + const log = memoryLogStore() + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [() => put('a.txt', 'two')], + persistence: { stores: { log, runs, metadata } }, + }) + + await session.prompt('Edit a.txt') + + const entries = await log.read(THREAD, { after: 0 }) + const records = entries + .map((entry) => entry.record) + .filter((record) => record.type === 'tanstack/snapshots:step') + expect(records).toMatchObject([{ files: ['a.txt'] }]) + await host.close() + }) + + it('gives the changed files and a unified diff between two trees', async () => { + await put('a.txt', 'one\n') + const { host, session, plugin } = await open({ + replies: editTurn('Done.'), + changes: [() => put('a.txt', 'two\n')], + }) + await session.prompt('Edit a.txt') + const [step] = savedSteps(session) + + const between = await plugin.diff(step.from, step.to) + expect(between.files).toEqual(['a.txt']) + expect(between.patch).toContain('\n-one\n+two\n') + + // Without `to`, the diff is against the files now. + await put('c.txt', 'later\n') + const sinceStep = await plugin.diff(step.to) + expect(sinceStep.files).toEqual(['c.txt']) + expect(sinceStep.patch).toContain('+later') + await host.close() + }) + + it('/undo restores the files and removes the last turn', async () => { + await put('a.txt', 'one') + await put('sub/keep.txt', 'keep') + const { host, session } = await open({ + replies: [() => text('Hi.'), ...editTurn('Done.')], + changes: [ + async () => { + await put('a.txt', 'two') + await put('new/deep/c.txt', 'added') + await rm(join(root, 'sub/keep.txt')) + }, + ], + }) + await session.prompt('Hello') + await session.prompt('Edit the files') + + expect(await session.command('undo')).toBe( + 'Undid the last turn. Files restored: 3.', + ) + expect(await read('a.txt')).toBe('one') + expect(await read('sub/keep.txt')).toBe('keep') + expect((await readdir(root)).sort()).toEqual(['a.txt', 'sub']) + expect(await session.transcript()).toMatchObject([ + { role: 'user', content: 'Hello' }, + { role: 'assistant', content: 'Hi.' }, + ]) + expect(await session.command('undo')).toBe('Nothing to undo.') + await host.close() + }) + + it('/redo brings back the files and the turn', async () => { + await put('a.txt', 'one') + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [ + async () => { + await put('a.txt', 'two') + await put('new/c.txt', 'added') + }, + ], + }) + await session.prompt('Edit the files') + const edited = await session.transcript() + await session.command('undo') + + expect(await session.command('redo')).toBe('Redid the last turn.') + expect(await read('a.txt')).toBe('two') + expect(await read('new/c.txt')).toBe('added') + expect(await session.transcript()).toEqual(edited) + // The turn is back, so it can be undone again. + await session.command('undo') + expect(await read('a.txt')).toBe('one') + await host.close() + }) + + it('/undo and /redo keep the files that the turn did not change', async () => { + await put('a.txt', 'one') + await put('c.txt', 'one') + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [() => put('a.txt', 'two')], + }) + await session.prompt('Edit a.txt') + // The user adds and edits files after the turn. + await put('b.txt', 'mine') + await put('c.txt', 'mine') + + expect(await session.command('undo')).toBe( + 'Undid the last turn. Files restored: 1.', + ) + expect(await read('a.txt')).toBe('one') + expect(await read('b.txt')).toBe('mine') + expect(await read('c.txt')).toBe('mine') + + await put('c.txt', 'later') + expect(await session.command('redo')).toBe('Redid the last turn.') + expect(await read('a.txt')).toBe('two') + expect(await read('b.txt')).toBe('mine') + expect(await read('c.txt')).toBe('later') + await host.close() + }) + + it('/undo removes a file that the turn added, and /redo brings it back', async () => { + await put('a.txt', 'one') + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [() => put('new/c.txt', 'added')], + }) + await session.prompt('Add a file') + + await session.command('undo') + expect(await readdir(root)).toEqual(['a.txt']) + // A file that the user adds between /undo and /redo stays. + await put('mine.txt', 'mine') + await session.command('redo') + expect(await read('new/c.txt')).toBe('added') + expect(await read('mine.txt')).toBe('mine') + await host.close() + }) + + it('restores paths with a space or a glob character', async () => { + await put('my file.txt', 'one') + await put('x.txt', 'one') + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [ + async () => { + await put('my file.txt', 'two') + await put('[x].txt', 'added') + }, + ], + }) + await session.prompt('Edit the files') + await put('x.txt', 'mine') + // No file has the name `[x].txt` now. As a glob, it matches `x.txt`, so + // git must read it as a plain name. + await rm(join(root, '[x].txt')) + + expect(await session.command('undo')).toBe( + 'Undid the last turn. Files restored: 2.', + ) + expect(await read('my file.txt')).toBe('one') + expect(await read('x.txt')).toBe('mine') + expect((await readdir(root)).sort()).toEqual(['my file.txt', 'x.txt']) + await host.close() + }) + + it('clears redo when a new turn starts', async () => { + await put('a.txt', 'one') + const { host, session } = await open({ + replies: [...editTurn('Done.'), () => text('Sure.')], + changes: [() => put('a.txt', 'two')], + }) + await session.prompt('Edit a.txt') + await session.command('undo') + await session.prompt('Something else') + + expect(await session.command('redo')).toBe('Nothing to redo.') + expect(await read('a.txt')).toBe('one') + await host.close() + }) + + it("leaves the project's own .git untouched", async () => { + // GIT_DIR keeps this init in the test folder, even inside a git hook. + const projectGit = join(root, '.git') + const init = await hostBackend.exec( + `git init -q ${quoteArg(root, process.platform)}`, + { env: { GIT_DIR: projectGit } }, + ) + expect(init.exitCode).toBe(0) + await put('a.txt', 'one') + const before = await filesUnder(projectGit) + const { host, session } = await open({ + replies: editTurn('Done.'), + changes: [() => put('b.txt', 'new')], + }) + + await session.prompt('Add b.txt') + await session.command('undo') + + expect(savedSteps(session)).toMatchObject([{ files: ['b.txt'] }]) + expect(await filesUnder(projectGit)).toEqual(before) + await host.close() + }) +}) + +/** Three turns. Each one changes the files once, with its own tool call. */ +const threeTurns = () => ({ + replies: (['First.', 'Second.', 'Third.'] as const).flatMap( + (answer, index): Array => [ + () => toolCall('change', {}, `call-${index}`), + () => text(answer), + ], + ), + changes: [ + () => put('a.txt', 'two'), + () => put('a.txt', 'three'), + () => put('b.txt', 'new'), + ], +}) + +/** Run the three turns of {@link threeTurns}. Gives the id of the first answer. */ +async function runThreeTurns(session: HarnessSession) { + await session.prompt('One') + await session.prompt('Two') + await session.prompt('Three') + const transcript = await session.transcript() + const first = transcript.find((message) => message.content === 'First.') + if (!first?.id) throw new Error('The first answer has no id.') + return { transcript, firstId: first.id } +} + +describe.skipIf(!hasGit)('session.revert', () => { + it('reverts two turns back: the files and the transcript', async () => { + await put('a.txt', 'one') + const { host, session } = await open(threeTurns()) + const { transcript, firstId } = await runThreeTurns(session) + expect(transcript).toHaveLength(12) + + expect(await session.revert(firstId)).toMatchObject({ status: 'accepted' }) + + expect(await read('a.txt')).toBe('two') + expect(await readdir(root)).toEqual(['a.txt']) + expect(await session.transcript()).toEqual(transcript.slice(0, 4)) + await host.close() + }) + + it('unrevert brings back the files and the messages', async () => { + await put('a.txt', 'one') + const { host, session } = await open(threeTurns()) + const { transcript, firstId } = await runThreeTurns(session) + await session.revert(firstId) + + expect(await session.unrevert()).toMatchObject({ status: 'accepted' }) + + expect(await read('a.txt')).toBe('three') + expect(await read('b.txt')).toBe('new') + expect(await session.transcript()).toEqual(transcript) + await host.close() + }) + + it('the next prompt drops the hidden messages for good', async () => { + await put('a.txt', 'one') + const turns = threeTurns() + const { host, session, calls } = await open({ + ...turns, + replies: [...turns.replies, () => text('Fourth.')], + }) + const { transcript, firstId } = await runThreeTurns(session) + await session.revert(firstId) + + await session.prompt('Four') + + expect(messageTexts(calls.at(-1))).toEqual([ + 'One', + 'null', + 'changed', + 'First.', + 'Four', + ]) + expect(await session.transcript()).toMatchObject([ + ...transcript.slice(0, 4), + { role: 'user', content: 'Four' }, + { role: 'assistant', content: 'Fourth.' }, + ]) + // The revert ended, so unrevert changes nothing. + await session.unrevert() + expect(await read('a.txt')).toBe('two') + expect(await session.transcript()).toHaveLength(6) + await host.close() + }) + + it('is refused while a turn runs, and for an unknown message', async () => { + const wait = gate() + const { host, session } = await open({ + replies: [() => text('Hi.'), after(wait.opened, 'Later.')], + }) + await session.prompt('Hello') + const [message] = await session.transcript() + const running = session.prompt('Again') + await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) + + expect(await session.revert(message?.id ?? '')).toMatchObject({ + status: 'rejected', + reason: 'busy', + }) + wait.open() + await running + expect(await session.revert('no-such-id')).toMatchObject({ + status: 'rejected', + reason: 'unknown_message', + }) + expect(await session.transcript()).toHaveLength(4) + await host.close() + }) + + it('keeps the revert after a restart on a durable host', async () => { + await put('a.txt', 'one') + const { runs, metadata } = memoryPersistence().stores + const persistence = { stores: { log: memoryLogStore(), runs, metadata } } + const first = await open({ ...threeTurns(), persistence }) + const { transcript, firstId } = await runThreeTurns(first.session) + await first.session.revert(firstId) + await first.host.close() + + const second = await open({ replies: [], persistence }) + + expect(await second.session.transcript()).toEqual(transcript.slice(0, 4)) + await second.session.unrevert() + expect(await read('a.txt')).toBe('three') + expect(await second.session.transcript()).toEqual(transcript) + await second.host.close() + }) + + it('keeps the revert after a restart on a host without a log', async () => { + await put('a.txt', 'one') + const persistence = memoryPersistence() + const first = await open({ ...threeTurns(), persistence }) + const { transcript, firstId } = await runThreeTurns(first.session) + await first.session.revert(firstId) + await first.host.close() + + const second = await open({ replies: [], persistence }) + + expect(await second.session.transcript()).toEqual(transcript.slice(0, 4)) + await second.session.unrevert() + expect(await read('b.txt')).toBe('new') + await second.host.close() + }) +}) + +describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { + it('keeps a user edit between two hidden turns', async () => { + await put('a.txt', 'one') + await put('c.txt', 'one') + const { host, session } = await open(threeTurns()) + await session.prompt('One') + await session.prompt('Two') + // The user edits a file that no turn changes. + await put('c.txt', 'mine') + await session.prompt('Three') + const first = (await session.transcript()).find( + (message) => message.content === 'First.', + ) + + await session.revert(first?.id ?? '') + + expect(await read('a.txt')).toBe('two') + expect(await read('c.txt')).toBe('mine') + expect((await readdir(root)).sort()).toEqual(['a.txt', 'c.txt']) + await host.close() + }) + + it('keeps a user edit after the last hidden step on unrevert', async () => { + await put('a.txt', 'one') + const { host, session } = await open(threeTurns()) + const { firstId } = await runThreeTurns(session) + // The user edits a file that a hidden turn changed. + await put('a.txt', 'mine') + + await session.revert(firstId) + expect(await read('a.txt')).toBe('two') + await session.unrevert() + + expect(await read('a.txt')).toBe('mine') + expect(await read('b.txt')).toBe('new') + await host.close() + }) + + it('holds a prompt that arrives during a revert', async () => { + await put('a.txt', 'one') + const wait = gate() + const reached = gate() + let blocking = false + // A backend that stops at the git checkout of the revert. + const backend: WorkspaceBackend = { + ...hostBackend, + exec: async (command, options) => { + if (blocking && / checkout /.test(command)) { + reached.open() + await wait.opened + } + return hostBackend.exec(command, options) + }, + } + const turns = threeTurns() + const { host, session } = await open({ + ...turns, + backend, + replies: [...turns.replies, () => text('Fourth.')], + }) + const { firstId } = await runThreeTurns(session) + blocking = true + const reverting = session.revert(firstId) + await reached.opened + + const prompting = session.prompt('Four') + await vi.waitFor(() => expect(session.snapshot().queuedTurns).toBe(1)) + expect(session.snapshot().status).toBe('idle') + blocking = false + wait.open() + await reverting + await prompting + + expect(await session.transcript()).toMatchObject([ + { role: 'user', content: 'One' }, + { role: 'assistant' }, + { role: 'tool' }, + { role: 'assistant', content: 'First.' }, + { role: 'user', content: 'Four' }, + { role: 'assistant', content: 'Fourth.' }, + ]) + await host.close() + }) + + it('adds a note that arrives during a revert after the commit', async () => { + await put('a.txt', 'one') + await mkdir(join(root, 'sub')) + const turns = threeTurns() + const { host, session } = await open({ + ...turns, + replies: [...turns.replies, () => text('Fourth.')], + }) + const { firstId } = await runThreeTurns(session) + await session.revert(firstId) + + await session.configure({ cwd: 'sub' }) + await session.prompt('Four') + + expect( + (await session.transcript()).slice(4).map((message) => message.content), + ).toEqual([ + 'The working folder is now sub. Paths are relative to it.', + 'Four', + 'Fourth.', + ]) + await host.close() + }) + + it('adds a note that arrives during a revert after unrevert', async () => { + await put('a.txt', 'one') + await mkdir(join(root, 'sub')) + const { host, session } = await open(threeTurns()) + const { transcript, firstId } = await runThreeTurns(session) + await session.revert(firstId) + + await session.configure({ cwd: 'sub' }) + await session.unrevert() + + expect( + (await session.transcript()).map((message) => message.content), + ).toEqual([ + ...transcript.map((message) => message.content), + 'The working folder is now sub. Paths are relative to it.', + ]) + await host.close() + }) + + it('refuses /undo and /redo while a revert stands', async () => { + await put('a.txt', 'one') + const { host, session } = await open(threeTurns()) + const { transcript, firstId } = await runThreeTurns(session) + await session.revert(firstId) + + const refused = 'A revert stands. Run unrevert, or send a prompt, first.' + expect(await session.command('undo')).toBe(refused) + expect(await session.command('redo')).toBe(refused) + + expect(await read('a.txt')).toBe('two') + await session.unrevert() + expect(await session.transcript()).toEqual(transcript) + expect(await read('a.txt')).toBe('three') + await host.close() + }) +}) + +describe('snapshots on a cmd.exe backend', () => { + it('quotes the git arguments for cmd.exe', async () => { + const commands: Array = [] + const backend = fakeBackend(async (command) => { + commands.push(command) + return { exitCode: 0, stdout: '', stderr: '' } + }, 'cmd') + const { host, session } = await open({ + replies: [() => text('Hi.')], + backend, + }) + + await session.prompt('Hello') + + // `git --version`, then the init of the shadow repository. + expect(commands[1]).toMatch(/^git init --bare -q \^".+\^"$/) + await host.close() + }) +}) + +describe('snapshots without git', () => { + it('adds no commands, warns once, and the turn still runs', async () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => {}) + // Every command fails, like a machine with no git. + const { host, session, plugin } = await open({ + replies: [() => text('Hi.')], + backend: fakeBackend(async () => ({ + exitCode: 127, + stdout: '', + stderr: 'not found', + })), + }) + + await session.prompt('Hello') + + expect(session.commands()).toEqual([]) + expect(await session.transcript()).toMatchObject([ + { role: 'user', content: 'Hello' }, + { role: 'assistant', content: 'Hi.' }, + ]) + expect(warn.mock.calls).toEqual([ + ['snapshots: git was not found. Snapshots are off.'], + ]) + await expect(plugin.diff('abc')).rejects.toThrow( + 'snapshots: git is not available.', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/coding-web.test.ts b/packages/ai-harness/tests/coding-web.test.ts new file mode 100644 index 0000000000..547b343ee5 --- /dev/null +++ b/packages/ai-harness/tests/coding-web.test.ts @@ -0,0 +1,384 @@ +import { createServer } from 'node:http' +import { + afterAll, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from 'vitest' +import { webTools } from '../src/first-party/coding/web' +import type { AddressInfo } from 'node:net' +import type { + SearchProvider, + WebToolsOptions, +} from '../src/first-party/coding/web' + +// DNS is the network boundary. Each lookup of a name takes its next +// answer, and the last answer repeats. +const dns = vi.hoisted(() => ({ + answers: new Map>>(), +})) + +vi.mock('node:dns/promises', () => ({ + lookup: async (hostname: string) => { + const answers = dns.answers.get(hostname) ?? [] + const addresses = answers.length > 1 ? answers.shift() : answers[0] + if (!addresses) throw new Error(`getaddrinfo ENOTFOUND ${hostname}`) + return addresses.map((address) => ({ + address, + family: address.includes(':') ? 6 : 4, + })) + }, +})) + +/** `host` resolves to `answers`, one answer a lookup. */ +function resolves(host: string, ...answers: Array>) { + dns.answers.set(host, answers) +} + +const PUBLIC_IP = '93.184.216.34' +const MB5 = 5 * 1024 * 1024 +const REFUSED = 'webfetch does not read private or local hosts' + +type Route = (init: RequestInit | undefined) => Response | Promise + +/** A route that answers with `body` and the `content-type` `type`. */ +function page(body: BodyInit, type: string) { + return () => new Response(body, { headers: { 'content-type': type } }) +} + +/** A route that redirects to `location`. */ +function redirect(location: string) { + return () => new Response(null, { status: 302, headers: { location } }) +} + +/** A fetch that answers from `routes`, by URL. `urls` lists each fetch. */ +function fakeFetch(routes: Record) { + const urls: Array = [] + const fetch: typeof globalThis.fetch = async (input, init) => { + const url = String(input) + urls.push(url) + const route = routes[url] + if (!route) throw new Error(`No route for ${url}`) + return route(init) + } + return { fetch, urls } +} + +/** Run the tool `name` from `webTools(options)` once. */ +async function run(name: string, args: object, options: WebToolsOptions = {}) { + const tool = webTools(options).find((item) => item.name === name) + if (!tool?.execute) throw new Error(`No ${name} tool`) + return tool.execute(args) +} + +/** Fetch `url` with a fetch that knows `routes`. */ +function webfetch( + args: { url: string; format?: string; timeoutMs?: number }, + routes: Record, + options: WebToolsOptions = {}, +) { + const { fetch } = fakeFetch(routes) + return run('webfetch', args, { fetch, ...options }) +} + +beforeEach(() => { + dns.answers.clear() + resolves('example.com', [PUBLIC_IP]) +}) + +describe('webfetch formats', () => { + const html = + 'Doc

Hello

Tom & Jerry <3 AB

' + const routes = { + 'https://example.com/': page(html, 'text/html; charset=utf-8'), + } + + it('gives back HTML as markdown when turndown is installed', async () => { + const result = await webfetch({ url: 'https://example.com/' }, routes) + expect(result).toBe('Doc\n\n# Hello\n\nTom & **Jerry** <3 AB') + }) + + it('gives back text without tags, scripts, or styles', async () => { + const result = await webfetch( + { url: 'https://example.com/', format: 'text' }, + routes, + ) + expect(result).toBe('Doc\nHello\nTom & Jerry <3 AB') + }) + + it('gives back the HTML as is with format html', async () => { + const result = await webfetch( + { url: 'https://example.com/', format: 'html' }, + routes, + ) + expect(result).toBe(html) + }) + + it('gives back JSON as it is', async () => { + const result = await webfetch( + { url: 'https://example.com/data' }, + { 'https://example.com/data': page('{"ok":true}', 'application/json') }, + ) + expect(result).toBe('{"ok":true}') + }) + + it('without turndown, refuses markdown and gives text by default', async () => { + vi.doMock('turndown', () => { + throw new Error('Cannot find package turndown') + }) + try { + await expect( + webfetch({ url: 'https://example.com/', format: 'markdown' }, routes), + ).rejects.toThrow('Install turndown to get markdown') + const result = await webfetch({ url: 'https://example.com/' }, routes) + expect(result).toBe('Doc\nHello\nTom & Jerry <3 AB') + } finally { + vi.doUnmock('turndown') + } + }) +}) + +describe('webfetch limits', () => { + it('stops reading at 5 MB and says the page was cut', async () => { + const body = 'a'.repeat(MB5 + 100) + const result = await webfetch( + { url: 'https://example.com/big.txt' }, + { 'https://example.com/big.txt': page(body, 'text/plain') }, + ) + expect(result).toBe( + `${'a'.repeat(MB5)}\n\n[Cut at 5 MB. The rest was not read.]`, + ) + }) + + it('stops after the timeout', async () => { + const hang: Route = (init) => + new Promise((_, reject) => { + init?.signal?.addEventListener('abort', () => + reject(new Error('aborted')), + ) + }) + await expect( + webfetch( + { url: 'https://example.com/slow', timeoutMs: 20 }, + { 'https://example.com/slow': hang }, + ), + ).rejects.toThrow('https://example.com/slow timed out after 20 ms.') + }) + + it('refuses binary content', async () => { + const png = new Uint8Array([0x89, 0x50, 0x4e, 0x47]) + await expect( + webfetch( + { url: 'https://example.com/logo.png' }, + { 'https://example.com/logo.png': page(png, 'image/png') }, + ), + ).rejects.toThrow( + 'https://example.com/logo.png is image/png. webfetch reads HTML, text, and JSON only.', + ) + }) + + it('refuses a body with no content type and a NUL byte', async () => { + const bytes = new Uint8Array([0x7f, 0x45, 0x4c, 0x46, 0x00, 0x01]) + await expect( + webfetch( + { url: 'https://example.com/tool' }, + { 'https://example.com/tool': () => new Response(bytes) }, + ), + ).rejects.toThrow('https://example.com/tool is binary.') + }) + + it('fails on an HTTP error status', async () => { + await expect( + webfetch( + { url: 'https://example.com/missing' }, + { + 'https://example.com/missing': () => + new Response('gone', { status: 404 }), + }, + ), + ).rejects.toThrow('https://example.com/missing answered with HTTP 404.') + }) +}) + +describe('webfetch hosts', () => { + it.each([ + 'http://localhost:3000/', + 'http://127.0.0.1/', + 'http://169.254.169.254/latest/meta-data/', + 'http://10.0.0.8/', + 'http://172.16.0.1/', + 'http://192.168.1.1/', + 'http://[::1]/', + 'http://[fd00::1]/', + 'http://[::ffff:127.0.0.1]/', + ])('refuses %s without fetching it', async (url) => { + const { fetch, urls } = fakeFetch({}) + await expect(run('webfetch', { url }, { fetch })).rejects.toThrow(REFUSED) + expect(urls).toEqual([]) + }) + + it('refuses a name when DNS gives a private address', async () => { + resolves('intranet.example', [PUBLIC_IP, '10.1.2.3']) + const { fetch, urls } = fakeFetch({}) + await expect( + run('webfetch', { url: 'https://intranet.example/' }, { fetch }), + ).rejects.toThrow(`${REFUSED}: intranet.example`) + expect(urls).toEqual([]) + }) + + it('refuses a name that DNS cannot resolve', async () => { + await expect( + webfetch({ url: 'https://nowhere.example/' }, {}), + ).rejects.toThrow('Could not resolve nowhere.example.') + }) + + it('reads private hosts with allowPrivateHosts', async () => { + const result = await webfetch( + { url: 'http://localhost:3000/' }, + { 'http://localhost:3000/': page('local', 'text/plain') }, + { allowPrivateHosts: true }, + ) + expect(result).toBe('local') + }) + + it('refuses a URL that is not http or https', async () => { + await expect( + webfetch({ url: 'file:///etc/passwd' }, {}, { allowPrivateHosts: true }), + ).rejects.toThrow('webfetch reads http and https URLs only, not file:') + }) +}) + +describe('webfetch redirects', () => { + it('follows a redirect to a public host', async () => { + const { fetch, urls } = fakeFetch({ + 'https://example.com/old': redirect('/new'), + 'https://example.com/new': page('moved', 'text/plain'), + }) + const result = await run( + 'webfetch', + { url: 'https://example.com/old' }, + { fetch }, + ) + expect(result).toBe('moved') + expect(urls).toEqual(['https://example.com/old', 'https://example.com/new']) + }) + + it('refuses a redirect to a private host', async () => { + const { fetch, urls } = fakeFetch({ + 'https://example.com/go': redirect('http://169.254.169.254/latest/'), + }) + await expect( + run('webfetch', { url: 'https://example.com/go' }, { fetch }), + ).rejects.toThrow(`${REFUSED}: 169.254.169.254`) + expect(urls).toEqual(['https://example.com/go']) + }) + + it('stops after 5 redirects', async () => { + const { fetch, urls } = fakeFetch({ + 'https://example.com/loop': redirect('/loop'), + }) + await expect( + run('webfetch', { url: 'https://example.com/loop' }, { fetch }), + ).rejects.toThrow('https://example.com/loop redirected more than 5 times.') + expect(urls).toHaveLength(6) + }) +}) + +describe('webfetch without a custom fetch', () => { + // A real server on 127.0.0.1. `paths` lists the requests that reach it. + const paths: Array = [] + const server = createServer((req, res) => { + paths.push(req.url ?? '') + if (req.url === '/old') { + res.writeHead(302, { location: '/new' }).end() + return + } + res.writeHead(200, { 'content-type': 'text/plain' }).end('hello') + }) + let port = 0 + + beforeAll(async () => { + await new Promise((done) => server.listen(0, '127.0.0.1', done)) + const address: AddressInfo | string | null = server.address() + port = typeof address === 'object' && address ? address.port : 0 + }) + + afterAll(async () => { + await new Promise((done) => server.close(() => done())) + }) + + beforeEach(() => { + paths.length = 0 + }) + + it('refuses a name whose answer turns private when it connects', async () => { + // The check sees a public address. The connection sees 127.0.0.1. + resolves('rebind.test', [PUBLIC_IP], ['127.0.0.1']) + await expect( + run('webfetch', { url: `http://rebind.test:${port}/` }), + ).rejects.toThrow(`${REFUSED}: rebind.test`) + expect(paths).toEqual([]) + }) + + it('reaches a private host with allowPrivateHosts and follows redirects', async () => { + resolves('rebind.test', ['127.0.0.1']) + const result = await run( + 'webfetch', + { url: `http://rebind.test:${port}/old` }, + { allowPrivateHosts: true }, + ) + expect(result).toBe('hello') + expect(paths).toEqual(['/old', '/new']) + }) +}) + +describe('websearch', () => { + /** A provider that gives back `results` and records each search. */ + function provider(results: Awaited>) { + const searches: Array<{ query: string; limit: number }> = [] + const search: SearchProvider = { + search: async (query, { limit }) => { + searches.push({ query, limit }) + return results + }, + } + return { search, searches } + } + + it('gives back the results as a numbered list', async () => { + const { search, searches } = provider([ + { + title: 'TanStack AI', + url: 'https://tanstack.com/ai', + snippet: 'Type-safe AI SDK', + }, + { title: 'TanStack Query', url: 'https://tanstack.com/query' }, + ]) + const result = await run( + 'websearch', + { query: 'tanstack', limit: 2 }, + { search }, + ) + expect(result).toBe( + '1. TanStack AI\n https://tanstack.com/ai\n Type-safe AI SDK\n\n2. TanStack Query\n https://tanstack.com/query', + ) + expect(searches).toEqual([{ query: 'tanstack', limit: 2 }]) + }) + + it('says when there are no results', async () => { + const { search } = provider([]) + const result = await run('websearch', { query: 'nothing' }, { search }) + expect(result).toBe('No results for "nothing".') + }) + + it('exists only with a provider', () => { + const { search } = provider([]) + const names = (options: WebToolsOptions) => + webTools(options).map((tool) => tool.name) + expect(names({})).toEqual(['webfetch']) + expect(names({ search })).toEqual(['webfetch', 'websearch']) + }) +}) diff --git a/packages/ai-harness/tests/edge-safety.test.ts b/packages/ai-harness/tests/edge-safety.test.ts index f585e5fd8a..0d87c8c97b 100644 --- a/packages/ai-harness/tests/edge-safety.test.ts +++ b/packages/ai-harness/tests/edge-safety.test.ts @@ -17,9 +17,13 @@ const srcDir = resolve(here, '..', 'src') /** * Node-only entry points, relative to `src`. A path that ends with `/` allows - * the whole folder. + * the whole folder: `first-party/coding/` is the Node-only `plugins/coding` entry. */ -const ALLOWED = ['first-party/files.ts', 'first-party/bound-output.ts'] +const ALLOWED = [ + 'first-party/files.ts', + 'first-party/bound-output.ts', + 'first-party/coding/', +] 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 index 1b6fc092a1..ed3044211d 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -1,4 +1,4 @@ -import { mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises' +import { mkdtemp, mkdir, readFile, 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' @@ -10,10 +10,12 @@ import { fileCommands, globToRegExp, modelPicker, + permissions, projectInstructions, todos, usage, } from '../src/first-party' +import { workspaceTools } from '../src/first-party/coding' import { messageTexts, mockAdapter, text, toolCall } from './helpers' import type { HarnessPlugin } from '../src' import type { AnyTextAdapter } from '@tanstack/ai' @@ -58,7 +60,53 @@ describe('decidePermission', () => { }) }) -describe('globToRegExp', () => { +describe('workspace tools with permissions', () => { + it('reads freely, asks before writing, and stays inside the root', async () => { + await writeFile(join(root, 'notes.txt'), 'hello\nworld\n') + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', { path: 'notes.txt' }, 'c1'), + () => + toolCall('write_file', { path: 'out.txt', content: 'written' }, 'c2'), + () => toolCall('read_file', { path: '../escape.txt' }, 'c3'), + () => text('done'), + ]) + const { host, session } = await open(adapter, [ + permissions(), + workspaceTools({ root }), + ]) + + const turn = session.prompt('work') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + expect(question?.message).toContain('write_file') + await session.answer(question!.questionId, { answer: 'once' }) + await turn + + expect(await readFile(join(root, 'out.txt'), 'utf8')).toBe('written') + const toolResults = JSON.stringify(calls[3].messages) + expect(toolResults).toContain('1\\thello') + expect(toolResults).toContain('outside the workspace') + await host.close() + }) + + it('denies edits in plan mode without asking', async () => { + const { adapter } = mockAdapter([ + () => toolCall('write_file', { path: 'x.txt', content: 'x' }, 'c1'), + () => text('ok'), + ]) + const { host, session } = await open(adapter, [ + permissions(), + workspaceTools({ root }), + ]) + expect(await session.command('mode', 'plan')).toBe('Mode: plan.') + await session.prompt('write') + await expect(readFile(join(root, 'x.txt'), 'utf8')).rejects.toThrow() + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) + 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) diff --git a/packages/ai-harness/tests/instructions.test.ts b/packages/ai-harness/tests/instructions.test.ts index 8f8ef03d1b..373e82bd5e 100644 --- a/packages/ai-harness/tests/instructions.test.ts +++ b/packages/ai-harness/tests/instructions.test.ts @@ -5,8 +5,9 @@ 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 { workspaceTools } from '../src/first-party/coding' +import { mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter, ModelMessage } from '@tanstack/ai' import type { HarnessPlugin } from '../src' // `~` in a global path is the home folder. Each test points it into its own @@ -231,4 +232,36 @@ describe('projectInstructions', () => { const [environment] = await promptsOfOneTurn({ root }) expect(environment).toMatch(/- Git repository: yes$/) }) + + it('adds a nested instruction file to read_file once', async () => { + await put('project/AGENTS.md', 'Root rule.') + await put('project/src/AGENTS.md', 'Src rule.') + await put('project/src/lib/AGENTS.md', 'Lib rule.') + await put('project/src/lib/a.ts', 'a') + await put('project/src/b.ts', 'b') + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', { path: 'src/b.ts' }, 'c1'), + () => toolCall('read_file', { path: 'src/lib/AGENTS.md' }, 'c2'), + () => toolCall('read_file', { path: 'src/lib/a.ts' }, 'c3'), + () => text('done'), + ]) + const { host, session } = await open(adapter, [ + workspaceTools({ root }), + projectInstructions({ root, env: false }), + ]) + + await session.prompt('read') + await host.close() + + // What each read added after the lines of the file. + const messages: Array = calls[3].messages + const added = messages + .filter((message) => message.role === 'tool') + .map((message) => String(message.content).split('\n\n').slice(1)) + expect(added).toEqual([ + ['Instructions from src/AGENTS.md:\nSrc rule.'], + [], + [], + ]) + }) }) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts new file mode 100644 index 0000000000..137ae6f20d --- /dev/null +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -0,0 +1,1037 @@ +import { mkdtemp, readFile, realpath, rm, symlink } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import * as nodePath from 'node:path' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { permissions } from '../src/first-party/permissions' +import { hostBackend } from '../src/first-party/coding/backend' +import { + createWorkspaceTools, + workspaceTools, +} from '../src/first-party/coding/workspace' +import { mockAdapter, text, toolCall } from './helpers' +import type { PlatformPath } from 'node:path' +import type { AnyTextAdapter, AnyTool } from '@tanstack/ai' +import type { HarnessPlugin } from '../src' +import type { PermissionRule } from '../src/first-party/permissions' +import type { WorkspaceBackend } from '../src/first-party/coding/backend' +import type { + OutsideAccess, + WorkspaceToolsOptions, +} from '../src/first-party/coding/workspace' + +const { posix } = nodePath + +/** + * A backend that keeps files in a map, by absolute POSIX path, like a Linux + * sandbox. Its commands answer with the command and the folder. `extra` + * replaces parts of it. + */ +function memoryBackend(extra: Partial = {}) { + const files = new Map() + const backend: WorkspaceBackend = { + shell: 'sh', + readFile: async (path) => { + const data = files.get(path) + if (!data) throw new Error(`No file ${path}`) + return data + }, + writeFile: async (path, data) => { + files.set( + path, + typeof data === 'string' ? new TextEncoder().encode(data) : data, + ) + }, + stat: async (path) => { + const data = files.get(path) + if (!data) return undefined + return { type: 'file', size: data.length, mtimeMs: 0 } + }, + readdir: async (path) => { + const prefix = path.endsWith('/') ? path : `${path}/` + const entries = new Map() + for (const file of files.keys()) { + if (!file.startsWith(prefix)) continue + const [name = '', ...rest] = file.slice(prefix.length).split('/') + entries.set(name, rest.length > 0 ? 'dir' : 'file') + } + // Sorted by name, like the host. + return [...entries] + .sort(([a], [b]) => (a < b ? -1 : 1)) + .map(([name, type]) => ({ name, type })) + }, + exec: async (command, options) => ({ + exitCode: 3, + stdout: `ran ${command} in ${options?.cwd}`, + stderr: 'warn', + }), + ...extra, + } + return { backend, files } +} + +/** A backend, its path functions, and how to make and remove a folder on it. */ +interface Fixture { + name: string + backend: WorkspaceBackend + paths: PlatformPath + folder: (prefix: string) => Promise + remove: (folder: string) => Promise +} + +let memoryFolders = 0 +const fixtures: Array = [ + { + name: 'host', + backend: hostBackend, + paths: nodePath, + folder: (prefix) => mkdtemp(nodePath.join(tmpdir(), prefix)), + remove: (folder) => rm(folder, { recursive: true, force: true }), + }, + { + name: 'memory', + backend: memoryBackend().backend, + paths: posix, + folder: async (prefix) => { + memoryFolders += 1 + return `/${prefix}${memoryFolders}` + }, + remove: async () => undefined, + }, +] + +const decoder = new TextDecoder() + +/** The tools by name, for `options` and the user's `access`. */ +function toolsOf(options: WorkspaceToolsOptions, access?: OutsideAccess) { + return new Map( + createWorkspaceTools(options, { access }).tools.map( + (tool): [string, AnyTool] => [tool.name, tool], + ), + ) +} + +const runIn = ( + set: Map, + name: string, + args: unknown, +): Promise => { + const tool = set.get(name) + if (!tool?.execute) throw new Error(`No tool ${name}`) + return Promise.resolve(tool.execute(args)) +} + +/** A session with `plugins`, whose model is `adapter`. */ +async function openSession( + adapter: AnyTextAdapter, + plugins: Array, +) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/workspace-tools', + adapter, + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +/** The names of the tools that the model got in `call`. */ +const toolNames = (call: { tools: Array<{ name: string }> }) => + call.tools.map((tool) => tool.name) + +/** Is the process `pid` still running? */ +function isRunning(pid: number) { + try { + process.kill(pid, 0) + return true + } catch { + return false + } +} + +describe.each(fixtures)('workspace tools on the $name backend', (fixture) => { + const { backend, paths } = fixture + let root: string + let tools: Map + + const run = (name: string, args: unknown) => runIn(tools, name, args) + const write = (path: string, text: string) => + backend.writeFile(paths.join(root, path), text) + const read = async (path: string) => + decoder.decode(await backend.readFile(paths.join(root, path))) + + beforeAll(async () => { + root = await fixture.folder('harness-workspace-') + await write('src/a.ts', 'one\ntwo\nthree\nfour\n') + await write('src/deep/b.ts', 'const two = 2\n') + await write('notes.md', 'two words\n') + await write('node_modules/pkg/index.js', 'two\n') + await write('big.txt', `two\n${'x'.repeat(1_100_000)}`) + tools = toolsOf({ root, backend }) + }) + + afterAll(async () => { + await fixture.remove(root) + }) + + describe('read_file', () => { + it('reads a window of numbered lines', async () => { + expect( + await run('read_file', { path: 'src/a.ts', offset: 2, limit: 2 }), + ).toBe('2\ttwo\n3\tthree') + expect( + await run('read_file', { path: 'src/a.ts', offset: -4, limit: 1 }), + ).toBe('1\tone') + }) + + it('refuses paths outside the workspace and missing arguments', async () => { + await expect( + run('read_file', { path: '../outside.txt' }), + ).rejects.toThrow('outside the workspace') + await expect( + run('read_file', { + path: paths.join(paths.dirname(root), 'elsewhere.txt'), + }), + ).rejects.toThrow('outside the workspace') + await expect(run('read_file', {})).rejects.toThrow( + 'Argument "path" must be a string.', + ) + await expect(run('read_file', null)).rejects.toThrow('must be a string') + }) + + it('cuts very long output', async () => { + await write('long.txt', 'y'.repeat(25_000)) + const text = await run('read_file', { path: 'long.txt' }) + expect(text).toContain('more characters]') + expect(text.length).toBeLessThan(21_000) + }) + }) + + describe('write_file and edit_file', () => { + it('creates folders and files', async () => { + expect( + await run('write_file', { path: 'new/dir/c.txt', content: 'hi' }), + ).toBe('Wrote new/dir/c.txt.') + expect(await read('new/dir/c.txt')).toBe('hi') + await expect(run('write_file', { path: 'x.txt' })).rejects.toThrow( + 'Argument "content" must be a string.', + ) + }) + + it('edits one match, refuses ambiguous or missing text, and replaces all when asked', async () => { + await write('e.txt', 'a b a') + await expect( + run('edit_file', { path: 'e.txt', old: 'zzz', new: 'q' }), + ).rejects.toThrow('not in the file') + await expect( + run('edit_file', { path: 'e.txt', old: 'a', new: 'q' }), + ).rejects.toThrow('appears 2 times') + expect( + await run('edit_file', { + path: 'e.txt', + old: 'a', + new: 'q', + replaceAll: true, + }), + ).toBe('Edited e.txt (2 changes).') + expect( + await run('edit_file', { path: 'e.txt', old: 'b', new: 'c' }), + ).toBe('Edited e.txt (1 change).') + expect(await read('e.txt')).toBe('q c q') + }) + + it('runs edits to one file one at a time, so none is lost', async () => { + await write('race.txt', 'a b') + await Promise.all([ + run('edit_file', { path: 'race.txt', old: 'a', new: 'x' }), + run('edit_file', { path: 'race.txt', old: 'b', new: 'y' }), + ]) + expect(await read('race.txt')).toBe('x y') + }) + }) + + describe('patch', () => { + it('adds a file in the workspace', async () => { + const patch = [ + '*** Begin Patch', + '*** Add File: patched/p.txt', + '+patched', + '*** End Patch', + ].join('\n') + expect(await run('patch', { patch })).toBe( + 'Applied the patch:\nadded patched/p.txt', + ) + expect(await read('patched/p.txt')).toBe('patched\n') + }) + }) + + describe('list_files and grep', () => { + it('lists files, skips dependency folders, and filters by glob', async () => { + const all = await run('list_files', {}) + expect(all).toContain('src/deep/b.ts') + expect(all).not.toContain('node_modules') + expect(await run('list_files', { pattern: 'src/**/*.ts' })).toBe( + 'src/a.ts\nsrc/deep/b.ts', + ) + expect(await run('list_files', { pattern: '*.none' })).toBe('No files.') + }) + + it('finds lines, skips big files, and filters by glob', async () => { + const hits = await run('grep', { pattern: 'two' }) + expect(hits).toContain('src/a.ts:2: two') + expect(hits).toContain('notes.md:1: two words') + expect(hits).not.toContain('big.txt') + expect(hits).not.toContain('node_modules') + expect(await run('grep', { pattern: 'two', glob: '**/*.md' })).toBe( + 'notes.md:1: two words', + ) + expect(await run('grep', { pattern: 'nothing-like-this' })).toBe( + 'No matches.', + ) + }) + }) + + describe('outside the workspace', () => { + let outside: string + beforeAll(async () => { + outside = await fixture.folder('harness-outside-') + await backend.writeFile(paths.join(outside, 'docs', 'a.md'), 'alpha') + await backend.writeFile(paths.join(outside, 'docs', 'b.md'), 'beta') + await backend.writeFile(paths.join(outside, 'top.md'), 'top') + }) + afterAll(async () => { + await fixture.remove(outside) + }) + + /** Tools that ask about outside paths. The user answers from `answers`. */ + function asking(answers: Array, mode = 'default') { + const questions: Array = [] + const set = toolsOf( + { root, backend, outside: 'ask' }, + { + ask: async (message) => { + questions.push(message) + return answers.shift() + }, + mode: () => mode, + }, + ) + return { set, questions } + } + + it('refuses outside paths without the outside option', async () => { + await expect( + run('read_file', { path: paths.join(outside, 'top.md') }), + ).rejects.toThrow('outside the workspace') + }) + + it('asks once for a folder, then allows it for the session', async () => { + const { set, questions } = asking([true]) + const docs = paths.join(outside, 'docs') + expect( + await runIn(set, 'read_file', { path: paths.join(docs, 'a.md') }), + ).toBe('1\talpha') + expect( + await runIn(set, 'read_file', { path: paths.join(docs, 'b.md') }), + ).toBe('1\tbeta') + expect(questions).toHaveLength(1) + // The question names the real folder: the temp folder can be a link. + const real = (await backend.realpath?.(docs)) ?? docs + expect(questions[0]).toContain(`Allow ${real} `) + }) + + it('asks again for another folder, and refuses after a no', async () => { + const { set, questions } = asking([true, 'n']) + await runIn(set, 'read_file', { + path: paths.join(outside, 'docs', 'a.md'), + }) + await expect( + runIn(set, 'read_file', { path: paths.join(outside, 'top.md') }), + ).rejects.toThrow('The user did not allow') + expect(questions).toHaveLength(2) + }) + + it('asks one question for calls at the same time', async () => { + const { set, questions } = asking([true]) + await Promise.all([ + runIn(set, 'read_file', { path: paths.join(outside, 'docs', 'a.md') }), + runIn(set, 'read_file', { path: paths.join(outside, 'docs', 'b.md') }), + ]) + expect(questions).toHaveLength(1) + }) + + it('lists and searches an allowed outside folder with full paths', async () => { + const { set } = asking([true]) + const folder = paths.join(outside, 'docs') + const shown = folder.split(paths.sep).join('/') + expect(await runIn(set, 'list_files', { path: folder })).toBe( + `${shown}/a.md\n${shown}/b.md`, + ) + expect(await runIn(set, 'grep', { pattern: 'beta', path: folder })).toBe( + `${shown}/b.md:1: beta`, + ) + }) + + it('does not ask in bypass mode', async () => { + const { set, questions } = asking([], 'bypass') + expect( + await runIn(set, 'read_file', { path: paths.join(outside, 'top.md') }), + ).toBe('1\ttop') + expect(questions).toHaveLength(0) + }) + }) +}) + +describe('a sandbox backend with POSIX paths', () => { + it('keeps the root /workspace, also on a Windows host', async () => { + const { backend, files } = memoryBackend() + const set = toolsOf({ root: '/workspace', backend }) + expect( + await runIn(set, 'write_file', { path: 'src/a.ts', content: 'one' }), + ).toBe('Wrote src/a.ts.') + expect([...files.keys()]).toEqual(['/workspace/src/a.ts']) + expect(await runIn(set, 'list_files', {})).toBe('src/a.ts') + expect(await runIn(set, 'bash', { command: 'make' })).toBe( + 'exit code: 3\nran make in /workspace\nstderr:\nwarn', + ) + await expect( + runIn(set, 'read_file', { path: '/etc/passwd' }), + ).rejects.toThrow('outside the workspace') + }) +}) + +describe('the tools of workspaceTools', () => { + const names = (options: Partial = {}) => + createWorkspaceTools({ + root: '/workspace', + backend: memoryBackend().backend, + ...options, + }).tools.map((tool) => tool.name) + const search = { search: async () => [] } + + it('has the file, shell, and webfetch tools by default', () => { + expect(names()).toEqual([ + 'read_file', + 'write_file', + 'edit_file', + 'patch', + 'list_files', + 'grep', + 'bash', + 'webfetch', + ]) + }) + + it('adds websearch with a search provider, and no web tools for web: false', () => { + expect(names({ web: { search } }).slice(-2)).toEqual([ + 'webfetch', + 'websearch', + ]) + expect(names({ web: false })).not.toContain('webfetch') + }) +}) + +describe('editStyle', () => { + const EDIT_TOOLS = ['write_file', 'edit_file', 'patch'] + + it.each([ + { editStyle: undefined, model: 'gpt-5', tools: ['write_file', 'patch'] }, + { editStyle: undefined, model: 'o3', tools: ['write_file', 'patch'] }, + { + editStyle: undefined, + model: 'openai/gpt-4.1', + tools: ['write_file', 'patch'], + }, + { + editStyle: undefined, + model: 'claude-sonnet-4-5', + tools: ['write_file', 'edit_file'], + }, + { + editStyle: 'edit' as const, + model: 'gpt-5', + tools: ['write_file', 'edit_file'], + }, + { + editStyle: 'patch' as const, + model: 'claude-sonnet-4-5', + tools: ['write_file', 'patch'], + }, + ])( + 'gives $model the edit tools $tools with editStyle $editStyle', + async ({ editStyle, model, tools }) => { + const { adapter, calls } = mockAdapter([() => text('ok')]) + const { host, session } = await openSession({ ...adapter, model }, [ + workspaceTools({ + root: '/workspace', + backend: memoryBackend().backend, + editStyle, + }), + ]) + await session.prompt('hi') + expect( + toolNames(calls[0]).filter((name) => EDIT_TOOLS.includes(name)), + ).toEqual(tools) + await host.close() + }, + ) +}) + +describe('permission resources', () => { + const { resources } = createWorkspaceTools({ + root: '/workspace', + backend: memoryBackend().backend, + }) + + it('gives each call its paths, from the root', () => { + expect(resources.read_file?.paths?.({ path: 'src/a.ts' })).toEqual([ + '/workspace/src/a.ts', + ]) + expect(resources.edit_file?.paths?.({ path: '/etc/hosts' })).toEqual([ + '/etc/hosts', + ]) + expect(resources.write_file?.paths?.({ path: '../x.txt' })).toEqual([ + '/x.txt', + ]) + expect(resources.list_files?.paths?.({})).toEqual(['/workspace']) + expect(resources.grep?.paths?.({ pattern: 'a', path: 'src' })).toEqual([ + '/workspace/src', + ]) + }) + + it('gives every path of a patch, also the target of a move', () => { + const patch = [ + '*** Begin Patch', + '*** Add File: a.ts', + '+a', + '*** Update File: old.ts', + '*** Move to: lib/new.ts', + '@@', + '-x', + '+y', + '*** End Patch', + ].join('\n') + expect(resources.patch?.paths?.({ patch })).toEqual([ + '/workspace/a.ts', + '/workspace/old.ts', + '/workspace/lib/new.ts', + ]) + }) + + it('gives the command parts of bash', () => { + expect(resources.bash?.commands?.({ command: 'ls && rm x' })).toEqual([ + 'ls', + 'rm x', + ]) + }) + + it('throws, so the call is refused, when the input has no path or patch', () => { + expect(() => resources.read_file?.paths?.({})).toThrow( + 'Argument "path" must be a string.', + ) + expect(() => resources.patch?.paths?.({ patch: 'not a patch' })).toThrow( + 'Patch line 1', + ) + }) + + it('reach permissions(), so an allow rule for a path skips the question', async () => { + const { backend, files } = memoryBackend() + const patch = (path: string) => + ['*** Begin Patch', `*** Add File: ${path}`, '+a', '*** End Patch'].join( + '\n', + ) + const { adapter } = mockAdapter([ + () => toolCall('patch', { patch: patch('src/a.ts') }, 'c1'), + () => text('done'), + ]) + const { host, session } = await openSession(adapter, [ + permissions({ + root: '/workspace', + rules: [{ tool: 'patch', resource: 'src/**', decision: 'allow' }], + }), + workspaceTools({ root: '/workspace', backend, editStyle: 'patch' }), + ]) + await session.prompt('work') + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect([...files.keys()]).toEqual(['/workspace/src/a.ts']) + await host.close() + }) +}) + +describe('permission rules, links, and grep', () => { + const SECRET = 'KEY=hunter2' + const SKIPPED_TWO = '[Skipped 2 files that the permission rules protect.]' + + /** + * A memory backend with `files` (from `/workspace`) and `links` (a link + * path to its target, both from `/workspace`). A path through a link reads + * as its target, and readdir marks the link, like the host. + */ + function linkedBackend( + files: Record, + links: Record = {}, + ) { + const base = memoryBackend() + for (const [path, text] of Object.entries(files)) { + base.files.set(`/workspace/${path}`, new TextEncoder().encode(text)) + } + const real = (path: string) => { + for (const [link, target] of Object.entries(links)) { + const from = `/workspace/${link}` + if (path === from || path.startsWith(`${from}/`)) { + return `/workspace/${target}${path.slice(from.length)}` + } + } + return path + } + const backend: WorkspaceBackend = { + ...base.backend, + realpath: async (path) => real(path), + readFile: (path) => base.backend.readFile(real(path)), + stat: (path) => base.backend.stat(real(path)), + readdir: async (path) => { + const entries = await base.backend.readdir(real(path)) + const linked = Object.keys(links) + .filter((link) => posix.dirname(`/workspace/${link}`) === path) + .map((link) => ({ + name: posix.basename(link), + type: 'link' as const, + })) + return [...entries, ...linked] + }, + } + return backend + } + + /** A session where the model calls `tools` in order, in `mode`. */ + async function callTools( + backend: WorkspaceBackend, + tools: Array<[string, Record]>, + options: { rules?: Array; mode?: string } = {}, + ) { + const { adapter, calls } = mockAdapter([ + ...tools.map( + ([name, args], index) => + () => + toolCall(name, args, `c${index}`), + ), + () => text('done'), + ]) + const { host, session } = await openSession(adapter, [ + permissions({ root: '/workspace', rules: options.rules }), + workspaceTools({ root: '/workspace', backend, web: false }), + ]) + if (options.mode) await session.command('mode', options.mode) + /** Everything the model saw in its last call. */ + const seen = () => JSON.stringify(calls.at(-1)?.messages) + /** The result of the tool call with id `id`, as the model saw it. */ + const result = (id: string) => + JSON.stringify( + calls + .at(-1) + ?.messages.find( + (message: { toolCallId?: string }) => message.toolCallId === id, + )?.content, + ) + return { host, session, seen, result } + } + + it('asks before read_file reads .env through a link with another name', async () => { + const backend = linkedBackend({ '.env': SECRET }, { notes: '.env' }) + const { host, session, seen } = await callTools(backend, [ + ['read_file', { path: 'notes' }], + ]) + const turn = session.prompt('read the notes') + 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('.env') + await session.answer(question.questionId, 'n') + await turn + expect(seen()).not.toContain('hunter2') + await host.close() + }) + + it('applies a deny rule to the real path of a linked folder', async () => { + const backend = linkedBackend( + { 'secrets/key.txt': SECRET }, + { docs: 'secrets' }, + ) + const { host, session, seen, result } = await callTools( + backend, + [['read_file', { path: 'docs/key.txt' }]], + { + rules: [ + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ], + }, + ) + await session.prompt('read the key') + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect(result('c0')).toContain('The permission rules do not allow') + expect(seen()).not.toContain('hunter2') + await host.close() + }) + + it('refuses a linked .env in plan mode without a question', async () => { + const backend = linkedBackend({ '.env': SECRET }, { notes: '.env' }) + const { host, session, seen } = await callTools( + backend, + [['read_file', { path: 'notes' }]], + { mode: 'plan' }, + ) + await session.prompt('read the notes') + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect(seen()).not.toContain('hunter2') + await host.close() + }) + + it.each(['default', 'plan'])( + 'grep in %s mode skips the files a read_file rule protects, whatever the pattern', + async (mode) => { + const backend = linkedBackend({ + '.env': SECRET, + 'secrets/key.txt': 'KEY=2', + 'src/a.ts': 'KEY=1', + }) + const { host, session, seen, result } = await callTools( + backend, + [ + ['grep', { pattern: 'KEY' }], + // Matches no protected file. The note must not tell that. + ['grep', { pattern: 'KEY=1' }], + ], + { + rules: [ + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ], + mode, + }, + ) + await session.prompt('search') + expect(session.snapshot().pendingQuestions).toHaveLength(0) + for (const id of ['c0', 'c1']) { + expect(result(id)).toBe(`"src/a.ts:1: KEY=1\\n${SKIPPED_TWO}"`) + } + expect(seen()).not.toContain('hunter2') + expect(seen()).not.toContain('KEY=2') + await host.close() + }, + ) + + it('grep in a linked folder checks the real paths of its files', async () => { + const backend = linkedBackend( + { 'secrets/key.txt': SECRET }, + { docs: 'secrets' }, + ) + const { host, session, seen, result } = await callTools( + backend, + [['grep', { pattern: 'KEY', path: 'docs' }]], + { + rules: [ + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ], + }, + ) + await session.prompt('search') + expect(result('c0')).toBe( + '"No matches.\\n[Skipped 1 file that the permission rules protect.]"', + ) + expect(seen()).not.toContain('hunter2') + await host.close() + }) +}) + +describe('webfetch approval', () => { + /** The page that every request gets, so no request leaves the test. */ + const fetchPage: typeof fetch = async () => + new Response('fetched body', { + headers: { 'content-type': 'text/plain' }, + }) + + /** A session where the model calls webfetch once. `rules` go to permissions(). */ + async function fetchOnce(rules: Array = []) { + const { adapter, calls } = mockAdapter([ + () => toolCall('webfetch', { url: 'https://example.com/' }, 'c1'), + () => text('done'), + ]) + const { host, session } = await openSession(adapter, [ + permissions({ root: '/workspace', rules }), + workspaceTools({ + root: '/workspace', + backend: memoryBackend().backend, + // So the URL check does no DNS lookup. + web: { fetch: fetchPage, allowPrivateHosts: true }, + }), + ]) + return { host, session, calls } + } + + it('asks before webfetch by default', async () => { + const { host, session } = await fetchOnce() + const turn = session.prompt('read the page') + 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('Allow webfetch') + await session.answer(question.questionId, { answer: 'reject' }) + await turn + await host.close() + }) + + it('runs webfetch without a question when a user rule allows it', async () => { + const { host, session, calls } = await fetchOnce([ + { tool: 'webfetch', decision: 'allow' }, + ]) + await session.prompt('read the page') + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect(JSON.stringify(calls[1].messages)).toContain('fetched body') + await host.close() + }) +}) + +describe('background bash in a session', () => { + it('notes the end of a job, and starts a turn for it', async () => { + const { backend } = memoryBackend({ + spawn: () => ({ + wait: async () => ({ exitCode: 0 }), + kill: () => undefined, + output: () => 'built\n', + }), + }) + const { adapter, calls } = mockAdapter([ + () => toolCall('bash', { command: 'make', background: true }, 'c1'), + () => text('started'), + () => text('saw the note'), + ]) + const { host, session } = await openSession(adapter, [ + workspaceTools({ root: '/workspace', backend }), + ]) + await session.prompt('build') + await vi.waitFor(() => expect(calls).toHaveLength(3)) + expect(JSON.stringify(calls[2].messages)).toContain( + 'Background job bash-1 ended.\\nexit code: 0\\nbuilt', + ) + await host.close() + }) + + it('kills a job that still runs when the session closes', async () => { + let killed = false + const { backend } = memoryBackend({ + spawn: () => { + let end: (exitCode: number) => void = () => undefined + const exited = new Promise<{ exitCode: number }>((done) => { + end = (exitCode) => done({ exitCode }) + }) + return { + wait: () => exited, + kill: () => { + killed = true + end(143) + }, + output: () => '', + } + }, + }) + const { adapter } = mockAdapter([ + () => toolCall('bash', { command: 'serve', background: true }, 'c1'), + () => text('started'), + ]) + const { host, session } = await openSession(adapter, [ + workspaceTools({ root: '/workspace', backend }), + ]) + await session.prompt('serve') + expect(killed).toBe(false) + await host.close() + expect(killed).toBe(true) + }) +}) + +describe('links on the host', () => { + let dir = '' + let root = '' + let outside = '' + // A link to a file can need admin rights on Windows. The tests that need + // one skip without it. + let hasFileLinks = false + let tools: Map + const run = (name: string, args: unknown) => runIn(tools, name, args) + const put = (path: string, text: string) => + hostBackend.writeFile(nodePath.join(dir, path), text) + + beforeAll(async () => { + dir = await mkdtemp(nodePath.join(tmpdir(), 'harness-links-')) + root = nodePath.join(dir, 'root') + outside = nodePath.join(dir, 'outside') + await put('root/src/a.ts', 'inside\n') + await put('outside/secret.txt', 'secret\n') + // A junction on Windows needs no admin rights. Elsewhere it is a link. + await symlink(outside, nodePath.join(root, 'linked'), 'junction') + await symlink( + nodePath.join(root, 'src'), + nodePath.join(root, 'inner'), + 'junction', + ) + hasFileLinks = await symlink( + nodePath.join(outside, 'missing.txt'), + nodePath.join(root, 'dangling.txt'), + 'file', + ).then( + () => true, + () => false, + ) + tools = toolsOf({ root, backend: hostBackend }) + }) + afterAll(async () => { + await rm(dir, { recursive: true, force: true }) + }) + + it('marks links in readdir', async () => { + const entries = await hostBackend.readdir(root) + const sorted = entries.sort((a, b) => (a.name < b.name ? -1 : 1)) + expect(sorted.filter((entry) => entry.name !== 'dangling.txt')).toEqual([ + { name: 'inner', type: 'link' }, + { name: 'linked', type: 'link' }, + { name: 'src', type: 'dir' }, + ]) + }) + + it('refuses a link inside the workspace that leads out of it', async () => { + await expect( + run('read_file', { path: 'linked/secret.txt' }), + ).rejects.toThrow('outside the workspace') + await expect(run('list_files', { path: 'linked' })).rejects.toThrow( + 'outside the workspace', + ) + await expect( + run('write_file', { path: 'linked/new.txt', content: 'x' }), + ).rejects.toThrow('outside the workspace') + expect(await hostBackend.stat(nodePath.join(outside, 'new.txt'))).toBe( + undefined, + ) + }) + + it('refuses a write through a link to a missing file', async (context) => { + if (!hasFileLinks) context.skip() + await expect( + run('write_file', { path: 'dangling.txt', content: 'x' }), + ).rejects.toThrow('ENOENT') + expect(await hostBackend.stat(nodePath.join(outside, 'missing.txt'))).toBe( + undefined, + ) + }) + + it('allows a link that stays inside the workspace', async () => { + expect(await run('read_file', { path: 'inner/a.ts' })).toBe('1\tinside') + }) + + it('applies a deny rule for the real path of a link with permissions()', async () => { + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', { path: 'inner/a.ts' }, 'c0'), + () => text('done'), + ]) + const { host, session } = await openSession(adapter, [ + permissions({ + root, + rules: [{ tool: 'read_file', resource: 'src/**', decision: 'deny' }], + }), + workspaceTools({ root, backend: hostBackend, web: false }), + ]) + await session.prompt('read through the link') + const seen = JSON.stringify(calls[1].messages) + expect(seen).toContain( + 'The permission rules do not allow read_file on src/a.ts, the real path of inner/a.ts.', + ) + expect(seen).not.toContain('inside') + await host.close() + }) + + it('asks about the real folder of a linked path with outside: ask', async () => { + const questions: Array = [] + const set = toolsOf( + { root, backend: hostBackend, outside: 'ask' }, + { + ask: async (message) => { + questions.push(message) + return 'n' + }, + mode: () => 'default', + }, + ) + await expect( + runIn(set, 'read_file', { path: 'linked/secret.txt' }), + ).rejects.toThrow('The user did not allow') + expect(questions[0]).toContain(`Allow ${await realpath(outside)} `) + }) +}) + +describe('bash', () => { + it('reports the exit code and stderr on the host', async () => { + // Its own folder: the shared system temp folder churns while the other + // tests run, and a process that starts there can take seconds. + const root = await mkdtemp(nodePath.join(tmpdir(), 'harness-bash-host-')) + try { + const set = toolsOf({ + root, + backend: hostBackend, + bashTimeoutMs: 20_000, + }) + const ok = await runIn(set, 'bash', { + command: `node -e "console.log('out'); console.error('warn')"`, + }) + expect(ok).toMatch(/^exit code: 0\nout/) + expect(ok).toContain('stderr:\nwarn') + const failed = await runIn(set, 'bash', { + command: `node -e "process.exit(3)"`, + }) + expect(failed).toMatch(/^exit code: 3/) + } finally { + await rm(root, { recursive: true, force: true }) + } + }) + + it('runs the command through the backend, in the workspace folder', async () => { + const set = toolsOf({ + root: '/harness-bash', + backend: memoryBackend().backend, + }) + expect(await runIn(set, 'bash', { command: 'make' })).toBe( + 'exit code: 3\nran make in /harness-bash\nstderr:\nwarn', + ) + }) + + it( + 'stops the commands that a command started when the timeout ends it', + { timeout: 60_000 }, + async () => { + const cwd = await mkdtemp(nodePath.join(tmpdir(), 'harness-exec-tree-')) + try { + // The shell starts node, then waits for it, so node is a child of + // the shell. A killed shell alone leaves node running. + const child = `node -e "require('fs').writeFileSync('child.pid', String(process.pid)); setTimeout(function () {}, 60000)"` + const result = await hostBackend.exec(`${child} && echo after`, { + cwd, + timeoutMs: 3000, + }) + expect(result.exitCode).toBe(124) + expect(result.stdout).not.toContain('after') + const pid = Number( + await readFile(nodePath.join(cwd, 'child.pid'), 'utf8'), + ) + await vi.waitFor(() => expect(isRunning(pid)).toBe(false), { + timeout: 5000, + }) + } finally { + await rm(cwd, { recursive: true, force: true, maxRetries: 5 }) + } + }, + ) +}) diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index 86b2a7a465..9e4184adf9 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,7 +28,11 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts', './src/first-party/index.ts'], + entry: [ + './src/index.ts', + './src/first-party/index.ts', + './src/first-party/coding/index.ts', + ], srcDir: './src', cjs: false, }), diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 8ca632dd08..4d36b98ac1 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2595,6 +2595,12 @@ importers: '@vitest/coverage-v8': specifier: 4.1.10 version: 4.1.10(vitest@4.1.11) + '@vscode/ripgrep': + specifier: ^1.18.0 + version: 1.18.0 + turndown: + specifier: ^7.2.4 + version: 7.2.4 zod: specifier: ^4.2.0 version: 4.6.5 @@ -6584,6 +6590,9 @@ packages: '@mistralai/mistralai@2.2.0': resolution: {integrity: sha512-JQUGIXjFWnw/J9LpTSf/ZXwVW3Sh8FBAcfTo5QvAHqkl4CfSiIwnjRJhMoAFcP6ncCe84YPU1ncDGX+p3OXnfg==} + '@mixmark-io/domino@2.2.0': + resolution: {integrity: sha512-Y28PR25bHXUg88kCV7nivXrP2Nj2RueZ3/l/jdx6J9f8J4nsEGcgX0Qe6lt7Pa+J79+kPiJU3LguR6O/6zrLOw==} + '@modelcontextprotocol/client@2.1.0': resolution: {integrity: sha512-mDVhoy5WjDb0U+4dQPzLcciC2erSex/GRVQqnoZdiuMoE3YiTZdiS7ezNcFwX4XEOWOaj7jZr0kLYnILnL8orA==} engines: {node: '>=20'} @@ -11750,6 +11759,69 @@ packages: '@volar/typescript@2.4.28': resolution: {integrity: sha512-Ja6yvWrbis2QtN4ClAKreeUZPVYMARDYZl9LMEv1iQ1QdepB6wn0jTRxA9MftYmYa4DQ4k/DaSZpFPUfxl8giw==} + '@vscode/ripgrep-darwin-arm64@1.18.0': + resolution: {integrity: sha512-r3ktHSvbFycQNF6sl7sNDPocpsI7J+mEzh1IaZFkY0spm3k2Z9t8hPAeOK7+p0l6p6/swkQC14XWX01low+94Q==} + cpu: [arm64] + os: [darwin] + + '@vscode/ripgrep-darwin-x64@1.18.0': + resolution: {integrity: sha512-25b4gWbL138dGuQU244ebCKKc0q05ULBMoFSz9oAEUHNeqK/lOJViDS7DRvbDazzAzSEdan391Znks/R5mkaTQ==} + cpu: [x64] + os: [darwin] + + '@vscode/ripgrep-linux-arm64@1.18.0': + resolution: {integrity: sha512-lQ/5zTG++U0E3IhVgS4EPTTn/U4okncaRMM5GOFfOYZywS4nuD31GhkHbNYlDk5CuDC68+hYJ0/eQeyCKJDA+g==} + cpu: [arm64] + os: [linux] + + '@vscode/ripgrep-linux-arm@1.18.0': + resolution: {integrity: sha512-GDAvufNDHu8zqLEmXstalQF0Wh6wQvdsBi/Vg3Yi3CK4a8XoFXqqXVEHEZ9xQz3t0NfoSEc9JbvK9DDS6FxyxQ==} + cpu: [arm] + os: [linux] + + '@vscode/ripgrep-linux-ia32@1.18.0': + resolution: {integrity: sha512-YWLkSUtFd4Jh5EepIhA9RJSfv3uMAVMo+2rBIGHPBnvgLrZciIs2cDKei1/p6Wc/aCzUoHyMAg2R6tw4ZCBKGg==} + cpu: [ia32] + os: [linux] + + '@vscode/ripgrep-linux-ppc64@1.18.0': + resolution: {integrity: sha512-quXVY8fwQ8O/lvU1yrSqSl3jlUzysRSb+AfUfCL/tRtphxsKlFvPAejryZ6vg4Bgvn8XL74xb4qMCDmWgYrT5w==} + cpu: [ppc64] + os: [linux] + + '@vscode/ripgrep-linux-riscv64@1.18.0': + resolution: {integrity: sha512-f5kBQBrWfQt8Q7OhSORuNDei5dkYagBj3y4jImSUXGMy8B/Ke7SltSRcUtjPv166FAFfHCAmWuZp3+cWnX2/Vw==} + cpu: [riscv64] + os: [linux] + + '@vscode/ripgrep-linux-s390x@1.18.0': + resolution: {integrity: sha512-rTOcJFGGcl2c07RUOWUo4U1ndnemKhY6A9hnMB18uk7jSgJc0d/QLBGWMWpumdtoJtpizn/wIv5mXIisJukusQ==} + cpu: [s390x] + os: [linux] + + '@vscode/ripgrep-linux-x64@1.18.0': + resolution: {integrity: sha512-mQ3bVrUpnD2vs7QT0vX90Lt0cnUq467uFtEktIdsJJmW296RoSULRGqWgzG1AKxyBpNDD6l4ZO4qKf6SgyC23Q==} + cpu: [x64] + os: [linux] + + '@vscode/ripgrep-win32-arm64@1.18.0': + resolution: {integrity: sha512-vfTIjq1OHnzUjxZcHVQAMbnggp8dpGf+0QKFOZHwWPqFwXxQC8eCWM+5NUdoJ6yrElCeMzoUTXoK/LdZaniB+Q==} + cpu: [arm64] + os: [win32] + + '@vscode/ripgrep-win32-ia32@1.18.0': + resolution: {integrity: sha512-//rfAE+BOw5AC2EMmepmiE36jUuevtQYNQqqlw1s3m9FlRxjxEut97RkRPHAu9BG4mSojatZx+kXZXNdyI9caQ==} + cpu: [ia32] + os: [win32] + + '@vscode/ripgrep-win32-x64@1.18.0': + resolution: {integrity: sha512-KNPvtElldqILHdnAetujPaowkNbpqJy3ssIGGN6F6Kve9Qi+nNLI2DN01O83JjCEVQbCzl8Ov3QZ9Eov3BR8Dg==} + cpu: [x64] + os: [win32] + + '@vscode/ripgrep@1.18.0': + resolution: {integrity: sha512-ns5lWe44tSfbTMbVUsyB+I1819PVSw4AdpgK0RNkzfWfwy6+3IUNSxwSrfTno1/oWaS/hERNz+XLWVyga2aJBQ==} + '@vue/compiler-core@3.5.25': resolution: {integrity: sha512-vay5/oQJdsNHmliWoZfHPoVZZRmnSWhug0BYT34njkYTPqClh3DNWLkZNJBVSjsNMrg0CCrBfoKkjZQPM/QVUw==} @@ -17868,6 +17940,10 @@ packages: tunnel-agent@0.6.0: resolution: {integrity: sha512-McnNiV1l8RYeY8tBgEpuodCC1mLUdbSN+CYBL7kJsJNInOP8UjDDEwdk6Mw60vdLLrr5NHKZhMAOSrR2NZuQ+w==} + turndown@7.2.4: + resolution: {integrity: sha512-I8yFsfRzmzK0WV1pNNOA4A7y4RDfFxPRxb3t+e3ui14qSGOxGtiSP6GjeX+Y6CHb7HYaFj7ECUD7VE5kQMZWGQ==} + engines: {node: '>=18', npm: '>=9'} + tw-animate-css@1.4.0: resolution: {integrity: sha512-7bziOlRqH0hJx80h/3mbicLW7o8qLsH5+RaLR2t+OHM3D0JlWGODQKQ4cxbK7WlvmUxpcj6Kgu6EKqjrGFe3QQ==} @@ -22742,6 +22818,8 @@ snapshots: - bufferutil - utf-8-validate + '@mixmark-io/domino@2.2.0': {} + '@modelcontextprotocol/client@2.1.0': dependencies: '@modelcontextprotocol/core': 2.1.0 @@ -28929,6 +29007,57 @@ snapshots: path-browserify: 1.0.1 vscode-uri: 3.1.0 + '@vscode/ripgrep-darwin-arm64@1.18.0': + optional: true + + '@vscode/ripgrep-darwin-x64@1.18.0': + optional: true + + '@vscode/ripgrep-linux-arm64@1.18.0': + optional: true + + '@vscode/ripgrep-linux-arm@1.18.0': + optional: true + + '@vscode/ripgrep-linux-ia32@1.18.0': + optional: true + + '@vscode/ripgrep-linux-ppc64@1.18.0': + optional: true + + '@vscode/ripgrep-linux-riscv64@1.18.0': + optional: true + + '@vscode/ripgrep-linux-s390x@1.18.0': + optional: true + + '@vscode/ripgrep-linux-x64@1.18.0': + optional: true + + '@vscode/ripgrep-win32-arm64@1.18.0': + optional: true + + '@vscode/ripgrep-win32-ia32@1.18.0': + optional: true + + '@vscode/ripgrep-win32-x64@1.18.0': + optional: true + + '@vscode/ripgrep@1.18.0': + optionalDependencies: + '@vscode/ripgrep-darwin-arm64': 1.18.0 + '@vscode/ripgrep-darwin-x64': 1.18.0 + '@vscode/ripgrep-linux-arm': 1.18.0 + '@vscode/ripgrep-linux-arm64': 1.18.0 + '@vscode/ripgrep-linux-ia32': 1.18.0 + '@vscode/ripgrep-linux-ppc64': 1.18.0 + '@vscode/ripgrep-linux-riscv64': 1.18.0 + '@vscode/ripgrep-linux-s390x': 1.18.0 + '@vscode/ripgrep-linux-x64': 1.18.0 + '@vscode/ripgrep-win32-arm64': 1.18.0 + '@vscode/ripgrep-win32-ia32': 1.18.0 + '@vscode/ripgrep-win32-x64': 1.18.0 + '@vue/compiler-core@3.5.25': dependencies: '@babel/parser': 7.29.9 @@ -37211,6 +37340,10 @@ snapshots: safe-buffer: 5.2.1 optional: true + turndown@7.2.4: + dependencies: + '@mixmark-io/domino': 2.2.0 + tw-animate-css@1.4.0: {} tweetnacl@0.14.5: {} diff --git a/testing/e2e/fixtures/tools-test/coding-tools.json b/testing/e2e/fixtures/tools-test/coding-tools.json new file mode 100644 index 0000000000..f323e1537f --- /dev/null +++ b/testing/e2e/fixtures/tools-test/coding-tools.json @@ -0,0 +1,178 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[coding-read] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { "name": "read_file", "arguments": "{\"path\":\"notes.txt\"}" } + ] + } + }, + { + "match": { + "userMessage": "[coding-read] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "notes.txt has two lines." + } + }, + { + "match": { + "userMessage": "[coding-edit] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "edit_file", + "arguments": "{\"path\":\"notes.txt\",\"old\":\"alpha\\nbeta\",\"new\":\"gamma\\nbeta\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-edit] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "Edited notes.txt." + } + }, + { + "match": { + "userMessage": "[coding-patch] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "patch", + "arguments": "{\"patch\":\"*** Begin Patch\\n*** Update File: notes.txt\\n-alpha\\n+delta\\n beta\\n*** Add File: added.txt\\n+added by patch\\n*** End Patch\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-patch] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "Patched notes.txt and added added.txt." + } + }, + { + "match": { + "userMessage": "[coding-grep] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [{ "name": "grep", "arguments": "{\"pattern\":\"beta\"}" }] + } + }, + { + "match": { + "userMessage": "[coding-grep] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "beta is in two files." + } + }, + { + "match": { + "userMessage": "[coding-bash] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "bash", + "arguments": "{\"command\":\"node -e \\\"require('fs').writeFileSync('bash.txt', 'AGENT=' + process.env.AGENT); console.log('done')\\\"\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-bash] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "The command ran." + } + }, + { + "match": { + "userMessage": "[coding-revert] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "edit_file", + "arguments": "{\"path\":\"notes.txt\",\"old\":\"alpha\",\"new\":\"gamma\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-revert] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "Edited notes.txt." + } + }, + { + "match": { + "userMessage": "[coding-revert-2] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "edit_file", + "arguments": "{\"path\":\"notes.txt\",\"old\":\"gamma\",\"new\":\"delta\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-revert-2] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "Edited notes.txt again." + } + }, + { + "match": { + "userMessage": "[coding-format-ask] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "write_file", + "arguments": "{\"path\":\"a.ts\",\"content\":\"let a=1\\n\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-format-ask] run test", + "sequenceIndex": 1 + }, + "response": { "content": "Wrote a.ts." } + } + ] +} diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index e604264202..58810f4e6f 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -74,6 +74,7 @@ import { Route as ApiByteplusSeedance1080pWireRouteImport } from './routes/api.b import { Route as ApiChatRouteImport } from './routes/api.chat' import { Route as ApiChatCompletionsIncompleteStreamRouteImport } from './routes/api.chat-completions-incomplete-stream' import { Route as ApiCloudflareBindingWireRouteImport } from './routes/api.cloudflare-binding-wire' +import { Route as ApiCodingToolsRouteImport } from './routes/api.coding-tools' import { Route as ApiCompactionUsageWireRouteImport } from './routes/api.compaction-usage-wire' import { Route as ApiCompactionWireRouteImport } from './routes/api.compaction-wire' import { Route as ApiContextOverflowRouteImport } from './routes/api.context-overflow' @@ -523,6 +524,11 @@ const ApiCloudflareBindingWireRoute = path: '/api/cloudflare-binding-wire', getParentRoute: () => rootRouteImport, } as any) +const ApiCodingToolsRoute = ApiCodingToolsRouteImport.update({ + id: '/api/coding-tools', + path: '/api/coding-tools', + getParentRoute: () => rootRouteImport, +} as any) const ApiCompactionUsageWireRoute = ApiCompactionUsageWireRouteImport.update({ id: '/api/compaction-usage-wire', path: '/api/compaction-usage-wire', @@ -1143,6 +1149,7 @@ export interface FileRoutesByFullPath { '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute + '/api/coding-tools': typeof ApiCodingToolsRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/context-overflow': typeof ApiContextOverflowRoute @@ -1316,6 +1323,7 @@ export interface FileRoutesByTo { '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute + '/api/coding-tools': typeof ApiCodingToolsRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/context-overflow': typeof ApiContextOverflowRoute @@ -1490,6 +1498,7 @@ export interface FileRoutesById { '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute + '/api/coding-tools': typeof ApiCodingToolsRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute '/api/context-overflow': typeof ApiContextOverflowRoute @@ -1665,6 +1674,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/chat-completions-incomplete-stream' | '/api/cloudflare-binding-wire' + | '/api/coding-tools' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -1838,6 +1848,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/chat-completions-incomplete-stream' | '/api/cloudflare-binding-wire' + | '/api/coding-tools' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -2011,6 +2022,7 @@ export interface FileRouteTypes { | '/api/chat' | '/api/chat-completions-incomplete-stream' | '/api/cloudflare-binding-wire' + | '/api/coding-tools' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -2185,6 +2197,7 @@ export interface RootRouteChildren { ApiChatRoute: typeof ApiChatRoute ApiChatCompletionsIncompleteStreamRoute: typeof ApiChatCompletionsIncompleteStreamRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute + ApiCodingToolsRoute: typeof ApiCodingToolsRoute ApiCompactionUsageWireRoute: typeof ApiCompactionUsageWireRoute ApiCompactionWireRoute: typeof ApiCompactionWireRoute ApiContextOverflowRoute: typeof ApiContextOverflowRoute @@ -2746,6 +2759,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiCloudflareBindingWireRouteImport parentRoute: typeof rootRouteImport } + '/api/coding-tools': { + id: '/api/coding-tools' + path: '/api/coding-tools' + fullPath: '/api/coding-tools' + preLoaderRoute: typeof ApiCodingToolsRouteImport + parentRoute: typeof rootRouteImport + } '/api/compaction-usage-wire': { id: '/api/compaction-usage-wire' path: '/api/compaction-usage-wire' @@ -3619,6 +3639,7 @@ const rootRouteChildren: RootRouteChildren = { ApiChatCompletionsIncompleteStreamRoute: ApiChatCompletionsIncompleteStreamRoute, ApiCloudflareBindingWireRoute: ApiCloudflareBindingWireRoute, + ApiCodingToolsRoute: ApiCodingToolsRoute, ApiCompactionUsageWireRoute: ApiCompactionUsageWireRoute, ApiCompactionWireRoute: ApiCompactionWireRoute, ApiContextOverflowRoute: ApiContextOverflowRoute, diff --git a/testing/e2e/src/routes/api.coding-tools.ts b/testing/e2e/src/routes/api.coding-tools.ts new file mode 100644 index 0000000000..f0ee9d60ac --- /dev/null +++ b/testing/e2e/src/routes/api.coding-tools.ts @@ -0,0 +1,143 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join, resolve, sep } from 'node:path' +import { createFileRoute } from '@tanstack/react-router' +import { EventType, toServerSentEventsResponse } from '@tanstack/ai' +import { + HARNESS_EVENTS, + createHarnessHost, + defineHarness, +} from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { + formatter, + snapshots, + workspaceTools, +} from '@tanstack/ai-harness/plugins/coding' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import { createTextAdapter } from '@/lib/providers' + +/** + * A `coding-*` scenario: one harness turn with `workspaceTools()` in `root`, + * as the chunks the client gets. `coding-patch` gets the `patch` tool, the + * other scenarios get `edit_file`. No `permissions()` is mounted, so no tool + * asks for approval. Only `coding-format-ask` mounts `formatter()` and + * `permissions()` in `acceptEdits` mode: the write runs, the formatter asks, + * and the route rejects the question. + */ +async function* codingTurn( + scenario: string, + root: string, + adapter: AnyTextAdapter, +): AsyncGenerator { + const formatAsk = scenario === 'coding-format-ask' + const host = createHarnessHost() + try { + const session = await host.open( + defineHarness({ + name: 'e2e/coding-tools', + adapter, + plugins: () => [ + workspaceTools({ + root, + editStyle: scenario === 'coding-patch' ? 'patch' : 'edit', + }), + ...(formatAsk ? [permissions({ root }), formatter({ root })] : []), + ], + }), + { threadId: scenario }, + ) + if (formatAsk) await session.command('mode', 'acceptEdits') + for await (const chunk of session + .prompt(`[${scenario}] run test`) + .stream()) { + yield chunk + const asks = + chunk.type === EventType.CUSTOM && + chunk.name === HARNESS_EVENTS.question + const [question] = asks ? session.snapshot().pendingQuestions : [] + if (question) await session.answer(question.questionId, 'reject') + } + } finally { + await host.close() + } +} + +/** + * The `coding-revert` scenario: two turns with `snapshots()`, each edits + * notes.txt. Between them, the user edits other.txt. Then a revert to the + * first prompt puts back the files of both turns. + */ +async function* revertTurns( + root: string, + adapter: AnyTextAdapter, +): AsyncGenerator { + const dataDir = await mkdtemp(join(tmpdir(), 'e2e-snapshots-')) + const host = createHarnessHost() + try { + const session = await host.open( + defineHarness({ + name: 'e2e/coding-tools-revert', + adapter, + plugins: () => [ + // The fixtures call `edit_file`, so the edit style must not depend + // on the model. + workspaceTools({ root, editStyle: 'edit' }), + snapshots({ root, dataDir }), + ], + }), + { threadId: 'coding-revert' }, + ) + yield* session.prompt('[coding-revert] run test').stream() + await writeFile(join(root, 'other.txt'), 'mine\n') + yield* session.prompt('[coding-revert-2] run test').stream() + const [prompt] = await session.transcript() + if (prompt?.id) await session.revert(prompt.id) + } finally { + await host.close() + await rm(dataDir, { recursive: true, force: true }) + } +} + +/** + * The coding tools of `@tanstack/ai-harness/plugins/coding` on real files. + * Body: `{ scenario, testId, root }`. `root` must be a folder in the temp + * folder, because the tools write files and run commands there. The answer + * is an SSE stream of the turn's chunks. + */ +export const Route = createFileRoute('/api/coding-tools')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json() + const field = (name: string) => + typeof body === 'object' && body !== null && name in body + ? Reflect.get(body, name) + : undefined + const scenario = field('scenario') + const testId = field('testId') + const rootField = field('root') + if (typeof scenario !== 'string' || !scenario.startsWith('coding-')) { + return new Response('unknown scenario', { status: 400 }) + } + const root = typeof rootField === 'string' ? resolve(rootField) : '' + if (!root.startsWith(resolve(tmpdir()) + sep)) { + return new Response('root must be in the temp folder', { + status: 400, + }) + } + const { adapter } = createTextAdapter( + 'openai', + 'gpt-5.5', + undefined, + typeof testId === 'string' ? testId : undefined, + ) + return toServerSentEventsResponse( + scenario === 'coding-revert' + ? revertTurns(root, adapter) + : codingTurn(scenario, root, adapter), + ) + }, + }, + }, +}) diff --git a/testing/e2e/tests/tools-test/coding-tools.spec.ts b/testing/e2e/tests/tools-test/coding-tools.spec.ts new file mode 100644 index 0000000000..c7016fa0d4 --- /dev/null +++ b/testing/e2e/tests/tools-test/coding-tools.spec.ts @@ -0,0 +1,193 @@ +import { mkdtemp, readFile, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { test, expect } from '../fixtures' +import type { APIRequestContext } from '@playwright/test' + +/** + * The coding tools of `workspaceTools()` from + * `@tanstack/ai-harness/plugins/coding`. A `coding-*` scenario of + * `/api/coding-tools` runs one harness turn in a temp folder that the test + * makes. aimock gives one tool call, and the tool works on real files. Each + * test checks the `TOOL_CALL_RESULT` that the client gets, and the files + * after the turn. No `permissions()` is mounted, so no tool waits for + * approval. Only `coding-format-ask` mounts `permissions()`: there the + * formatter asks, and the route rejects the question. + */ + +type Chunk = { + type: string + delta?: string + content?: string + name?: string + value?: { message?: string } +} + +/** The chunks of an SSE body. */ +const sse = (body: string): Array => + body + .split('\n') + .filter((line) => line.startsWith('data: ')) + .map((line) => JSON.parse(line.slice('data: '.length))) + +/** The workspace of the running test. */ +let root = '' +const file = (path: string) => join(root, path) +const read = (path: string) => readFile(file(path), 'utf8') + +/** Run `scenario` in `root`. Gives the chunks the client got. */ +async function post( + request: APIRequestContext, + scenario: string, + testId: string, +) { + const response = await request.post('/api/coding-tools', { + data: { scenario, testId, root }, + }) + const body = await response.text() + expect(response.ok(), body).toBe(true) + const chunks = sse(body) + expect( + chunks.filter((chunk) => chunk.type === 'RUN_ERROR'), + body, + ).toEqual([]) + return chunks +} + +/** Run `scenario` in `root`. Gives the tool results and the text the client got. */ +async function runScenario( + request: APIRequestContext, + scenario: string, + testId: string, +) { + const chunks = await post(request, scenario, testId) + return { + results: chunks + .filter((chunk) => chunk.type === 'TOOL_CALL_RESULT') + .map((chunk) => chunk.content), + text: chunks + .filter((chunk) => chunk.type === 'TEXT_MESSAGE_CONTENT') + .map((chunk) => chunk.delta) + .join(''), + } +} + +test.describe('Coding tools', () => { + test.beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'e2e-coding-')) + }) + + test.afterEach(async () => { + await rm(root, { recursive: true, force: true }) + }) + + test('read_file gives the lines with their numbers', async ({ + request, + testId, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + + expect(await runScenario(request, 'coding-read', testId)).toEqual({ + results: ['1\talpha\n2\tbeta'], + text: 'notes.txt has two lines.', + }) + expect(await read('notes.txt')).toBe('alpha\nbeta\n') + }) + + test('edit_file finds the old text when the file has spaces at the end of a line', async ({ + request, + testId, + }) => { + // The old text is `alpha\nbeta`. The exact text is not in the file. + await writeFile(file('notes.txt'), 'alpha \nbeta\n') + + expect(await runScenario(request, 'coding-edit', testId)).toEqual({ + results: ['Edited notes.txt (1 change).'], + text: 'Edited notes.txt.', + }) + // The match covers the trailing spaces, so they go too. + expect(await read('notes.txt')).toBe('gamma\nbeta\n') + }) + + test('patch updates one file and adds another', async ({ + request, + testId, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + + expect(await runScenario(request, 'coding-patch', testId)).toEqual({ + results: ['Applied the patch:\nupdated notes.txt\nadded added.txt'], + text: 'Patched notes.txt and added added.txt.', + }) + expect(await read('notes.txt')).toBe('delta\nbeta\n') + expect(await read('added.txt')).toBe('added by patch\n') + }) + + test('grep gives each match as file:line: text', async ({ + request, + testId, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + await writeFile(file('more.txt'), 'more beta\n') + + expect(await runScenario(request, 'coding-grep', testId)).toEqual({ + results: ['more.txt:1: more beta\nnotes.txt:2: beta'], + text: 'beta is in two files.', + }) + expect(await read('notes.txt')).toBe('alpha\nbeta\n') + }) + + test('bash runs the command in the workspace with AGENT=1', async ({ + request, + testId, + }) => { + // The command is `node -e "..."`: it writes bash.txt and prints `done`. + expect(await runScenario(request, 'coding-bash', testId)).toEqual({ + results: ['exit code: 0\ndone\n'], + text: 'The command ran.', + }) + expect(await read('bash.txt')).toBe('AGENT=1') + }) + + test('revert puts back the files of the turns and keeps the edit of the user', async ({ + request, + testId, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + await writeFile(file('other.txt'), 'original\n') + + // Two turns edit notes.txt. Between them, the user edits other.txt. + // Then the route reverts to the first prompt. + expect(await runScenario(request, 'coding-revert', testId)).toEqual({ + results: ['Edited notes.txt (1 change).', 'Edited notes.txt (1 change).'], + text: 'Edited notes.txt.Edited notes.txt again.', + }) + expect(await read('notes.txt')).toBe('alpha\nbeta\n') + expect(await read('other.txt')).toBe('mine\n') + }) + + test('in acceptEdits mode, the formatter asks before it runs, and a reject keeps the write', async ({ + request, + testId, + }) => { + // The config file makes prettier the formatter of a.ts. + await writeFile(file('.prettierrc'), '{}') + + const chunks = await post(request, 'coding-format-ask', testId) + + const questions = chunks + .filter((chunk) => chunk.name === 'harness.question') + .map((chunk) => chunk.value?.message) + expect(questions).toEqual([ + `Allow formatter:prettier on ${file('a.ts')}? It runs code from the project. Answer once, always, or reject.`, + ]) + expect( + chunks + .filter((chunk) => chunk.type === 'TOOL_CALL_RESULT') + .map((chunk) => chunk.content), + ).toEqual([ + 'Wrote a.ts.\nThe formatter prettier did not run: it is not allowed.', + ]) + expect(await read('a.ts')).toBe('let a=1\n') + }) +}) From 4d87928d53c43334f35a33d472d24d200bbd60c6 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 00:41:09 +0200 Subject: [PATCH 04/15] docs(ai-harness): end the terminal loop of the coding samples on exit --- docs/harness/coding-agent.md | 30 +++++++++++++++++------------- docs/harness/coding-tools.md | 9 ++++++--- 2 files changed, 23 insertions(+), 16 deletions(-) diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md index c5f88ff349..a0d33794c3 100644 --- a/docs/harness/coding-agent.md +++ b/docs/harness/coding-agent.md @@ -70,26 +70,30 @@ const terminal = createInterface({ output: process.stdout, }) -for (;;) { - const line = await terminal.question('\n> ') - // `/undo`, `/model fast`, `/mode plan`, and the other plugin commands. +// Type `exit` to stop. +let line = await terminal.question('\n> ') +while (line !== 'exit') { if (line.startsWith('/')) { + // `/undo`, `/model fast`, `/mode plan`, and the other plugin commands. const [name = '', ...args] = line.slice(1).split(' ') console.log(await session.command(name, args.join(' ') || undefined)) - continue - } - for await (const chunk of session.prompt(line).stream()) { - if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) - // A risky call waits for an answer: once, always, or reject. - if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { - const [question] = session.snapshot().pendingQuestions - if (question) { - const answer = await terminal.question(`\n${question.message} `) - await session.answer(question.questionId, answer) + } else { + for await (const chunk of session.prompt(line).stream()) { + if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) + // A risky call waits for an answer: once, always, or reject. + if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { + const [question] = session.snapshot().pendingQuestions + if (question) { + const answer = await terminal.question(`\n${question.message} `) + await session.answer(question.questionId, answer) + } } } } + line = await terminal.question('\n> ') } +terminal.close() +await host.close() ``` Both plugin entries run only in Node. The coding tools use the file system and the shell of the machine. A line that starts with `/` runs a command of a plugin. Any other line is a prompt. diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index 55431056e8..579029005a 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -77,9 +77,9 @@ const terminal = createInterface({ }) // Send each line as a prompt. Print the answer. Ask the user each -// question that a risky call waits on. -for (;;) { - const line = await terminal.question('\n> ') +// question that a risky call waits on. Type `exit` to stop. +let line = await terminal.question('\n> ') +while (line !== 'exit') { for await (const chunk of session.prompt(line).stream()) { if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { @@ -90,7 +90,10 @@ for (;;) { } } } + line = await terminal.question('\n> ') } +terminal.close() +await host.close() ``` `@tanstack/ai-harness/plugins/coding` runs only in Node. The tools use the file system and the shell of the machine. From 9bdd8600090271e9aa9ece6e229835f3c434d005 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:39:03 +0200 Subject: [PATCH 05/15] 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 06/15] 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 07/15] 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 1a1a7db3be226507d887dcf89bed8b9935e7ddfa Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:51:03 +0200 Subject: [PATCH 08/15] fix(ai-harness, ai-e2e): adopt object params from #1779 The coding tool tests and the coding-tools E2E route call session.answer with one object. --- packages/ai-harness/tests/coding-formatter.test.ts | 2 +- packages/ai-harness/tests/first-party.test.ts | 10 ++++++++-- packages/ai-harness/tests/workspace-tools.test.ts | 7 +++++-- testing/e2e/src/routes/api.coding-tools.ts | 6 +++++- 4 files changed, 19 insertions(+), 6 deletions(-) diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index 01e0760fa2..18f229337e 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -357,7 +357,7 @@ describe('formatter with permissions()', () => { ) 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/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index bc136b24d4..f386a6d93c 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -81,7 +81,10 @@ describe('workspace tools with permissions', () => { ) const [question] = session.snapshot().pendingQuestions expect(question?.message).toContain('write_file') - await session.answer(question!.questionId, { answer: 'once' }) + await session.answer({ + questionId: question!.questionId, + value: { answer: 'once' }, + }) await turn expect(await readFile(join(root, 'out.txt'), 'utf8')).toBe('written') @@ -125,7 +128,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() }) diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index 137ae6f20d..26bc868246 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -652,7 +652,7 @@ describe('permission rules, links, and grep', () => { const [question] = session.snapshot().pendingQuestions if (!question) throw new Error('No question.') expect(question.message).toContain('.env') - await session.answer(question.questionId, 'n') + await session.answer({ questionId: question.questionId, value: 'n' }) await turn expect(seen()).not.toContain('hunter2') await host.close() @@ -782,7 +782,10 @@ describe('webfetch approval', () => { const [question] = session.snapshot().pendingQuestions if (!question) throw new Error('No question.') expect(question.message).toContain('Allow webfetch') - await session.answer(question.questionId, { answer: 'reject' }) + await session.answer({ + questionId: question.questionId, + value: { answer: 'reject' }, + }) await turn await host.close() }) diff --git a/testing/e2e/src/routes/api.coding-tools.ts b/testing/e2e/src/routes/api.coding-tools.ts index f0ee9d60ac..7d147b00d3 100644 --- a/testing/e2e/src/routes/api.coding-tools.ts +++ b/testing/e2e/src/routes/api.coding-tools.ts @@ -56,7 +56,11 @@ async function* codingTurn( chunk.type === EventType.CUSTOM && chunk.name === HARNESS_EVENTS.question const [question] = asks ? session.snapshot().pendingQuestions : [] - if (question) await session.answer(question.questionId, 'reject') + if (question) + await session.answer({ + questionId: question.questionId, + value: 'reject', + }) } } finally { await host.close() From b33f052820a3b004159e3376be7a9c7aa9c48026 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:51:03 +0200 Subject: [PATCH 09/15] docs(ai-harness): print and answer with onEvent in the coding samples The coding agent and coding tools samples print the text with turn.onEvent and answer questions with session.onEvent, instead of a loop over prompt().stream(). They call session.answer with one object. --- docs/config.json | 6 ++++-- docs/harness/coding-agent.md | 26 +++++++++++++++----------- docs/harness/coding-tools.md | 28 ++++++++++++++++------------ 3 files changed, 35 insertions(+), 25 deletions(-) diff --git a/docs/config.json b/docs/config.json index beb99bae36..db49a35cd5 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1003,7 +1003,8 @@ { "label": "Give your agent coding tools", "to": "harness/coding-tools", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-14" }, { "label": "Extend the coding tools", @@ -1018,7 +1019,8 @@ { "label": "Build a coding agent", "to": "harness/coding-agent", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-14" }, { "label": "Switch between agent profiles", diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md index a0d33794c3..145c90492c 100644 --- a/docs/harness/coding-agent.md +++ b/docs/harness/coding-agent.md @@ -70,6 +70,15 @@ const terminal = createInterface({ output: process.stdout, }) +// A risky call waits for an answer: once, always, or reject. +const stopAsking = session.onEvent('CUSTOM', async ({ event }) => { + if (event.name !== HARNESS_EVENTS.question) return + const [question] = session.snapshot().pendingQuestions + if (!question) return + const answer = await terminal.question(`\n${question.message} `) + await session.answer({ questionId: question.questionId, value: answer }) +}) + // Type `exit` to stop. let line = await terminal.question('\n> ') while (line !== 'exit') { @@ -78,20 +87,15 @@ while (line !== 'exit') { const [name = '', ...args] = line.slice(1).split(' ') console.log(await session.command(name, args.join(' ') || undefined)) } else { - for await (const chunk of session.prompt(line).stream()) { - if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) - // A risky call waits for an answer: once, always, or reject. - if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { - const [question] = session.snapshot().pendingQuestions - if (question) { - const answer = await terminal.question(`\n${question.message} `) - await session.answer(question.questionId, answer) - } - } - } + const turn = session.prompt(line) + turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { + process.stdout.write(event.delta) + }) + await turn } line = await terminal.question('\n> ') } +stopAsking() terminal.close() await host.close() ``` diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index 579029005a..e550e14f3c 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -76,22 +76,26 @@ const terminal = createInterface({ output: process.stdout, }) -// Send each line as a prompt. Print the answer. Ask the user each -// question that a risky call waits on. Type `exit` to stop. +// Ask the user each question that a risky call waits on. +const stopAsking = session.onEvent('CUSTOM', async ({ event }) => { + if (event.name !== HARNESS_EVENTS.question) return + const [question] = session.snapshot().pendingQuestions + if (!question) return + const answer = await terminal.question(`\n${question.message} `) + await session.answer({ questionId: question.questionId, value: answer }) +}) + +// Send each line as a prompt. Print the answer. Type `exit` to stop. let line = await terminal.question('\n> ') while (line !== 'exit') { - for await (const chunk of session.prompt(line).stream()) { - if (chunk.type === 'TEXT_MESSAGE_CONTENT') process.stdout.write(chunk.delta) - if (chunk.type === 'CUSTOM' && chunk.name === HARNESS_EVENTS.question) { - const [question] = session.snapshot().pendingQuestions - if (question) { - const answer = await terminal.question(`\n${question.message} `) - await session.answer(question.questionId, answer) - } - } - } + const turn = session.prompt(line) + turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { + process.stdout.write(event.delta) + }) + await turn line = await terminal.question('\n> ') } +stopAsking() terminal.close() 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 10/15] 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 332a8fcfa03f7eea92d2696044471af6ae1218b5 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 20:58:27 +0200 Subject: [PATCH 11/15] docs(ai-harness): ask one question at a time in the coding samples Tool calls can run in parallel, so two questions can be open at the same time. The terminal samples now ask each open question once, one at a time. A failed turn prints its error and the loop goes on. --- docs/harness/coding-agent.md | 24 ++++++++++++++++++------ docs/harness/coding-tools.md | 26 +++++++++++++++++++------- 2 files changed, 37 insertions(+), 13 deletions(-) diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md index 145c90492c..8a98aa9c00 100644 --- a/docs/harness/coding-agent.md +++ b/docs/harness/coding-agent.md @@ -71,12 +71,19 @@ const terminal = createInterface({ }) // A risky call waits for an answer: once, always, or reject. -const stopAsking = session.onEvent('CUSTOM', async ({ event }) => { +// Calls can run in parallel. Ask one question at a time, each once. +const asked = new Set() +let asking = Promise.resolve() +const stopAsking = session.onEvent('CUSTOM', ({ event }) => { if (event.name !== HARNESS_EVENTS.question) return - const [question] = session.snapshot().pendingQuestions - if (!question) return - const answer = await terminal.question(`\n${question.message} `) - await session.answer({ questionId: question.questionId, value: answer }) + asking = asking.then(async () => { + for (const question of session.snapshot().pendingQuestions) { + if (asked.has(question.questionId)) continue + asked.add(question.questionId) + const answer = await terminal.question(`\n${question.message} `) + await session.answer({ questionId: question.questionId, value: answer }) + } + }) }) // Type `exit` to stop. @@ -91,7 +98,12 @@ while (line !== 'exit') { turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { process.stdout.write(event.delta) }) - await turn + try { + await turn + } catch (error) { + // A failed turn rejects. Show the error and keep the loop. + console.error(error) + } } line = await terminal.question('\n> ') } diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index e550e14f3c..84d5a711e9 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -76,13 +76,20 @@ const terminal = createInterface({ output: process.stdout, }) -// Ask the user each question that a risky call waits on. -const stopAsking = session.onEvent('CUSTOM', async ({ event }) => { +// Ask the user each question that a risky call waits on. Calls can run +// in parallel, so ask one question at a time, each once. +const asked = new Set() +let asking = Promise.resolve() +const stopAsking = session.onEvent('CUSTOM', ({ event }) => { if (event.name !== HARNESS_EVENTS.question) return - const [question] = session.snapshot().pendingQuestions - if (!question) return - const answer = await terminal.question(`\n${question.message} `) - await session.answer({ questionId: question.questionId, value: answer }) + asking = asking.then(async () => { + for (const question of session.snapshot().pendingQuestions) { + if (asked.has(question.questionId)) continue + asked.add(question.questionId) + const answer = await terminal.question(`\n${question.message} `) + await session.answer({ questionId: question.questionId, value: answer }) + } + }) }) // Send each line as a prompt. Print the answer. Type `exit` to stop. @@ -92,7 +99,12 @@ while (line !== 'exit') { turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { process.stdout.write(event.delta) }) - await turn + try { + await turn + } catch (error) { + // A failed turn rejects. Show the error and keep the loop. + console.error(error) + } line = await terminal.question('\n> ') } stopAsking() From 7ae9f5c0ea790040b5cf92ddfdf57ef50f61a611 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Thu, 15 Oct 2026 13:35:44 +0200 Subject: [PATCH 12/15] feat(ai-harness, ai-sandbox): codingAgents and sandboxWorkspaceBackend in @tanstack/ai-sandbox/harness (#1771) * feat(ai-sandbox): codingAgents and sandboxWorkspaceBackend in @tanstack/ai-sandbox/harness codingAgents({ sandbox, agents, workspace? }) lets a harness hand tasks to Claude Code, Codex, Grok Build, or any ACP agent in a sandbox, with a saved session per agent and read-only plan mode. sandboxWorkspaceBackend runs the harness coding tools in a sandbox. Split out of #1555. * test(ai-sandbox): an agent without plan options does not run in plan mode * fix(ai-sandbox): adopt object params from #1779 The coding agents test calls session.setConfig with one object. --- .changeset/ai-sandbox-harness.md | 10 + docs/config.json | 13 +- docs/harness/coding-agent.md | 6 +- docs/harness/coding-agents.md | 129 +++++ docs/harness/coding-tools-backends.md | 67 ++- docs/harness/coding-tools.md | 4 +- docs/harness/permissions.md | 1 + docs/harness/snapshots.md | 20 + packages/ai-sandbox/package.json | 9 + packages/ai-sandbox/src/harness.ts | 218 +++++++++ packages/ai-sandbox/src/workspace-backend.ts | 154 ++++++ .../ai-sandbox/tests/coding-agents.test.ts | 383 +++++++++++++++ .../tests/workspace-backend.test.ts | 456 ++++++++++++++++++ packages/ai-sandbox/vite.config.ts | 7 +- pnpm-lock.yaml | 3 + .../e2e/fixtures/tools-test/coding-tools.json | 39 +- testing/e2e/src/routes/api.coding-tools.ts | 40 +- .../e2e/tests/tools-test/coding-tools.spec.ts | 14 + 18 files changed, 1558 insertions(+), 15 deletions(-) create mode 100644 .changeset/ai-sandbox-harness.md create mode 100644 docs/harness/coding-agents.md create mode 100644 packages/ai-sandbox/src/harness.ts create mode 100644 packages/ai-sandbox/src/workspace-backend.ts create mode 100644 packages/ai-sandbox/tests/coding-agents.test.ts create mode 100644 packages/ai-sandbox/tests/workspace-backend.test.ts diff --git a/.changeset/ai-sandbox-harness.md b/.changeset/ai-sandbox-harness.md new file mode 100644 index 0000000000..10f2443916 --- /dev/null +++ b/.changeset/ai-sandbox-harness.md @@ -0,0 +1,10 @@ +--- +'@tanstack/ai-sandbox': minor +--- + +Add `@tanstack/ai-sandbox/harness` for `@tanstack/ai-harness`: + +- `codingAgents({ sandbox, agents, workspace? })` is a harness plugin. The lead model gets one tool per coding agent (Claude Code, Codex, Grok Build, or any ACP agent). Each agent works in the sandbox, keeps its own session per thread (also after a restart), and starts read-only in the `plan` mode of `permissions()`. `/fresh [agent]` starts a new session. `workspace: 'shared'` (default) runs the agents of a thread one at a time in one sandbox. `'per-agent'` gives each agent its own sandbox. +- `sandboxWorkspaceBackend(sandbox)` runs the coding tools of `@tanstack/ai-harness/plugins/coding` (`workspaceTools()`, `snapshots()`, `formatter()`) in a sandbox. + +`@tanstack/ai-harness` is an optional peer dependency. diff --git a/docs/config.json b/docs/config.json index db49a35cd5..0baf38d22e 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1007,14 +1007,16 @@ "updatedAt": "2026-10-14" }, { - "label": "Extend the coding tools", + "label": "Sandbox and extend the coding tools", "to": "harness/coding-tools-backends", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" }, { "label": "Undo a turn", "to": "harness/snapshots", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" }, { "label": "Build a coding agent", @@ -1022,6 +1024,11 @@ "addedAt": "2026-10-13", "updatedAt": "2026-10-14" }, + { + "label": "Delegate to coding agents", + "to": "harness/coding-agents", + "addedAt": "2026-10-13" + }, { "label": "Switch between agent profiles", "to": "harness/agents", diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md index 8a98aa9c00..0effd1f088 100644 --- a/docs/harness/coding-agent.md +++ b/docs/harness/coding-agent.md @@ -184,7 +184,11 @@ export const noDeploys = definePlugin({ A trailing `*` matches every tool that starts with the text. The last matching rule wins. [Ask before risky tool calls](./permissions) explains the rules, the modes, and the saved answers. -The workspace tools run on your machine with your permissions. To run them on another machine, [write your own backend](./coding-tools-backends#write-your-own-backend). +The workspace tools run on your machine with your permissions. For code that you do not trust, [run the tools in a sandbox](./coding-tools-backends#run-the-tools-in-a-sandbox). + +## Hand work to Claude Code or Codex + +Your agent can also give tasks to coding agents you already use. [Delegate to coding agents](./coding-agents) shows how. ## What you have now diff --git a/docs/harness/coding-agents.md b/docs/harness/coding-agents.md new file mode 100644 index 0000000000..994c0c52ca --- /dev/null +++ b/docs/harness/coding-agents.md @@ -0,0 +1,129 @@ +--- +title: Delegate to coding agents +id: harness-coding-agents +order: 11 +description: "Let a harness hand coding work to Claude Code, Codex, Grok Build, or any ACP agent. Each one works in a sandbox and keeps its own session." +keywords: + - tanstack ai + - harness + - claude code + - codex + - grok build + - sandbox + - subagents +--- + +Your lead agent plans the work, but you want Claude Code or Codex to do the edits, because they are good at code and you already use them. `codingAgents` gives the lead model one tool per coding agent. Each agent works in a sandbox and keeps its own session between calls. + +## 1. Add the plugin + + + +react: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +vue: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +solid: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +svelte: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +preact: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +angular: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +octane: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex +vanilla: @tanstack/ai-harness @tanstack/ai-sandbox @tanstack/ai-sandbox-local-process @tanstack/ai-claude-code @tanstack/ai-codex + + + +```ts group=harness-coding-agents +import { defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { claudeCodeText } from '@tanstack/ai-claude-code' +import { codexText } from '@tanstack/ai-codex' +import { openaiText } from '@tanstack/ai-openai' +import { defineSandbox, defineWorkspace, localSource } from '@tanstack/ai-sandbox' +import { codingAgents } from '@tanstack/ai-sandbox/harness' +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' + +const repo = '/path/to/your/repo' + +export const lead = defineHarness({ + name: 'acme/lead', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + permissions(), + codingAgents({ + sandbox: defineSandbox({ + id: 'repo', + provider: localProcessSandbox({ dir: repo }), + workspace: defineWorkspace({ source: localSource(repo) }), + }), + agents: { + claude_code: { + adapter: claudeCodeText('claude-opus-4-8', { permissionMode: 'acceptEdits' }), + description: 'Larger changes, refactors, and reviews', + }, + codex: { + adapter: codexText('gpt-5.3-codex', { sandboxMode: 'workspace-write' }), + description: 'Quick fixes and tests', + }, + }, + }), + ], +}) +``` + +`localProcessSandbox({ dir: repo })` works in your repo folder itself, so every thread and every agent shares that folder. Leave out `dir` to give each sandbox its own copy of the repo. + +The lead model now has a `claude_code` tool and a `codex` tool. Each tool takes one `task`: the whole job, in words. The agent sees only that text and the files in its sandbox. + +## 2. Ask for work + +Open a session, and give the lead a task for both agents: + +```ts group=harness-coding-agents +import { createHarnessHost } from '@tanstack/ai-harness' +import { memoryPersistence } from '@tanstack/ai-persistence' + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(lead, { threadId: 'date-parser' }) + +const turn = await session.prompt( + 'Have claude_code add a test for the date parser, then have codex fix what fails.', +) +console.log(turn.text) +``` + +The lead calls `claude_code` with a task, and gets back its answer. Then it calls `codex`. Each agent works in the same sandbox, one after the other. + +## Sessions + +Each agent keeps its own session per harness thread. The next task for `claude_code` resumes the same Claude Code session, so it still knows the files it read. The session ids live in the plugin state, so they also survive a restart of the host. + +To start over, run `/fresh claude_code`, or `/fresh` for every agent. + +## Workspaces + +`workspace` chooses where the agents work: + +| Value | Where each agent works | At the same time | +| --- | --- | --- | +| `'shared'` (default) | One sandbox per harness thread | One agent at a time | +| `'per-agent'` | A sandbox per agent | Yes | + +Use `'shared'` when the agents build on each other's changes. Use `'per-agent'` when they work on separate copies and you merge the results. + +## Plan mode + +When the `permissions()` plugin is in [`plan` mode](./permissions#pick-a-mode), the agents start read-only: + +- Claude Code gets `permissionMode: 'plan'`. +- Codex gets `sandboxMode: 'read-only'`. +- Any other agent gets its `planModelOptions`, for example `{ permissionMode: 'default' }` for an ACP agent that asks before each edit. +- An agent with no `planModelOptions` does not run in `plan` mode. Its tool call gets an error. + +## Other agents and sandboxes + +- `adapter` takes any coding-agent adapter: `grokBuildText` from `@tanstack/ai-grok-build`, or `acpCompatibleText` from `@tanstack/ai-acp` for any ACP agent. +- `sandbox` takes any sandbox provider. [Sandbox providers](../sandbox/providers) lists them. +- `modelOptions` on an agent is added to every call, for example a fixed `permissionMode`. + +## What you have now + +- A lead model that hands tasks to Claude Code and Codex. +- One sandbox per thread, and one saved session per agent. diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md index da9e702484..86718fdda2 100644 --- a/docs/harness/coding-tools-backends.md +++ b/docs/harness/coding-tools-backends.md @@ -1,18 +1,76 @@ --- -title: Extend the coding tools +title: Sandbox and extend the coding tools id: harness-coding-tools-backends order: 7 -description: "Run the coding tools on your own backend, format each written file, and keep tool output small." +description: "Run the coding tools in a sandbox or on your own backend, format each written file, and keep tool output small." keywords: - tanstack ai - harness - coding agent - workspace tools + - sandbox - formatter - workspace hooks --- -Your agent must work on files that are not on this machine. Or you want code to run after each write, like a formatter. The [coding tools](./coding-tools) send every read, write, and command through a backend, and they run hooks after each read and write. This page shows how to change both. +Your agent must work on files that are not on this machine, or on code that you do not trust. Or you want code to run after each write, like a formatter. The [coding tools](./coding-tools) send every read, write, and command through a backend, and they run hooks after each read and write. This page shows how to change both. + +## Run the tools in a sandbox + +By default, the tools run on your machine, with the rights of your user. For code that you do not trust, run them in a sandbox. `sandboxWorkspaceBackend()` sends every read, write, and command to the sandbox. + +1. Install the sandbox packages: + + + +react: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +vue: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +solid: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +svelte: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +preact: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +angular: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +octane: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker +vanilla: @tanstack/ai-sandbox @tanstack/ai-sandbox-docker + + + +2. Start a sandbox, and give its backend to `workspaceTools()`: + +```ts group=harness-coding-tools-sandbox +import { defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { openaiText } from '@tanstack/ai-openai' +import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' +import { dockerSandbox } from '@tanstack/ai-sandbox-docker' + +const sandbox = await dockerSandbox({ image: 'node:22' }).create({}) +// Without `dir`, the clone goes to the workspace root, /workspace. +await sandbox.git.clone({ url: 'https://github.com/acme/web-app.git' }) + +const root = '/workspace' + +export const sandboxed = defineHarness({ + name: 'acme/sandboxed-coder', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + permissions({ root }), + workspaceTools({ root, backend: sandboxWorkspaceBackend(sandbox) }), + ], +}) +``` + +3. When you are done, stop the sandbox with `await sandbox.destroy()`. + +Every provider works. Get the handle from `create()`, or from `resume()` for a sandbox that runs already. See [Providers](../sandbox/providers). + +In a sandbox, some things are different: + +- Paths are POSIX paths in the sandbox, like `/workspace/src/a.ts`, also on a Windows host. +- The tools do not check where links lead. The sandbox backend has no `realpath`. +- `background: true` needs a sandbox with `capabilities.backgroundProcesses`. +- A sandbox with `capabilities.killableProcesses` off cannot stop a command after the time limit. The call returns, and the command continues in the sandbox. +- `list_files` and `grep` use `rg` in the sandbox. The `@vscode/ripgrep` binary is a host file, so the sandbox does not use it. ## Write your own backend @@ -152,7 +210,7 @@ export const withToml = formatter({ - `when(project)` limits a formatter to some projects. Without it, the formatter is used in every project. - `builtins: false` turns the built-in formatters off. - `timeoutMs` limits each run. The default is 20 seconds. -- On your own backend, give `formatter()` the same `backend` as `workspaceTools()`. +- In a sandbox or on your own backend, give `formatter()` the same `backend` as `workspaceTools()`. A formatter that fails, stops at the time limit, or is not installed does not fail the tool call. The file stays as the tool wrote it, and the plugin sends `FormatFailed`: @@ -208,6 +266,7 @@ export const bounded = defineHarness({ ## What you have now +- The coding tools in a sandbox, with one option. - A backend of your own for another place. - Code that runs after each read and each write. - Formatted files after each write, and tool results that fit the context. diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index 84d5a711e9..ba3646f952 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -112,7 +112,7 @@ terminal.close() await host.close() ``` -`@tanstack/ai-harness/plugins/coding` runs only in Node. The tools use the file system and the shell of the machine. +`@tanstack/ai-harness/plugins/coding` runs only in Node. The tools use the file system and the shell of the machine. For code that you do not trust, [run the tools in a sandbox](./coding-tools-backends#run-the-tools-in-a-sandbox). ## 3. Run it @@ -301,4 +301,4 @@ await session.configure({ cwd: 'web-app' }) - A question before each edit, command, and web fetch. - A working folder for each thread, and a question before a path outside `root`. -Next: run the same tools on your own backend, or format each file after a write. See [Extend the coding tools](./coding-tools-backends). +Next: run the same tools in a sandbox, or format each file after a write. See [Sandbox and extend the coding tools](./coding-tools-backends). diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index fb57124d58..251232eaf7 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -294,6 +294,7 @@ The same check runs in every agent run of the session: - 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. +- [Coding agents](./coding-agents#plan-mode) such as Claude Code start read-only in `plan` mode. ## Rules for your own tool diff --git a/docs/harness/snapshots.md b/docs/harness/snapshots.md index 2c5fc79dd4..001f0b6883 100644 --- a/docs/harness/snapshots.md +++ b/docs/harness/snapshots.md @@ -111,6 +111,26 @@ await session.unrevert() In your own UI, call `client.revert(messageId)` and `client.unrevert()`. The client needs `'undo'` in `expose.commands`. More about reverts: [Go back to an earlier message](./sessions#go-back-to-an-earlier-message). +## In a sandbox + +When the [coding tools run in a sandbox](./coding-tools-backends#run-the-tools-in-a-sandbox), give `snapshots()` the same `backend`: + +```ts group=harness-snapshots-sandbox +import { snapshots } from '@tanstack/ai-harness/plugins/coding' +import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' +import { dockerSandbox } from '@tanstack/ai-sandbox-docker' + +const sandbox = await dockerSandbox({ image: 'node:22' }).create({}) + +export const sandboxHistory = snapshots({ + root: '/workspace', + dataDir: '/tmp/snapshots', + backend: sandboxWorkspaceBackend(sandbox), +}) +``` + +Git then runs in the sandbox, and `dataDir` is a path in the sandbox. The sandbox image needs `git`. + ## Known limits - One level of undo. A second `/undo` answers `Nothing to undo.` After a `/redo`, `/undo` works again. diff --git a/packages/ai-sandbox/package.json b/packages/ai-sandbox/package.json index 2e113c4ab6..0dea764a2c 100644 --- a/packages/ai-sandbox/package.json +++ b/packages/ai-sandbox/package.json @@ -30,6 +30,10 @@ "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" }, + "./harness": { + "types": "./dist/esm/harness.d.ts", + "import": "./dist/esm/harness.js" + }, "./ngrok": { "types": "./dist/esm/ngrok.d.ts", "import": "./dist/esm/ngrok.js" @@ -84,10 +88,14 @@ "peerDependencies": { "@ngrok/ngrok": "^1.0.0", "@tanstack/ai": "workspace:^", + "@tanstack/ai-harness": "workspace:^", "@tanstack/ai-persistence": "workspace:^", "vitest": "^4.1.10" }, "peerDependenciesMeta": { + "@tanstack/ai-harness": { + "optional": true + }, "@ngrok/ngrok": { "optional": true }, @@ -101,6 +109,7 @@ "devDependencies": { "@ngrok/ngrok": "^1.7.0", "@tanstack/ai": "workspace:*", + "@tanstack/ai-harness": "workspace:*", "@tanstack/ai-persistence": "workspace:*", "@vitest/coverage-v8": "4.1.10", "vitest": "^4.1.10" diff --git a/packages/ai-sandbox/src/harness.ts b/packages/ai-sandbox/src/harness.ts new file mode 100644 index 0000000000..b1b6a14265 --- /dev/null +++ b/packages/ai-sandbox/src/harness.ts @@ -0,0 +1,218 @@ +import { EventType, defineAgent } from '@tanstack/ai' +import { defineCommand, definePlugin } from '@tanstack/ai-harness' +import { withSandbox } from './middleware' +import type { + AnyChatMiddleware, + AnyTextAdapter, + JSONSchema, + StreamChunk, +} from '@tanstack/ai' +import type { SandboxMiddlewareOptions } from './middleware' +import type { SandboxDefinition } from './sandbox' + +export { sandboxWorkspaceBackend } from './workspace-backend' + +export interface CodingAgentConfig { + /** A coding-agent adapter: `claudeCodeText`, `codexText`, `grokBuildText`, `acpCompatibleText`, and more. */ + adapter: AnyTextAdapter + /** When the lead model should pick this agent. */ + description: string + /** Options for every call, for example Claude Code's `permissionMode`. */ + modelOptions?: Record + /** + * Options added when the harness `mode` setting is `plan`. Default: Claude + * Code `{ permissionMode: 'plan' }` and Codex `{ sandboxMode: 'read-only' }`. + */ + planModelOptions?: Record +} + +export interface CodingAgentsOptions { + /** Where the agents work. One sandbox per harness thread (or per agent). */ + sandbox: SandboxDefinition + /** The agents, keyed by the tool name the lead model sees. */ + agents: Record + /** + * `shared` (default): every agent of a thread works in one sandbox, one + * agent at a time. `per-agent`: each agent gets its own sandbox. + */ + workspace?: 'shared' | 'per-agent' + /** Passed to `withSandbox`. */ + sandboxOptions?: SandboxMiddlewareOptions +} + +const PLAN_DEFAULTS: Record> = { + 'claude-code': { permissionMode: 'plan' }, + codex: { sandboxMode: 'read-only' }, +} + +const TASK_SCHEMA: JSONSchema = { + type: 'object', + properties: { + task: { + type: 'string', + description: + 'The whole task: what to change, where, and how to check it. The agent sees only this text and the workspace.', + }, + }, + required: ['task'], +} + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null +} + +/** + * Run `open` while holding the lock for `key`. The lock is taken when the + * stream is first read, so a stream nobody reads never holds it. + */ +async function* serialized( + tails: Map>, + key: string, + open: () => AsyncIterable, +): AsyncGenerator { + const previous = tails.get(key) ?? Promise.resolve() + let release = () => {} + const current = new Promise((resolve) => (release = resolve)) + const tail = previous.then(() => current) + tails.set(key, tail) + await previous + try { + yield* open() + } finally { + release() + if (tails.get(key) === tail) tails.delete(key) + } +} + +/** + * A harness plugin that lets the lead model hand coding work to Claude Code, + * Codex, Grok Build, or any ACP agent: + * + * - each agent runs in `sandbox`, so nobody wires `withSandbox` by hand; + * - each agent keeps its own coding session per thread, also after a restart + * (`/fresh [agent]` starts over); + * - the harness `plan` mode starts the agents read-only. + * + * @example + * ```ts + * codingAgents({ + * sandbox: defineSandbox({ id: 'code', provider: localProcessSandbox() }), + * agents: { + * claude_code: { adapter: claudeCodeText('claude-opus-4-8'), description: 'Large refactors' }, + * codex: { adapter: codexText('gpt-5.5'), description: 'Quick fixes and tests' }, + * }, + * }) + * ``` + */ +export function codingAgents(options: CodingAgentsOptions) { + const shared = (options.workspace ?? 'shared') === 'shared' + const sandboxOptions = options.sandboxOptions ?? {} + // ponytail: an in-process lock per sandbox. Several hosts on one sandbox need a shared lock store. + const tails = new Map>() + const names = Object.keys(options.agents) + + return definePlugin({ + name: 'tanstack/coding-agents', + setup: (ctx) => { + const state = ctx.state<{ sessions: Record }>({ + sessions: {}, + }) + // Saves the session id the adapter reports (`claude-code.session-id`, ...). + const capture = (agent: string): AnyChatMiddleware => ({ + name: 'tanstack/coding-agents:session', + onChunk: async (_run, chunk) => { + if (chunk.type !== EventType.CUSTOM) return + if (!chunk.name.endsWith('.session-id')) return + const value: unknown = chunk.value + if (!isRecord(value) || typeof value.sessionId !== 'string') return + const sessionId = value.sessionId + await state.update((current) => ({ + sessions: { ...current.sessions, [agent]: sessionId }, + })) + }, + }) + + const subagents = Object.entries(options.agents).map(([name, config]) => + defineAgent({ + name, + description: config.description, + inputSchema: TASK_SCHEMA, + run: async (run) => { + const input: unknown = run.input + const task = + isRecord(input) && typeof input.task === 'string' + ? input.task + : '' + const plan = + ctx.config.get('mode') === 'plan' + ? (config.planModelOptions ?? + PLAN_DEFAULTS[config.adapter.name]) + : {} + // Fail closed: an agent with no read-only options must not edit. + if (plan === undefined) { + throw new Error( + `${name} cannot run in plan mode: it has no planModelOptions.`, + ) + } + // The harness thread, not `run.threadId`: a child's own thread is + // `:` by default. + const thread = ctx.session.threadId + const threadId = shared ? thread : `${thread}:${name}` + // Inside the lock, so a queued call resumes the session of the + // call before it. + async function* open(): AsyncGenerator { + const sessionId = (await state.get()).sessions[name] + const modelOptions = { + ...config.modelOptions, + ...plan, + ...(sessionId ? { sessionId } : {}), + } + yield* run.chat({ + adapter: config.adapter, + messages: [{ role: 'user', content: task }], + threadId, + ...(Object.keys(modelOptions).length > 0 + ? { modelOptions } + : {}), + middleware: [ + withSandbox(options.sandbox, sandboxOptions), + capture(name), + ], + }) as AsyncIterable + } + // One call at a time per sandbox, also for the same agent in + // `per-agent` mode. + return serialized(tails, threadId, open) + }, + }), + ) + + return { + subagents, + commands: { + fresh: defineCommand({ + description: `Start new coding-agent sessions (all, or one of: ${names.join(', ')})`, + run: async (input: unknown) => { + const name = typeof input === 'string' ? input.trim() : '' + if (name && !names.includes(name)) { + return `Unknown coding agent "${name}". Agents: ${names.join(', ')}.` + } + await state.update((current) => ({ + sessions: name + ? Object.fromEntries( + Object.entries(current.sessions).filter( + ([agent]) => agent !== name, + ), + ) + : {}, + })) + return name + ? `${name} starts a new session next time.` + : 'Every coding agent starts a new session next time.' + }, + }), + }, + } + }, + }) +} diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts new file mode 100644 index 0000000000..b52ef87841 --- /dev/null +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -0,0 +1,154 @@ +import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' +import type { ExecResult, SandboxHandle } from './contracts' + +/** + * A `WorkspaceBackend` over a live sandbox. With it, `workspaceTools()` + * reads and writes files and runs commands in the sandbox, not on this + * machine. You start the sandbox before, and you stop it after. + * + * - Paths are POSIX paths in the sandbox, like `/workspace/src/a.ts`. + * - `exec` gives exit code 124 when `timeoutMs` stops a command, and 1 when + * `signal` stops it. A sandbox with `capabilities.killableProcesses` off + * cannot stop the command, so it continues in the sandbox. + * - `spawn` is there only when the sandbox has + * `capabilities.backgroundProcesses`. + * - `stat` gives `mtimeMs: 0`. The sandbox file system does not give the + * time of the last change. + * - There is no `realpath`, so the tools do not check where links go. + * + * @param handle - The sandbox, from `provider.create()` or `provider.resume()`. + * + * @example + * ```ts + * const handle = await localProcessSandbox().create({}) + * workspaceTools({ + * root: '/workspace', + * backend: sandboxWorkspaceBackend(handle), + * }) + * ``` + */ +export function sandboxWorkspaceBackend(handle: SandboxHandle) { + const { fs } = handle + + // The read of the sandbox follows links, so this follows links too. A + // path that `readBytes` cannot read is a folder. + // ponytail: reads the whole file for its size. Only links, and a sandbox + // without `lstat`, come here. + const follow: WorkspaceBackend['stat'] = async (path) => { + if (!(await fs.exists(path))) return undefined + const bytes = await fs.readBytes(path).catch(() => undefined) + if (!bytes) return { type: 'dir', size: 0, mtimeMs: 0 } + return { type: 'file', size: bytes.byteLength, mtimeMs: 0 } + } + + const stat: WorkspaceBackend['stat'] = async (path) => { + if (!fs.lstat) return follow(path) + const info = await fs.lstat(path) + if (!info) return undefined + switch (info.type) { + case 'file': + return { type: 'file', size: info.size, mtimeMs: 0 } + case 'dir': + return { type: 'dir', size: 0, mtimeMs: 0 } + // Like `fs.stat` on the host, a fifo or a socket is an empty file. + case 'other': + return { type: 'file', size: 0, mtimeMs: 0 } + case 'symlink': + return follow(path) + } + } + + const spawn: WorkspaceBackend['spawn'] = (command, options) => { + // ponytail: keeps all output in memory, like `hostBackend`. + let output = '' + const collect = async (stream: AsyncIterable) => { + for await (const chunk of stream) output += chunk + } + const started = handle.process.spawn(command, options) + // Wait for the end of the output too, so `output()` is complete. + const exited = started + .then(async (job) => { + const [exitCode] = await Promise.all([ + job.wait(), + collect(job.stdout), + collect(job.stderr), + ]) + return { exitCode } + }) + .catch((error: unknown) => { + output += String(error) + return { exitCode: 1 } + }) + return { + wait: () => exited, + // Before the process starts, this kills it when it starts. + kill: () => { + started.then((job) => job.kill()).catch(() => undefined) + }, + output: () => output, + } + } + + return { + shell: 'sh', + readFile: (path) => fs.readBytes(path), + writeFile: (path, data) => fs.write(path, data), + // The sandbox `remove` also removes folders, and it does not fail for a + // missing path. This backend removes one file only. + remove: async (path) => { + const isFile = (await stat(path))?.type === 'file' + if (!isFile) throw new Error(`Cannot remove ${path}: no such file.`) + await fs.remove(path) + }, + stat, + readdir: async (path) => { + const entries = await fs.list(path) + const { lstat } = fs + // ponytail: one `lstat` for each entry, so a big folder is slow in a + // remote sandbox. Without `lstat`, a link has the type of its target. + return Promise.all( + entries.map(async (entry) => { + const isLink = + lstat !== undefined && (await lstat(entry.path))?.type === 'symlink' + const type = isLink ? ('link' as const) : entry.type + return { name: entry.name, type } + }), + ) + }, + exec: (command, options = {}) => { + const { cwd, env, timeoutMs, signal } = options + const controller = new AbortController() + // Stop with an exit code. The first reason wins. The result does not + // wait for the sandbox, because some sandboxes cannot stop a command. + let stop = (_exitCode: number) => {} + const stopped = new Promise((resolve) => { + stop = (exitCode) => { + resolve({ exitCode, stdout: '', stderr: '' }) + controller.abort() + } + }) + // A timeout of 0 means no timeout. + const timer = timeoutMs + ? setTimeout(() => stop(124), timeoutMs) + : undefined + const onAbort = () => stop(1) + signal?.addEventListener('abort', onAbort) + if (signal?.aborted) onAbort() + // A command that does not start gives exit code 1, like `hostBackend`. + const run = handle.process + .exec(command, { cwd, env, signal: controller.signal }) + .catch( + (error: unknown): ExecResult => ({ + exitCode: 1, + stdout: '', + stderr: String(error), + }), + ) + return Promise.race([stopped, run]).finally(() => { + clearTimeout(timer) + signal?.removeEventListener('abort', onAbort) + }) + }, + ...(handle.capabilities.backgroundProcesses ? { spawn } : {}), + } satisfies WorkspaceBackend +} diff --git a/packages/ai-sandbox/tests/coding-agents.test.ts b/packages/ai-sandbox/tests/coding-agents.test.ts new file mode 100644 index 0000000000..da467217fb --- /dev/null +++ b/packages/ai-sandbox/tests/coding-agents.test.ts @@ -0,0 +1,383 @@ +import { mkdtempSync, writeFileSync } from 'node:fs' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { codingAgents } from '../src/harness' +import { defineSandbox } from '../src/sandbox' +import { InMemorySandboxInstanceStore } from '../src/instance-store' +import { defineWorkspace } from '../src/workspace' +import { makeFakeProvider } from './fakes' +import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' + +const now = () => Date.now() +// A small local project the sandboxes start from. +const project = mkdtempSync(join(tmpdir(), 'coding-agents-')) +writeFileSync(join(project, 'README.md'), '# demo\n') +const 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, +} + +function textChunks(content: string): Array { + return [ + { + type: EventType.TEXT_MESSAGE_START, + messageId: `m-${content}`, + role: 'assistant', + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: `m-${content}`, + delta: content, + timestamp: now(), + }, + { + type: EventType.TEXT_MESSAGE_END, + messageId: `m-${content}`, + timestamp: now(), + }, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'stop' } }, + }, + ] +} + +/** A lead model that calls the given tools in one step, then answers. */ +function lead(steps: Array>) { + let call = 0 + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'lead', + model: 'lead', + '~types': TYPES, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: () => { + const step = steps[call] + call += 1 + return (async function* (): AsyncGenerator { + yield { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + } + if (!step) { + yield* textChunks('lead done') + return + } + for (const [index, [tool, task]] of step.entries()) { + const id = `call-${call}-${index}` + yield { + type: EventType.TOOL_CALL_START, + toolCallId: id, + toolCallName: tool, + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: id, + delta: JSON.stringify({ task }), + timestamp: now(), + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId: id, + timestamp: now(), + } + } + yield { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason: 'tool_calls' } }, + } + })() + }, + } + return adapter +} + +/** + * A stand-in coding agent. It records each call, reports a new session id the + * way Claude Code does (`.session-id`), and answers with the task. + */ +function codingAgent(name: string, options: { wait?: Promise } = {}) { + const calls: Array = [] + const log: Array = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name, + model: `${name}-model`, + '~types': TYPES, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (chatOptions: any) => { + calls.push(chatOptions) + const index = calls.length + return (async function* (): AsyncGenerator { + log.push(`${name} start ${index}`) + yield { + type: EventType.RUN_STARTED, + runId: 'c', + threadId: 't', + timestamp: now(), + } + yield { + type: EventType.CUSTOM, + name: `${name}.session-id`, + value: { sessionId: `${name}-session-${index}` }, + timestamp: now(), + } + if (options.wait) await options.wait + const task = String(chatOptions.messages.at(-1)?.content ?? '') + yield* textChunks(`${name} did: ${task}`) + log.push(`${name} end ${index}`) + })() + }, + } + return { adapter, calls, log } +} + +async function open(input: { + leadSteps: Array> + agents: Record< + string, + { adapter: AnyTextAdapter; planModelOptions?: Record } + > + workspace?: 'shared' | 'per-agent' + persistence?: ReturnType +}) { + const provider = makeFakeProvider() + const persistence = input.persistence ?? memoryPersistence() + const host = createHarnessHost({ persistence }) + const harness = defineHarness({ + name: 'test/coding-lead', + adapter: lead(input.leadSteps), + plugins: () => [ + permissions(), + codingAgents({ + sandbox: defineSandbox({ + id: 'code', + provider, + workspace: defineWorkspace({ + source: { type: 'local', path: project }, + }), + }), + agents: Object.fromEntries( + Object.entries(input.agents).map(([name, agent]) => [ + name, + { description: `The ${name} agent`, ...agent }, + ]), + ), + ...(input.workspace ? { workspace: input.workspace } : {}), + // A store per test, so no test resumes a sandbox of another test. + sandboxOptions: { instances: new InMemorySandboxInstanceStore() }, + }), + ], + }) + const session = await host.open(harness, { threadId: 't' }) + return { host, session, provider, persistence, harness } +} + +describe('codingAgents', () => { + it('runs a delegate in the sandbox and resumes its session on the next call', async () => { + const claude = codingAgent('claude-code') + const { host, session, provider } = await open({ + leadSteps: [ + [['claude_code', 'add a test']], + [], + [['claude_code', 'fix the test']], + ], + agents: { claude_code: { adapter: claude.adapter } }, + }) + await session.prompt('first') + await session.prompt('second') + + expect(claude.calls).toHaveLength(2) + expect(claude.calls[0].messages).toEqual([ + { role: 'user', content: 'add a test' }, + ]) + expect(claude.calls[0].modelOptions?.sessionId).toBeUndefined() + // The second call resumes the session the first call reported. + expect(claude.calls[1].modelOptions).toEqual({ + sessionId: 'claude-code-session-1', + }) + expect(claude.calls[0].threadId).toBe('t') + // One sandbox for the thread, created once. + expect(provider.calls.create).toBe(1) + await host.close() + }) + + it('lists the delegates as session agents and forgets sessions with /fresh', async () => { + const claude = codingAgent('claude-code') + const codex = codingAgent('codex') + const { host, session } = await open({ + leadSteps: [ + [['claude_code', 'one']], + [], + [['codex', 'two']], + [], + [['claude_code', 'three']], + ], + agents: { + claude_code: { adapter: claude.adapter }, + codex: { adapter: codex.adapter }, + }, + }) + expect(session.registry.list().map((agent) => agent.name)).toEqual([ + 'claude_code', + 'codex', + ]) + await session.prompt('a') + await session.prompt('b') + expect(await session.command('fresh', 'nope')).toBe( + 'Unknown coding agent "nope". Agents: claude_code, codex.', + ) + expect(await session.command('fresh', 'claude_code')).toBe( + 'claude_code starts a new session next time.', + ) + await session.prompt('c') + expect(claude.calls[1].modelOptions?.sessionId).toBeUndefined() + expect(await session.command('fresh')).toBe( + 'Every coding agent starts a new session next time.', + ) + await host.close() + }) + + it('keeps sessions across a host restart', async () => { + const persistence = memoryPersistence() + const first = codingAgent('claude-code') + const before = await open({ + leadSteps: [[['claude_code', 'start']]], + agents: { claude_code: { adapter: first.adapter } }, + persistence, + }) + await before.session.prompt('go') + await before.host.close() + + const second = codingAgent('claude-code') + const after = await open({ + leadSteps: [[['claude_code', 'continue']]], + agents: { claude_code: { adapter: second.adapter } }, + persistence, + }) + await after.session.prompt('go on') + expect(second.calls[0].modelOptions).toEqual({ + sessionId: 'claude-code-session-1', + }) + await after.host.close() + }) + + it('runs shared-workspace delegates one at a time, and per-agent ones side by side', async () => { + let release = () => {} + const gate = new Promise((resolve) => (release = resolve)) + const claude = codingAgent('claude-code', { wait: gate }) + const codex = codingAgent('codex') + const shared = await open({ + leadSteps: [], + agents: { + claude_code: { adapter: claude.adapter }, + codex: { adapter: codex.adapter }, + }, + }) + // Two background runs at once, like two /agent commands. + const first = shared.session.agent('claude_code')?.start({ task: 'edit a' }) + const second = shared.session.agent('codex')?.start({ task: 'edit b' }) + await new Promise((resolve) => setTimeout(resolve, 30)) + // Codex waits for the lock while Claude Code holds the shared workspace. + expect(claude.calls).toHaveLength(1) + expect(codex.calls).toHaveLength(0) + release() + await first + await second + expect(codex.calls).toHaveLength(1) + expect(codex.calls[0].threadId).toBe('t') + expect(shared.provider.calls.create).toBe(1) + await shared.host.close() + + let releaseSlow = () => {} + const slowGate = new Promise((resolve) => (releaseSlow = resolve)) + const slow = codingAgent('claude-code', { wait: slowGate }) + const fast = codingAgent('codex') + const apart = await open({ + leadSteps: [], + agents: { + claude_code: { adapter: slow.adapter }, + codex: { adapter: fast.adapter }, + }, + workspace: 'per-agent', + }) + const slowRun = apart.session + .agent('claude_code') + ?.start({ task: 'edit a' }) + const fastRun = apart.session.agent('codex')?.start({ task: 'edit b' }) + await fastRun + // Codex finished while Claude Code still works in its own sandbox. + expect(fast.calls[0].threadId).toBe('t:codex') + expect(slow.log).toEqual(['claude-code start 1']) + releaseSlow() + await slowRun + expect(slow.calls[0].threadId).toBe('t:claude_code') + expect(apart.provider.calls.create).toBe(2) + await apart.host.close() + }) + + it('starts delegates read-only in plan mode, and refuses one without plan options', async () => { + const claude = codingAgent('claude-code') + const codex = codingAgent('codex') + const other = codingAgent('grok-build') + const custom = codingAgent('acp') + const { host, session } = await open({ + // Two steps: tool calls in one step run at the same time, and the + // harness runs at most 3 subagents at once. + leadSteps: [ + [ + ['claude_code', 'look'], + ['codex', 'look'], + ['grok', 'look'], + ], + [['custom', 'look']], + ], + agents: { + claude_code: { adapter: claude.adapter }, + codex: { adapter: codex.adapter }, + grok: { adapter: other.adapter }, + custom: { + adapter: custom.adapter, + planModelOptions: { readOnly: true }, + }, + }, + }) + await session.setConfig({ key: 'mode', value: 'plan' }) + await session.prompt('plan it') + expect(claude.calls[0].modelOptions).toMatchObject({ + permissionMode: 'plan', + }) + expect(codex.calls[0].modelOptions).toMatchObject({ + sandboxMode: 'read-only', + }) + // No plan options for grok-build: it does not run in plan mode. + expect(other.calls).toHaveLength(0) + expect(custom.calls[0].modelOptions).toMatchObject({ readOnly: true }) + await host.close() + }) +}) diff --git a/packages/ai-sandbox/tests/workspace-backend.test.ts b/packages/ai-sandbox/tests/workspace-backend.test.ts new file mode 100644 index 0000000000..78958d92dd --- /dev/null +++ b/packages/ai-sandbox/tests/workspace-backend.test.ts @@ -0,0 +1,456 @@ +import { describe, expect, it } from 'vitest' +import { EventType } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { workspaceTools } from '@tanstack/ai-harness/plugins/coding' +import { sandboxWorkspaceBackend } from '../src/harness' +import { FULL_CAPS, makeFakeHandle } from './fakes' +import type { AnyTextAdapter, StreamChunk, TextOptions } from '@tanstack/ai' +import type { + ExecResult, + SandboxCapabilities, + SandboxFs, + SandboxHandle, +} from '../src/contracts' + +const encoder = new TextEncoder() +const decoder = new TextDecoder() + +interface FsSeed { + files?: Record + dirs?: Array + /** A link and its target. */ + links?: Record + /** Paths that are not a file, a folder, or a link, like a fifo. */ + others?: Array +} + +/** + * A sandbox file system in memory, by POSIX path. Like a real sandbox: + * `readBytes`, `exists`, and `list` follow links, `lstat` does not, `write` + * makes the parent folders, and `remove` is `rm -rf`. With + * `withLstat: false`, it has no `lstat`. + */ +function memoryFs(seed: FsSeed, withLstat = true) { + const files = new Map( + Object.entries(seed.files ?? {}).map(([path, text]) => [ + path, + encoder.encode(text), + ]), + ) + const dirs = new Set(['/', ...(seed.dirs ?? [])]) + const links = new Map(Object.entries(seed.links ?? {})) + const others = new Set(seed.others ?? []) + const target = (path: string) => links.get(path) ?? path + const parent = (path: string) => path.slice(0, path.lastIndexOf('/')) || '/' + const baseName = (path: string) => path.slice(path.lastIndexOf('/') + 1) + const paths = () => [...files.keys(), ...dirs, ...links.keys(), ...others] + + const lstat: SandboxFs['lstat'] = async (path) => { + const data = files.get(path) + if (links.has(path)) return { type: 'symlink', mode: 0o120777 } + if (data) return { type: 'file', mode: 0o100644, size: data.byteLength } + if (dirs.has(path)) return { type: 'dir', mode: 0o40755 } + if (others.has(path)) return { type: 'other', mode: 0o10644 } + return undefined + } + const readBytes: SandboxFs['readBytes'] = async (path) => { + const data = files.get(target(path)) + if (data) return data + throw new Error(`${dirs.has(target(path)) ? 'EISDIR' : 'ENOENT'}: ${path}`) + } + const fs: SandboxFs = { + read: async (path) => decoder.decode(await readBytes(path)), + readBytes, + write: async (path, data) => { + for (let dir = parent(path); !dirs.has(dir); dir = parent(dir)) { + dirs.add(dir) + } + files.set(path, typeof data === 'string' ? encoder.encode(data) : data) + }, + list: async (path) => { + if (!dirs.has(target(path))) throw new Error(`ENOENT: ${path}`) + const names = paths() + .filter((each) => each !== '/' && parent(each) === target(path)) + .map(baseName) + .sort() + return names.map((name) => { + const type: 'file' | 'dir' = dirs.has(target(`${path}/${name}`)) + ? 'dir' + : 'file' + return { name, path: `${path}/${name}`, type } + }) + }, + mkdir: async (path) => { + dirs.add(path) + }, + remove: async (path) => { + const gone = paths().filter( + (each) => each === path || each.startsWith(`${path}/`), + ) + for (const each of gone) { + files.delete(each) + dirs.delete(each) + links.delete(each) + } + }, + rename: () => Promise.reject(new Error('rename is not used')), + exists: async (path) => files.has(target(path)) || dirs.has(target(path)), + ...(withLstat ? { lstat } : {}), + } + return { fs, files, dirs } +} + +async function* fromChunks(chunks: Array) { + for (const chunk of chunks) yield chunk +} + +/** + * Sandbox processes. `exec` answers like a shell that prints the command, + * the folder, and `$AGENT`. A `sleep` command never ends, like in a sandbox + * that cannot stop it, and only notes the abort in `aborted`. + * + * `spawn` prints ` out` and ` err`. `serve` runs until it + * is killed, then ends with 143. `nope` does not start. + */ +function fakeProcesses() { + const aborted: Array = [] + const killed: Array = [] + const process: SandboxHandle['process'] = { + exec: (command, options = {}) => { + if (command.startsWith('sleep')) { + options.signal?.addEventListener('abort', () => aborted.push(command)) + return new Promise(() => {}) + } + return Promise.resolve({ + exitCode: 3, + stdout: `ran ${command} in ${options.cwd} with AGENT=${options.env?.AGENT}`, + stderr: 'warn', + }) + }, + spawn: async (command) => { + if (command === 'nope') throw new Error('nope: not found') + let end = () => {} + const ended = + command === 'serve' + ? new Promise((resolve) => (end = resolve)) + : Promise.resolve() + async function* stdout() { + yield `${command} out\n` + await ended + } + return { + pid: 7, + stdout: stdout(), + stderr: fromChunks([`${command} err\n`]), + stdin: { write: async () => {}, end: async () => {} }, + wait: async () => { + await ended + return command === 'serve' ? 143 : 0 + }, + kill: async () => { + killed.push(command) + end() + }, + } + }, + } + return { process, aborted, killed } +} + +/** A sandbox handle with `memoryFs` and `fakeProcesses`. */ +function fakeSandbox( + seed: FsSeed = {}, + options: { withLstat?: boolean; caps?: SandboxCapabilities } = {}, +) { + const { fs, files, dirs } = memoryFs(seed, options.withLstat) + const { process, aborted, killed } = fakeProcesses() + const handle: SandboxHandle = { + ...makeFakeHandle('sandbox-1', 'fake', options.caps), + fs, + process, + } + const backend = sandboxWorkspaceBackend(handle) + return { backend, files, dirs, aborted, killed } +} + +const SEED: FsSeed = { + files: { '/workspace/a.txt': 'hello' }, + dirs: ['/workspace', '/workspace/src'], + links: { + '/workspace/to-file': '/workspace/a.txt', + '/workspace/to-src': '/workspace/src', + '/workspace/broken': '/workspace/gone', + }, + others: ['/workspace/pipe'], +} + +describe('sandboxWorkspaceBackend files', () => { + it('uses POSIX paths and quoting, and checks no links', () => { + const { backend } = fakeSandbox() + expect(backend.shell).toBe('sh') + expect(backend).not.toHaveProperty('realpath') + }) + + it('reads and writes the bytes of a file in the sandbox', async () => { + const { backend, files } = fakeSandbox(SEED) + expect(decoder.decode(await backend.readFile('/workspace/a.txt'))).toBe( + 'hello', + ) + await backend.writeFile('/workspace/new/b.txt', 'made') + await backend.writeFile('/workspace/c.bin', new Uint8Array([1, 2])) + expect(decoder.decode(files.get('/workspace/new/b.txt'))).toBe('made') + expect(files.get('/workspace/c.bin')).toEqual(new Uint8Array([1, 2])) + }) + + it.each([ + { withLstat: true, path: '/workspace/a.txt', type: 'file', size: 5 }, + { withLstat: true, path: '/workspace/src', type: 'dir', size: 0 }, + { withLstat: true, path: '/workspace/to-file', type: 'file', size: 5 }, + { withLstat: true, path: '/workspace/to-src', type: 'dir', size: 0 }, + { withLstat: true, path: '/workspace/pipe', type: 'file', size: 0 }, + { withLstat: false, path: '/workspace/a.txt', type: 'file', size: 5 }, + { withLstat: false, path: '/workspace/src', type: 'dir', size: 0 }, + { withLstat: false, path: '/workspace/to-file', type: 'file', size: 5 }, + ])( + 'stat gives a $type of size $size for $path (lstat: $withLstat)', + async ({ withLstat, path, type, size }) => { + const { backend } = fakeSandbox(SEED, { withLstat }) + expect(await backend.stat(path)).toEqual({ type, size, mtimeMs: 0 }) + }, + ) + + it.each([true, false])( + 'stat gives undefined for a missing path and a broken link (lstat: %s)', + async (withLstat) => { + const { backend } = fakeSandbox(SEED, { withLstat }) + expect(await backend.stat('/workspace/gone')).toBeUndefined() + expect(await backend.stat('/workspace/broken')).toBeUndefined() + }, + ) + + it('lists a folder with links as links when the sandbox has lstat', async () => { + const { backend } = fakeSandbox(SEED) + expect(await backend.readdir('/workspace')).toEqual([ + { name: 'a.txt', type: 'file' }, + { name: 'broken', type: 'link' }, + { name: 'pipe', type: 'file' }, + { name: 'src', type: 'dir' }, + { name: 'to-file', type: 'link' }, + { name: 'to-src', type: 'link' }, + ]) + }) + + it('lists a link with the type of its target when the sandbox has no lstat', async () => { + const { backend } = fakeSandbox(SEED, { withLstat: false }) + expect(await backend.readdir('/workspace')).toEqual([ + { name: 'a.txt', type: 'file' }, + { name: 'broken', type: 'file' }, + { name: 'pipe', type: 'file' }, + { name: 'src', type: 'dir' }, + { name: 'to-file', type: 'file' }, + { name: 'to-src', type: 'dir' }, + ]) + }) + + it('removes one file, and refuses a folder or a missing file', async () => { + const { backend, files, dirs } = fakeSandbox({ + files: { '/workspace/a.txt': 'a', '/workspace/src/b.ts': 'b' }, + dirs: ['/workspace', '/workspace/src'], + }) + await backend.remove('/workspace/a.txt') + expect(files.has('/workspace/a.txt')).toBe(false) + await expect(backend.remove('/workspace/src')).rejects.toThrow( + 'Cannot remove /workspace/src: no such file.', + ) + await expect(backend.remove('/workspace/gone')).rejects.toThrow( + 'Cannot remove /workspace/gone: no such file.', + ) + expect(dirs.has('/workspace/src')).toBe(true) + expect(files.has('/workspace/src/b.ts')).toBe(true) + }) +}) + +describe('sandboxWorkspaceBackend exec', () => { + it('runs a command in the sandbox with the folder and the env', async () => { + const { backend } = fakeSandbox() + expect( + await backend.exec('ls', { cwd: '/workspace', env: { AGENT: '1' } }), + ).toEqual({ + exitCode: 3, + stdout: 'ran ls in /workspace with AGENT=1', + stderr: 'warn', + }) + }) + + it('stops a command at the timeout with exit code 124', async () => { + const { backend, aborted } = fakeSandbox() + expect(await backend.exec('sleep 60', { timeoutMs: 20 })).toEqual({ + exitCode: 124, + stdout: '', + stderr: '', + }) + expect(aborted).toEqual(['sleep 60']) + }) + + it('stops a command with exit code 1 when the signal aborts', async () => { + const { backend, aborted } = fakeSandbox() + const controller = new AbortController() + const result = backend.exec('sleep 60', { + timeoutMs: 60_000, + signal: controller.signal, + }) + controller.abort() + expect(await result).toEqual({ exitCode: 1, stdout: '', stderr: '' }) + expect(aborted).toEqual(['sleep 60']) + }) +}) + +describe('sandboxWorkspaceBackend spawn', () => { + it('gives the exit code and all output of a background command', async () => { + const { backend } = fakeSandbox() + const job = backend.spawn?.('build', { cwd: '/workspace' }) + expect(await job?.wait()).toEqual({ exitCode: 0 }) + expect(job?.output()).toBe('build out\nbuild err\n') + }) + + it('kills a background command, also before it started', async () => { + const { backend, killed } = fakeSandbox() + const job = backend.spawn?.('serve') + job?.kill() + expect(await job?.wait()).toEqual({ exitCode: 143 }) + expect(killed).toEqual(['serve']) + expect(job?.output()).toBe('serve out\nserve err\n') + }) + + it('gives exit code 1 and the error when the command does not start', async () => { + const { backend } = fakeSandbox() + const job = backend.spawn?.('nope') + expect(await job?.wait()).toEqual({ exitCode: 1 }) + expect(job?.output()).toBe('Error: nope: not found') + }) + + it('has no spawn when the sandbox cannot run background processes', () => { + const { backend } = fakeSandbox( + {}, + { caps: { ...FULL_CAPS, backgroundProcesses: false } }, + ) + expect(backend.spawn).toBeUndefined() + }) +}) + +const now = () => Date.now() + +/** The chunks of one model reply: `chunks` in one run. */ +function reply( + finishReason: 'tool_calls' | 'stop', + chunks: Array = [], +) { + const all: Array = [ + { + type: EventType.RUN_STARTED, + runId: 'r', + threadId: 't', + timestamp: now(), + }, + ...chunks, + { + type: EventType.RUN_FINISHED, + runId: 'r', + threadId: 't', + timestamp: now(), + metadata: { tanstack: { finishReason } }, + }, + ] + return all +} + +/** The chunks of one model reply that calls `tool` with `args`. */ +function toolCall(id: string, tool: string, args: Record) { + return reply('tool_calls', [ + { + type: EventType.TOOL_CALL_START, + toolCallId: id, + toolCallName: tool, + timestamp: now(), + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: id, + delta: JSON.stringify(args), + timestamp: now(), + }, + { type: EventType.TOOL_CALL_END, toolCallId: id, timestamp: now() }, + ]) +} + +/** + * A model that gives `replies` in order, then ends the turn. `calls` has the + * options of each call. + */ +function scriptedModel(replies: Array>) { + const calls: Array = [] + const adapter: AnyTextAdapter = { + kind: 'text', + name: 'scripted', + model: 'scripted', + '~types': { + providerOptions: {}, + inputModalities: ['text'], + messageMetadataByModality: {}, + toolCapabilities: [], + toolCallMetadata: undefined, + systemPromptMetadata: undefined, + }, + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + chatStream: (options) => { + calls.push(options) + const chunks = replies[calls.length - 1] ?? reply('stop') + return (async function* () { + yield* chunks + })() + }, + } + return { adapter, calls } +} + +describe('sandboxWorkspaceBackend with workspaceTools', () => { + it('writes and reads a file under /workspace, with POSIX paths for the model', async () => { + const sandbox = fakeSandbox({ dirs: ['/workspace'] }) + const { adapter, calls } = scriptedModel([ + toolCall('c1', 'write_file', { + path: 'src/a.ts', + content: 'export const a = 1\n', + }), + toolCall('c2', 'read_file', { path: 'src/a.ts' }), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/sandbox-workspace', + adapter, + plugins: () => [ + workspaceTools({ + root: '/workspace', + backend: sandbox.backend, + web: false, + }), + ], + }), + { threadId: 't' }, + ) + await session.prompt('work') + await host.close() + + expect(decoder.decode(sandbox.files.get('/workspace/src/a.ts'))).toBe( + 'export const a = 1\n', + ) + const results = calls[2]?.messages + .filter((message) => message.role === 'tool') + .map((message) => message.content) + expect(results).toEqual(['Wrote src/a.ts.', '1\texport const a = 1']) + expect(calls[0]?.systemPrompts).toContainEqual( + expect.stringContaining('Your workspace is /workspace.'), + ) + }) +}) diff --git a/packages/ai-sandbox/vite.config.ts b/packages/ai-sandbox/vite.config.ts index 5eaa0e507a..e7d52c2e2a 100644 --- a/packages/ai-sandbox/vite.config.ts +++ b/packages/ai-sandbox/vite.config.ts @@ -34,7 +34,12 @@ export default mergeConfig( // (`@tanstack/ai-sandbox/ngrok`), which lazy-loads the optional `@ngrok/ngrok` // peer dep so the core never pulls in its native binary + the SandboxInstanceStore // conformance testkit (`@tanstack/ai-sandbox/testkit`). - entry: ['./src/index.ts', './src/ngrok.ts', './src/testkit/conformance.ts'], + entry: [ + './src/index.ts', + './src/harness.ts', + './src/ngrok.ts', + './src/testkit/conformance.ts', + ], srcDir: './src', // The conformance testkit imports Vitest; keep it external so the built // artifact references the consumer's Vitest at test time. diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 00201dc041..eba60e1814 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -3138,6 +3138,9 @@ 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 diff --git a/testing/e2e/fixtures/tools-test/coding-tools.json b/testing/e2e/fixtures/tools-test/coding-tools.json index f323e1537f..12d6c61760 100644 --- a/testing/e2e/fixtures/tools-test/coding-tools.json +++ b/testing/e2e/fixtures/tools-test/coding-tools.json @@ -7,7 +7,10 @@ }, "response": { "toolCalls": [ - { "name": "read_file", "arguments": "{\"path\":\"notes.txt\"}" } + { + "name": "read_file", + "arguments": "{\"path\":\"notes.txt\"}" + } ] } }, @@ -72,7 +75,12 @@ "sequenceIndex": 0 }, "response": { - "toolCalls": [{ "name": "grep", "arguments": "{\"pattern\":\"beta\"}" }] + "toolCalls": [ + { + "name": "grep", + "arguments": "{\"pattern\":\"beta\"}" + } + ] } }, { @@ -172,7 +180,32 @@ "userMessage": "[coding-format-ask] run test", "sequenceIndex": 1 }, - "response": { "content": "Wrote a.ts." } + "response": { + "content": "Wrote a.ts." + } + }, + { + "match": { + "userMessage": "[coding-sandbox] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "edit_file", + "arguments": "{\"path\":\"notes.txt\",\"old\":\"alpha\",\"new\":\"sandboxed\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[coding-sandbox] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "Edited notes.txt in the sandbox." + } } ] } diff --git a/testing/e2e/src/routes/api.coding-tools.ts b/testing/e2e/src/routes/api.coding-tools.ts index 7d147b00d3..bb75d5c5c9 100644 --- a/testing/e2e/src/routes/api.coding-tools.ts +++ b/testing/e2e/src/routes/api.coding-tools.ts @@ -14,6 +14,8 @@ import { snapshots, workspaceTools, } from '@tanstack/ai-harness/plugins/coding' +import { sandboxWorkspaceBackend } from '@tanstack/ai-sandbox/harness' +import { localProcessSandbox } from '@tanstack/ai-sandbox-local-process' import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' import { createTextAdapter } from '@/lib/providers' @@ -103,6 +105,40 @@ async function* revertTurns( } } +/** + * The `coding-sandbox` scenario: one `edit_file` turn with the tools in a + * local-process sandbox on `root`. The tools see `/workspace`, and the + * sandbox maps it to `root`. + */ +async function* sandboxTurn( + root: string, + adapter: AnyTextAdapter, +): AsyncGenerator { + const sandbox = await localProcessSandbox({ dir: root }).create({}) + const host = createHarnessHost() + try { + const session = await host.open( + defineHarness({ + name: 'e2e/coding-tools-sandbox', + adapter, + plugins: () => [ + workspaceTools({ + root: '/workspace', + editStyle: 'edit', + backend: sandboxWorkspaceBackend(sandbox), + }), + ], + }), + { threadId: 'coding-sandbox' }, + ) + yield* session.prompt('[coding-sandbox] run test').stream() + } finally { + await host.close() + // With `dir`, destroy keeps the folder. + await sandbox.destroy() + } +} + /** * The coding tools of `@tanstack/ai-harness/plugins/coding` on real files. * Body: `{ scenario, testId, root }`. `root` must be a folder in the temp @@ -139,7 +175,9 @@ export const Route = createFileRoute('/api/coding-tools')({ return toServerSentEventsResponse( scenario === 'coding-revert' ? revertTurns(root, adapter) - : codingTurn(scenario, root, adapter), + : scenario === 'coding-sandbox' + ? sandboxTurn(root, adapter) + : codingTurn(scenario, root, adapter), ) }, }, diff --git a/testing/e2e/tests/tools-test/coding-tools.spec.ts b/testing/e2e/tests/tools-test/coding-tools.spec.ts index c7016fa0d4..da0ccbcdb4 100644 --- a/testing/e2e/tests/tools-test/coding-tools.spec.ts +++ b/testing/e2e/tests/tools-test/coding-tools.spec.ts @@ -149,6 +149,20 @@ test.describe('Coding tools', () => { expect(await read('bash.txt')).toBe('AGENT=1') }) + test('edit_file works in a sandbox with sandboxWorkspaceBackend', async ({ + request, + testId, + }) => { + await writeFile(file('notes.txt'), 'alpha\nbeta\n') + + // The sandbox maps its /workspace to the test folder. + expect(await runScenario(request, 'coding-sandbox', testId)).toEqual({ + results: ['Edited notes.txt (1 change).'], + text: 'Edited notes.txt in the sandbox.', + }) + expect(await read('notes.txt')).toBe('sandboxed\nbeta\n') + }) + test('revert puts back the files of the turns and keeps the edit of the user', async ({ request, testId, From c75aefd9f624873effa9f6eb602a8ec78358e23e Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 16 Oct 2026 11:35:39 +0200 Subject: [PATCH 13/15] fix(ai-harness, ai-sandbox, ai-e2e): use session.send() after the core send() change --- docs/harness/coding-agent.md | 2 +- docs/harness/coding-agents.md | 2 +- docs/harness/coding-tools.md | 2 +- docs/harness/snapshots.md | 4 +- .../ai-harness/tests/coding-backend.test.ts | 2 +- .../ai-harness/tests/coding-formatter.test.ts | 12 ++--- .../ai-harness/tests/coding-snapshots.test.ts | 52 +++++++++---------- packages/ai-harness/tests/first-party.test.ts | 4 +- .../ai-harness/tests/instructions.test.ts | 2 +- .../ai-harness/tests/workspace-tools.test.ts | 24 ++++----- .../ai-sandbox/tests/coding-agents.test.ts | 16 +++--- .../tests/workspace-backend.test.ts | 2 +- testing/e2e/src/routes/api.coding-tools.ts | 10 ++-- 13 files changed, 66 insertions(+), 68 deletions(-) diff --git a/docs/harness/coding-agent.md b/docs/harness/coding-agent.md index 0effd1f088..a8bffe5f11 100644 --- a/docs/harness/coding-agent.md +++ b/docs/harness/coding-agent.md @@ -94,7 +94,7 @@ while (line !== 'exit') { const [name = '', ...args] = line.slice(1).split(' ') console.log(await session.command(name, args.join(' ') || undefined)) } else { - const turn = session.prompt(line) + const turn = session.send(line) turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { process.stdout.write(event.delta) }) diff --git a/docs/harness/coding-agents.md b/docs/harness/coding-agents.md index 994c0c52ca..125577e14c 100644 --- a/docs/harness/coding-agents.md +++ b/docs/harness/coding-agents.md @@ -83,7 +83,7 @@ import { memoryPersistence } from '@tanstack/ai-persistence' const host = createHarnessHost({ persistence: memoryPersistence() }) const session = await host.open(lead, { threadId: 'date-parser' }) -const turn = await session.prompt( +const turn = await session.send( 'Have claude_code add a test for the date parser, then have codex fix what fails.', ) console.log(turn.text) diff --git a/docs/harness/coding-tools.md b/docs/harness/coding-tools.md index ba3646f952..577d2a0a99 100644 --- a/docs/harness/coding-tools.md +++ b/docs/harness/coding-tools.md @@ -95,7 +95,7 @@ const stopAsking = session.onEvent('CUSTOM', ({ event }) => { // Send each line as a prompt. Print the answer. Type `exit` to stop. let line = await terminal.question('\n> ') while (line !== 'exit') { - const turn = session.prompt(line) + const turn = session.send(line) turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { process.stdout.write(event.delta) }) diff --git a/docs/harness/snapshots.md b/docs/harness/snapshots.md index 001f0b6883..cf8ab2425f 100644 --- a/docs/harness/snapshots.md +++ b/docs/harness/snapshots.md @@ -51,7 +51,7 @@ const session = await host.open(coder, { threadId: 'fix-login' }) ## 2. Undo and redo a turn -1. Send a prompt that changes files: `await session.prompt('Fix the login form.')`. +1. Send a prompt that changes files: `await session.send('Fix the login form.')`. 2. Run `await session.command('undo')`. It returns the number of files that it put back, for example `Undid the last turn. Files restored: 3.` 3. Look at the files that the turn changed: they are as they were before the turn. The turn is gone from the conversation too. 4. To bring back the files and the turn, run `await session.command('redo')`. It returns `Redid the last turn.` @@ -73,7 +73,7 @@ Each model step that changed files leaves a step: `{ from, to, files, toolCallId On a [durable host](./durable-sessions), like the host above, the steps are `tanstack/snapshots:step` records in the session log: ```ts group=harness-snapshots -await session.prompt('Fix the login form.') +await session.send('Fix the login form.') for (const { record } of await log.read('fix-login')) { if (record.type !== 'tanstack/snapshots:step') continue diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index a52e55e95d..f71916bbdc 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -172,7 +172,7 @@ describe('workspace hooks', () => { { threadId: 't' }, ) - await session.prompt('work') + await session.send('work') expect(written).toEqual(['one', 'two']) expect(JSON.stringify(calls[3].messages)).toContain( diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index 18f229337e..0b9f9e0664 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -296,7 +296,7 @@ describe('formatter', () => { { threadId: 't' }, ) - await session.prompt('work') + await session.send('work') await host.close() const command = `npx --no-install prettier --write ${quoted('a.ts')}` @@ -366,7 +366,7 @@ describe('formatter with permissions()', () => { 'prettier.config.js', ]) - const turn = session.prompt('change the config') + const turn = session.send('change the config') const message = await answer(session, { answer: 'reject' }) await turn @@ -386,7 +386,7 @@ describe('formatter with permissions()', () => { ['prettier.config.js', 'a.ts'], persistence, ) - const turn = first.session.prompt('write twice') + const turn = first.session.send('write twice') await answer(first.session, { answer: 'always' }) await turn expect(first.runs.map((run) => run.command)).toEqual([ @@ -399,7 +399,7 @@ describe('formatter with permissions()', () => { const second = await openWith('acceptEdits', ['b.ts'], persistence, { 'prettier.config.js': '', }) - await second.session.prompt('write again') + await second.session.send('write again') expect(second.runs.map((run) => run.command)).toEqual([ `npx --no-install prettier --write ${quoted('b.ts')}`, ]) @@ -410,7 +410,7 @@ describe('formatter with permissions()', () => { const { host, session, runs } = await openWith('bypass', [ 'prettier.config.js', ]) - await session.prompt('change the config') + await session.send('change the config') expect(runs.map((run) => run.command)).toEqual([ `npx --no-install prettier --write ${quoted('prettier.config.js')}`, ]) @@ -436,7 +436,7 @@ describe('formatter with permissions()', () => { }), { threadId: 't' }, ) - await session.prompt('change the config') + await session.send('change the config') expect(runs.map((run) => run.command)).toEqual([ `npx --no-install prettier --write ${quoted('prettier.config.js')}`, ]) diff --git a/packages/ai-harness/tests/coding-snapshots.test.ts b/packages/ai-harness/tests/coding-snapshots.test.ts index 9ec08635ca..9696515752 100644 --- a/packages/ai-harness/tests/coding-snapshots.test.ts +++ b/packages/ai-harness/tests/coding-snapshots.test.ts @@ -152,7 +152,7 @@ describe.skipIf(!hasGit)('snapshots', () => { }) const exec = vi.spyOn(hostBackend, 'exec') - await session.prompt('Edit the file') + await session.send('Edit the file') const commands = exec.mock.calls.map(([command]) => command) // `git --version` and `git init` once, the turn start, and one snapshot @@ -178,7 +178,7 @@ describe.skipIf(!hasGit)('snapshots', () => { ], }) - await session.prompt('Edit the files') + await session.send('Edit the files') const steps = savedSteps(session) expect(steps).toMatchObject([{ files: ['a.txt', 'src/b.txt'] }]) @@ -198,7 +198,7 @@ describe.skipIf(!hasGit)('snapshots', () => { persistence: { stores: { log, runs, metadata } }, }) - await session.prompt('Edit a.txt') + await session.send('Edit a.txt') const entries = await log.read(THREAD, { after: 0 }) const records = entries @@ -214,7 +214,7 @@ describe.skipIf(!hasGit)('snapshots', () => { replies: editTurn('Done.'), changes: [() => put('a.txt', 'two\n')], }) - await session.prompt('Edit a.txt') + await session.send('Edit a.txt') const [step] = savedSteps(session) const between = await plugin.diff(step.from, step.to) @@ -242,8 +242,8 @@ describe.skipIf(!hasGit)('snapshots', () => { }, ], }) - await session.prompt('Hello') - await session.prompt('Edit the files') + await session.send('Hello') + await session.send('Edit the files') expect(await session.command('undo')).toBe( 'Undid the last turn. Files restored: 3.', @@ -270,7 +270,7 @@ describe.skipIf(!hasGit)('snapshots', () => { }, ], }) - await session.prompt('Edit the files') + await session.send('Edit the files') const edited = await session.transcript() await session.command('undo') @@ -291,7 +291,7 @@ describe.skipIf(!hasGit)('snapshots', () => { replies: editTurn('Done.'), changes: [() => put('a.txt', 'two')], }) - await session.prompt('Edit a.txt') + await session.send('Edit a.txt') // The user adds and edits files after the turn. await put('b.txt', 'mine') await put('c.txt', 'mine') @@ -317,7 +317,7 @@ describe.skipIf(!hasGit)('snapshots', () => { replies: editTurn('Done.'), changes: [() => put('new/c.txt', 'added')], }) - await session.prompt('Add a file') + await session.send('Add a file') await session.command('undo') expect(await readdir(root)).toEqual(['a.txt']) @@ -341,7 +341,7 @@ describe.skipIf(!hasGit)('snapshots', () => { }, ], }) - await session.prompt('Edit the files') + await session.send('Edit the files') await put('x.txt', 'mine') // No file has the name `[x].txt` now. As a glob, it matches `x.txt`, so // git must read it as a plain name. @@ -362,9 +362,9 @@ describe.skipIf(!hasGit)('snapshots', () => { replies: [...editTurn('Done.'), () => text('Sure.')], changes: [() => put('a.txt', 'two')], }) - await session.prompt('Edit a.txt') + await session.send('Edit a.txt') await session.command('undo') - await session.prompt('Something else') + await session.send('Something else') expect(await session.command('redo')).toBe('Nothing to redo.') expect(await read('a.txt')).toBe('one') @@ -386,7 +386,7 @@ describe.skipIf(!hasGit)('snapshots', () => { changes: [() => put('b.txt', 'new')], }) - await session.prompt('Add b.txt') + await session.send('Add b.txt') await session.command('undo') expect(savedSteps(session)).toMatchObject([{ files: ['b.txt'] }]) @@ -412,9 +412,9 @@ const threeTurns = () => ({ /** Run the three turns of {@link threeTurns}. Gives the id of the first answer. */ async function runThreeTurns(session: HarnessSession) { - await session.prompt('One') - await session.prompt('Two') - await session.prompt('Three') + await session.send('One') + await session.send('Two') + await session.send('Three') const transcript = await session.transcript() const first = transcript.find((message) => message.content === 'First.') if (!first?.id) throw new Error('The first answer has no id.') @@ -460,7 +460,7 @@ describe.skipIf(!hasGit)('session.revert', () => { const { transcript, firstId } = await runThreeTurns(session) await session.revert(firstId) - await session.prompt('Four') + await session.send('Four') expect(messageTexts(calls.at(-1))).toEqual([ 'One', @@ -486,9 +486,9 @@ describe.skipIf(!hasGit)('session.revert', () => { const { host, session } = await open({ replies: [() => text('Hi.'), after(wait.opened, 'Later.')], }) - await session.prompt('Hello') + await session.send('Hello') const [message] = await session.transcript() - const running = session.prompt('Again') + const running = session.send('Again') await vi.waitFor(() => expect(session.snapshot().status).toBe('running')) expect(await session.revert(message?.id ?? '')).toMatchObject({ @@ -545,11 +545,11 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { await put('a.txt', 'one') await put('c.txt', 'one') const { host, session } = await open(threeTurns()) - await session.prompt('One') - await session.prompt('Two') + await session.send('One') + await session.send('Two') // The user edits a file that no turn changes. await put('c.txt', 'mine') - await session.prompt('Three') + await session.send('Three') const first = (await session.transcript()).find( (message) => message.content === 'First.', ) @@ -605,7 +605,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { const reverting = session.revert(firstId) await reached.opened - const prompting = session.prompt('Four') + const prompting = session.send('Four') await vi.waitFor(() => expect(session.snapshot().queuedTurns).toBe(1)) expect(session.snapshot().status).toBe('idle') blocking = false @@ -636,7 +636,7 @@ describe.skipIf(!hasGit)('session.revert keeps the edits of the user', () => { await session.revert(firstId) await session.configure({ cwd: 'sub' }) - await session.prompt('Four') + await session.send('Four') expect( (await session.transcript()).slice(4).map((message) => message.content), @@ -697,7 +697,7 @@ describe('snapshots on a cmd.exe backend', () => { backend, }) - await session.prompt('Hello') + await session.send('Hello') // `git --version`, then the init of the shadow repository. expect(commands[1]).toMatch(/^git init --bare -q \^".+\^"$/) @@ -718,7 +718,7 @@ describe('snapshots without git', () => { })), }) - await session.prompt('Hello') + await session.send('Hello') expect(session.commands()).toEqual([]) expect(await session.transcript()).toMatchObject([ diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index 12e8140bbc..e213680260 100644 --- a/packages/ai-harness/tests/first-party.test.ts +++ b/packages/ai-harness/tests/first-party.test.ts @@ -75,7 +75,7 @@ describe('workspace tools with permissions', () => { workspaceTools({ root }), ]) - const turn = session.prompt('work') + const turn = session.send('work') await vi.waitFor(() => expect(session.snapshot().pendingQuestions).toHaveLength(1), ) @@ -104,7 +104,7 @@ describe('workspace tools with permissions', () => { workspaceTools({ root }), ]) expect(await session.command('mode', 'plan')).toBe('Mode: plan.') - await session.prompt('write') + await session.send('write') await expect(readFile(join(root, 'x.txt'), 'utf8')).rejects.toThrow() expect(session.snapshot().pendingQuestions).toHaveLength(0) await host.close() diff --git a/packages/ai-harness/tests/instructions.test.ts b/packages/ai-harness/tests/instructions.test.ts index ad508f05c4..235b3871c0 100644 --- a/packages/ai-harness/tests/instructions.test.ts +++ b/packages/ai-harness/tests/instructions.test.ts @@ -250,7 +250,7 @@ describe('projectInstructions', () => { projectInstructions({ root, env: false }), ]) - await session.prompt('read') + await session.send('read') await host.close() // What each read added after the lines of the file. diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index 26bc868246..e572b04a74 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -472,7 +472,7 @@ describe('editStyle', () => { editStyle, }), ]) - await session.prompt('hi') + await session.send('hi') expect( toolNames(calls[0]).filter((name) => EDIT_TOOLS.includes(name)), ).toEqual(tools) @@ -555,7 +555,7 @@ describe('permission resources', () => { }), workspaceTools({ root: '/workspace', backend, editStyle: 'patch' }), ]) - await session.prompt('work') + await session.send('work') expect(session.snapshot().pendingQuestions).toHaveLength(0) expect([...files.keys()]).toEqual(['/workspace/src/a.ts']) await host.close() @@ -645,7 +645,7 @@ describe('permission rules, links, and grep', () => { const { host, session, seen } = await callTools(backend, [ ['read_file', { path: 'notes' }], ]) - const turn = session.prompt('read the notes') + const turn = session.send('read the notes') await vi.waitFor(() => expect(session.snapshot().pendingQuestions).toHaveLength(1), ) @@ -672,7 +672,7 @@ describe('permission rules, links, and grep', () => { ], }, ) - await session.prompt('read the key') + await session.send('read the key') expect(session.snapshot().pendingQuestions).toHaveLength(0) expect(result('c0')).toContain('The permission rules do not allow') expect(seen()).not.toContain('hunter2') @@ -686,7 +686,7 @@ describe('permission rules, links, and grep', () => { [['read_file', { path: 'notes' }]], { mode: 'plan' }, ) - await session.prompt('read the notes') + await session.send('read the notes') expect(session.snapshot().pendingQuestions).toHaveLength(0) expect(seen()).not.toContain('hunter2') await host.close() @@ -714,7 +714,7 @@ describe('permission rules, links, and grep', () => { mode, }, ) - await session.prompt('search') + await session.send('search') expect(session.snapshot().pendingQuestions).toHaveLength(0) for (const id of ['c0', 'c1']) { expect(result(id)).toBe(`"src/a.ts:1: KEY=1\\n${SKIPPED_TWO}"`) @@ -739,7 +739,7 @@ describe('permission rules, links, and grep', () => { ], }, ) - await session.prompt('search') + await session.send('search') expect(result('c0')).toBe( '"No matches.\\n[Skipped 1 file that the permission rules protect.]"', ) @@ -775,7 +775,7 @@ describe('webfetch approval', () => { it('asks before webfetch by default', async () => { const { host, session } = await fetchOnce() - const turn = session.prompt('read the page') + const turn = session.send('read the page') await vi.waitFor(() => expect(session.snapshot().pendingQuestions).toHaveLength(1), ) @@ -794,7 +794,7 @@ describe('webfetch approval', () => { const { host, session, calls } = await fetchOnce([ { tool: 'webfetch', decision: 'allow' }, ]) - await session.prompt('read the page') + await session.send('read the page') expect(session.snapshot().pendingQuestions).toHaveLength(0) expect(JSON.stringify(calls[1].messages)).toContain('fetched body') await host.close() @@ -818,7 +818,7 @@ describe('background bash in a session', () => { const { host, session } = await openSession(adapter, [ workspaceTools({ root: '/workspace', backend }), ]) - await session.prompt('build') + await session.send('build') await vi.waitFor(() => expect(calls).toHaveLength(3)) expect(JSON.stringify(calls[2].messages)).toContain( 'Background job bash-1 ended.\\nexit code: 0\\nbuilt', @@ -851,7 +851,7 @@ describe('background bash in a session', () => { const { host, session } = await openSession(adapter, [ workspaceTools({ root: '/workspace', backend }), ]) - await session.prompt('serve') + await session.send('serve') expect(killed).toBe(false) await host.close() expect(killed).toBe(true) @@ -948,7 +948,7 @@ describe('links on the host', () => { }), workspaceTools({ root, backend: hostBackend, web: false }), ]) - await session.prompt('read through the link') + await session.send('read through the link') const seen = JSON.stringify(calls[1].messages) expect(seen).toContain( 'The permission rules do not allow read_file on src/a.ts, the real path of inner/a.ts.', diff --git a/packages/ai-sandbox/tests/coding-agents.test.ts b/packages/ai-sandbox/tests/coding-agents.test.ts index da467217fb..a799a4167a 100644 --- a/packages/ai-sandbox/tests/coding-agents.test.ts +++ b/packages/ai-sandbox/tests/coding-agents.test.ts @@ -209,8 +209,8 @@ describe('codingAgents', () => { ], agents: { claude_code: { adapter: claude.adapter } }, }) - await session.prompt('first') - await session.prompt('second') + await session.send('first') + await session.send('second') expect(claude.calls).toHaveLength(2) expect(claude.calls[0].messages).toEqual([ @@ -247,15 +247,15 @@ describe('codingAgents', () => { 'claude_code', 'codex', ]) - await session.prompt('a') - await session.prompt('b') + await session.send('a') + await session.send('b') expect(await session.command('fresh', 'nope')).toBe( 'Unknown coding agent "nope". Agents: claude_code, codex.', ) expect(await session.command('fresh', 'claude_code')).toBe( 'claude_code starts a new session next time.', ) - await session.prompt('c') + await session.send('c') expect(claude.calls[1].modelOptions?.sessionId).toBeUndefined() expect(await session.command('fresh')).toBe( 'Every coding agent starts a new session next time.', @@ -271,7 +271,7 @@ describe('codingAgents', () => { agents: { claude_code: { adapter: first.adapter } }, persistence, }) - await before.session.prompt('go') + await before.session.send('go') await before.host.close() const second = codingAgent('claude-code') @@ -280,7 +280,7 @@ describe('codingAgents', () => { agents: { claude_code: { adapter: second.adapter } }, persistence, }) - await after.session.prompt('go on') + await after.session.send('go on') expect(second.calls[0].modelOptions).toEqual({ sessionId: 'claude-code-session-1', }) @@ -368,7 +368,7 @@ describe('codingAgents', () => { }, }) await session.setConfig({ key: 'mode', value: 'plan' }) - await session.prompt('plan it') + await session.send('plan it') expect(claude.calls[0].modelOptions).toMatchObject({ permissionMode: 'plan', }) diff --git a/packages/ai-sandbox/tests/workspace-backend.test.ts b/packages/ai-sandbox/tests/workspace-backend.test.ts index 78958d92dd..cda25d693d 100644 --- a/packages/ai-sandbox/tests/workspace-backend.test.ts +++ b/packages/ai-sandbox/tests/workspace-backend.test.ts @@ -439,7 +439,7 @@ describe('sandboxWorkspaceBackend with workspaceTools', () => { }), { threadId: 't' }, ) - await session.prompt('work') + await session.send('work') await host.close() expect(decoder.decode(sandbox.files.get('/workspace/src/a.ts'))).toBe( diff --git a/testing/e2e/src/routes/api.coding-tools.ts b/testing/e2e/src/routes/api.coding-tools.ts index bb75d5c5c9..c0039be999 100644 --- a/testing/e2e/src/routes/api.coding-tools.ts +++ b/testing/e2e/src/routes/api.coding-tools.ts @@ -50,9 +50,7 @@ async function* codingTurn( { threadId: scenario }, ) if (formatAsk) await session.command('mode', 'acceptEdits') - for await (const chunk of session - .prompt(`[${scenario}] run test`) - .stream()) { + for await (const chunk of session.send(`[${scenario}] run test`).stream()) { yield chunk const asks = chunk.type === EventType.CUSTOM && @@ -94,9 +92,9 @@ async function* revertTurns( }), { threadId: 'coding-revert' }, ) - yield* session.prompt('[coding-revert] run test').stream() + yield* session.send('[coding-revert] run test').stream() await writeFile(join(root, 'other.txt'), 'mine\n') - yield* session.prompt('[coding-revert-2] run test').stream() + yield* session.send('[coding-revert-2] run test').stream() const [prompt] = await session.transcript() if (prompt?.id) await session.revert(prompt.id) } finally { @@ -131,7 +129,7 @@ async function* sandboxTurn( }), { threadId: 'coding-sandbox' }, ) - yield* session.prompt('[coding-sandbox] run test').stream() + yield* session.send('[coding-sandbox] run test').stream() } finally { await host.close() // With `dir`, destroy keeps the folder. From e9ff04564b6161c230ba9d1a4cf891da5a382136 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 16 Oct 2026 11:41:41 +0200 Subject: [PATCH 14/15] fix(ai-harness, ai-sandbox): review fixes for the coding tools - A background job keeps the last 10 MiB of its output, not all of it. The host backend decodes UTF-8 on the stream, so a character split over two chunks stays whole. - A Windows network path (\server\share) goes through the boundary check before realpath, which connects to the server. - patch refuses an Add File on a folder before it writes any file. - The formatter checks the project again after a failed check, so one failed readdir does not stop formatting for the session. --- .../src/first-party/coding/backend.ts | 15 +++++++---- .../src/first-party/coding/formatter.ts | 6 ++++- .../src/first-party/coding/patch.ts | 4 +++ .../src/first-party/coding/workspace.ts | 8 ++++++ .../ai-harness/tests/coding-formatter.test.ts | 19 ++++++++++++++ .../ai-harness/tests/coding-patch.test.ts | 25 +++++++++++++++++-- .../ai-harness/tests/workspace-tools.test.ts | 20 +++++++++++++++ packages/ai-sandbox/src/workspace-backend.ts | 10 ++++++-- 8 files changed, 97 insertions(+), 10 deletions(-) diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index 0ccb42321e..f316f0797a 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -298,13 +298,14 @@ export const hostBackend = { }), spawn: (command, options = {}) => { const { child, kill } = startShell(command, options) - // ponytail: keeps all output in memory. Cap it if long jobs need that. + // ponytail: keeps the last 10 MiB of a long job, like `exec` keeps the first. let output = '' - const collect = (chunk: Buffer) => { - output += chunk.toString() + const collect = (chunk: string) => { + output = keepTail(output + chunk) } - child.stdout.on('data', collect) - child.stderr.on('data', collect) + // The stream decodes UTF-8, so a character split over two chunks stays whole. + child.stdout.setEncoding('utf8').on('data', collect) + child.stderr.setEncoding('utf8').on('data', collect) const exited = new Promise<{ exitCode: number }>((done) => { child.on('close', (code) => done({ exitCode: code ?? 1 })) child.on('error', (error) => { @@ -316,6 +317,10 @@ export const hostBackend = { }, } satisfies WorkspaceBackend +/** Keep the last `MAX_EXEC_OUTPUT` characters. Cuts at twice that, so it copies rarely. */ +const keepTail = (text: string) => + text.length > 2 * MAX_EXEC_OUTPUT ? text.slice(-MAX_EXEC_OUTPUT) : text + const MAX_OUTPUT = 20_000 /** Cut text after 20,000 characters, and say how many were cut. */ diff --git a/packages/ai-harness/src/first-party/coding/formatter.ts b/packages/ai-harness/src/first-party/coding/formatter.ts index c7b3aafc75..ad46a78431 100644 --- a/packages/ai-harness/src/first-party/coding/formatter.ts +++ b/packages/ai-harness/src/first-party/coding/formatter.ts @@ -175,7 +175,11 @@ export function formatOnWrite( const format = async (path: string) => { const extension = extname(path).toLowerCase() - const formatters = await (detected ??= detect()) + // A failed detection is not kept, so the next write tries again. + const formatters = await (detected ??= detect().catch((error: unknown) => { + detected = undefined + throw error + })) const match = formatters.find((candidate) => candidate.extensions.includes(extension), ) diff --git a/packages/ai-harness/src/first-party/coding/patch.ts b/packages/ai-harness/src/first-party/coding/patch.ts index c617459f14..8caed0d643 100644 --- a/packages/ai-harness/src/first-party/coding/patch.ts +++ b/packages/ai-harness/src/first-party/coding/patch.ts @@ -319,6 +319,10 @@ async function applyAll(env: ToolEnv, steps: ReadonlyArray) { const shown = env.shown(full) switch (operation.type) { case 'add': + // Checked here, so a patch that does not fit writes nothing. + if ((await backend.stat(full))?.type === 'dir') { + throw new Error(`Cannot add ${shown}: it is a folder.`) + } planned.set(full, operation.content) changes.push(`added ${shown}`) break diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index 8f2f62d3ed..cb810cee5f 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -75,6 +75,9 @@ export interface OutsideAccess { | undefined } +/** A Windows network path (`\\server\share`). Not `\\?\` or `\\.\`. */ +const isUnc = (path: string) => /^[\\/]{2}[^\\/?.]/.test(path) + /** Is `full` the folder `dir`, or a path in it? `sep` is the separator. */ function within(dir: string, full: string, sep: string) { return full === dir || full.startsWith(dir.endsWith(sep) ? dir : dir + sep) @@ -269,6 +272,11 @@ export function createWorkspaceTools( // leads out of it counts as outside. const reach: ToolEnv['reach'] = async (path, tool, kind) => { const full = paths.resolve(here(), path) + // `realpath` on a UNC path connects to its server, and Windows can send + // the user's credentials. So the boundary check runs first. + if (isUnc(full) && !isUnc(root)) { + await allowOutside(path, tool, kind, full, full) + } const real = await realPath(full) const folders = [await realPath(root), ...allowed] const isAllowed = folders.some((dir) => within(dir, real, paths.sep)) diff --git a/packages/ai-harness/tests/coding-formatter.test.ts b/packages/ai-harness/tests/coding-formatter.test.ts index 0b9f9e0664..6dbe9baa9b 100644 --- a/packages/ai-harness/tests/coding-formatter.test.ts +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -216,6 +216,25 @@ describe('formatOnWrite', () => { expect(commands).toEqual([]) }) + it('checks the project again at the next write when the check failed', async () => { + const { backend, runs } = project(prettier) + const readdir = backend.readdir + let fail = true + backend.readdir = async (dir) => { + if (fail) throw new Error('EBUSY') + return readdir(dir) + } + const failures: Array = [] + const afterWrite = formatOnWrite({ root, backend }, (failure) => + failures.push(failure), + ) + await afterWrite(join(root, 'a.ts')) + fail = false + await afterWrite(join(root, 'a.ts')) + expect(failures).toEqual([{ path: join(root, 'a.ts'), message: 'EBUSY' }]) + expect(runs).toHaveLength(1) + }) + it('checks the config files once, at the first write', async () => { const { backend, runs, put } = project({}) const afterWrite = formatOnWrite({ root, backend }, () => undefined) diff --git a/packages/ai-harness/tests/coding-patch.test.ts b/packages/ai-harness/tests/coding-patch.test.ts index 9ff5a176aa..0b8a373cce 100644 --- a/packages/ai-harness/tests/coding-patch.test.ts +++ b/packages/ai-harness/tests/coding-patch.test.ts @@ -28,9 +28,12 @@ function patch(...lines: Array) { /** * The patch tool on files in memory that start as `files`. `written` lists * the files that the `afterWrite` hook saw, in order. With `canRemove: - * false`, the backend has no `remove`. + * false`, the backend has no `remove`. `dirs` are folders. */ -function workspace(files: Record, { canRemove = true } = {}) { +function workspace( + files: Record, + { canRemove = true, dirs = [] as Array } = {}, +) { const store = new Map() for (const [path, text] of Object.entries(files)) { store.set(join(root, path), encoder.encode(text)) @@ -50,6 +53,9 @@ function workspace(files: Record, { canRemove = true } = {}) { } : undefined, stat: async (path) => { + if (dirs.some((dir) => join(root, dir) === path)) { + return { type: 'dir', size: 0, mtimeMs: 0 } + } const data = store.get(path) return data && { type: 'file', size: data.length, mtimeMs: 0 } }, @@ -392,6 +398,21 @@ describe('patch tool', () => { expect(files.written).toEqual([]) }) + it('changes no file when an added file is a folder', async () => { + const files = workspace({ 'a.txt': 'a\n' }, { dirs: ['src'] }) + const text = patch( + '*** Update File: a.txt', + '-a', + '+A', + '*** Add File: src', + '+c', + ) + await expect(files.run(text)).rejects.toThrow( + 'Cannot add src: it is a folder.', + ) + expect(files.files()).toEqual({ 'a.txt': 'a\n' }) + }) + it('runs the afterWrite hook once for each written file', async () => { const files = workspace({ 'a.txt': 'a\n', 'b.txt': 'b\n' }) await files.run( diff --git a/packages/ai-harness/tests/workspace-tools.test.ts b/packages/ai-harness/tests/workspace-tools.test.ts index e572b04a74..1ea8f8c01a 100644 --- a/packages/ai-harness/tests/workspace-tools.test.ts +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -1038,3 +1038,23 @@ describe('bash', () => { }, ) }) + +describe('a network path', () => { + it('is refused before realpath can connect to its server', async () => { + const realpath = vi.fn(async (path: string) => path) + const backend: WorkspaceBackend = { + shell: 'cmd', + realpath, + readFile: async () => new Uint8Array(), + writeFile: async () => undefined, + stat: async () => undefined, + readdir: async () => [], + exec: async () => ({ exitCode: 0, stdout: '', stderr: '' }), + } + const tools = toolsOf({ root: 'C:/work', backend }) + await expect( + runIn(tools, 'read_file', { path: '//attacker/share/a.txt' }), + ).rejects.toThrow('outside the workspace') + expect(realpath.mock.calls.flat().join()).not.toContain('attacker') + }) +}) diff --git a/packages/ai-sandbox/src/workspace-backend.ts b/packages/ai-sandbox/src/workspace-backend.ts index b52ef87841..b5e3c51665 100644 --- a/packages/ai-sandbox/src/workspace-backend.ts +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -1,6 +1,9 @@ import type { WorkspaceBackend } from '@tanstack/ai-harness/plugins/coding' import type { ExecResult, SandboxHandle } from './contracts' +/** The most output a background job keeps, in characters. */ +const MAX_OUTPUT = 10 * 1024 * 1024 + /** * A `WorkspaceBackend` over a live sandbox. With it, `workspaceTools()` * reads and writes files and runs commands in the sandbox, not on this @@ -59,10 +62,13 @@ export function sandboxWorkspaceBackend(handle: SandboxHandle) { } const spawn: WorkspaceBackend['spawn'] = (command, options) => { - // ponytail: keeps all output in memory, like `hostBackend`. + // ponytail: keeps the last 10 MiB of a long job, like `hostBackend`. let output = '' const collect = async (stream: AsyncIterable) => { - for await (const chunk of stream) output += chunk + for await (const chunk of stream) { + output += chunk + if (output.length > 2 * MAX_OUTPUT) output = output.slice(-MAX_OUTPUT) + } } const started = handle.process.spawn(command, options) // Wait for the end of the output too, so `output()` is complete. From d4287887cf87bc83fe85d03f75c833ea1d6d5d03 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Fri, 16 Oct 2026 12:20:59 +0200 Subject: [PATCH 15/15] perf(ai-harness): faster file reads, edits, and searches in the coding tools - edit_file skips a loose match level that changes neither the file nor the old text. A miss in a 5 MB file takes about 12 ms, not 1 s. - hostBackend.readLines reads one page of a file. It counts lines on the bytes and decodes only the page. read_file uses it when the backend has it, and reads the whole file when it does not. - Without rg, list_files and grep read up to 16 folders or files at a time, and keep the order. - The real path of the root is read once, the permission check reuses the real path of the call, read_file needs no stat when the file is there, and writeFile makes folders only when one is missing. --- docs/config.json | 2 +- docs/harness/coding-tools-backends.md | 1 + .../src/first-party/coding/backend.ts | 93 +++++++++++++++- .../ai-harness/src/first-party/coding/edit.ts | 9 +- .../ai-harness/src/first-party/coding/read.ts | 66 +++++++----- .../src/first-party/coding/search.ts | 101 ++++++++++++------ .../src/first-party/coding/workspace.ts | 41 ++++--- .../ai-harness/tests/coding-backend.test.ts | 42 ++++++++ packages/ai-harness/tests/coding-edit.test.ts | 12 +++ .../ai-harness/tests/coding-search.test.ts | 31 ++++++ 10 files changed, 327 insertions(+), 71 deletions(-) diff --git a/docs/config.json b/docs/config.json index 2c5a840167..98cb63da0c 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1040,7 +1040,7 @@ "label": "Sandbox and extend the coding tools", "to": "harness/coding-tools-backends", "addedAt": "2026-10-13", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-16" }, { "label": "Undo a turn", diff --git a/docs/harness/coding-tools-backends.md b/docs/harness/coding-tools-backends.md index 86718fdda2..1641fe93e3 100644 --- a/docs/harness/coding-tools-backends.md +++ b/docs/harness/coding-tools-backends.md @@ -94,6 +94,7 @@ The backend can also have: | --- | --- | --- | | `shell` | `'sh'` or `'cmd'`. Sets the path style and how arguments are quoted. | `sh` quotes, and the path style of this machine. | | `remove(path)` | Removes one file. | `patch` cannot delete or move files. | +| `readLines(path, { offset, limit })` | Gives one page of a file: `{ head, size, lines, total }`. `head` is the first 8 KiB. | `read_file` reads the whole file for each page. | | `realpath(path)` | Gives the path with every link resolved. | The tools do not check links. | | `spawn(command, options)` | Starts a command in the background. | `background: true` gives an error. | diff --git a/packages/ai-harness/src/first-party/coding/backend.ts b/packages/ai-harness/src/first-party/coding/backend.ts index f316f0797a..545ab72801 100644 --- a/packages/ai-harness/src/first-party/coding/backend.ts +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -21,6 +21,24 @@ export interface WorkspaceBackend { shell?: 'sh' | 'cmd' /** The bytes of a file. Throws when there is no file. */ readFile: (path: string) => Promise + /** + * One page of a file, without the whole file in memory: the first 8 KiB + * (`head`), the size in bytes, the lines from line `offset` (from 1) as + * UTF-8 text, and the number of lines. It gives at most `limit` lines, and + * fewer when the lines are long: no more lines after 64 KiB. A final + * newline ends the last line. It does not start a new one. Throws when + * there is no file. Optional: without it, `read_file` reads the whole + * file. + */ + readLines?: ( + path: string, + options: { offset: number; limit: number }, + ) => Promise<{ + head: Uint8Array + size: number + lines: Array + total: number + }> /** Create or replace a file. Missing parent folders are created. */ writeFile: (path: string, data: Uint8Array | string) => Promise /** @@ -143,6 +161,15 @@ export function commandEnv( const isMissing = (error: unknown) => error instanceof Error && 'code' in error && error.code === 'ENOENT' +/** `readLines` reads this many bytes at a time. */ +const READ_CHUNK = 1024 * 1024 +/** The first bytes that `readLines` gives, for the file type check. */ +const HEAD_BYTES = 8 * 1024 +/** `readLines` gives no more lines after this many bytes. */ +const MAX_PAGE_BYTES = 64 * 1024 +// Drops a byte order mark, like the whole-file read of `read_file`. +const pageDecoder = new TextDecoder() + /** `exec` stops a command that writes more than 10 MiB of output. */ const MAX_EXEC_OUTPUT = 10 * 1024 * 1024 @@ -209,9 +236,71 @@ function startShell( export const hostBackend = { shell: process.platform === 'win32' ? 'cmd' : 'sh', readFile: (path) => fs.readFile(path), + readLines: async (path, { offset, limit }) => { + const handle = await fs.open(path, 'r') + try { + const { size } = await handle.stat() + const chunk = Buffer.allocUnsafe(Math.max(1, Math.min(size, READ_CHUNK))) + let head: Uint8Array | undefined + // The line of the next byte. At the end, one more than the newlines. + let line = 1 + let last = -1 + // The first line after the page. A long page ends early. + let end = offset + limit + const page: Array = [] + let pageBytes = 0 + // ponytail: reads up to the size at the start. Bytes added later wait + // for the next read. + for (let position = 0; position < size;) { + const { bytesRead } = await handle.read( + chunk, + 0, + chunk.length, + position, + ) + if (bytesRead === 0) break + position += bytesRead + const data = chunk.subarray(0, bytesRead) + // A copy: the next read writes over `chunk`. + head ??= Buffer.from(data.subarray(0, HEAD_BYTES)) + last = data[bytesRead - 1] ?? last + // The page bytes of this read. The page lines are one after another. + let pageFrom = -1 + let pageTo = -1 + // Count the lines with `indexOf`. Only the page is decoded. + for (let from = 0; from < bytesRead;) { + const newline = data.indexOf(10, from) + const to = newline === -1 ? bytesRead : newline + 1 + if (line >= offset && line < end) { + if (pageFrom === -1) pageFrom = from + pageTo = to + pageBytes += to - from + if (newline !== -1 && pageBytes > MAX_PAGE_BYTES) end = line + 1 + } + if (newline === -1) break + line += 1 + from = to + } + if (pageFrom !== -1) + page.push(Buffer.from(data.subarray(pageFrom, pageTo))) + } + // An empty file has one empty line, like `''.split('\n')`. + const total = line - 1 + (size === 0 || last !== 10 ? 1 : 0) + const text = pageDecoder.decode(Buffer.concat(page)) + const lines = offset > total ? [] : text.split('\n') + if (text.endsWith('\n')) lines.pop() + return { head: head ?? Buffer.alloc(0), size, lines, total } + } finally { + await handle.close() + } + }, writeFile: async (path, data) => { - await fs.mkdir(nodePath.dirname(path), { recursive: true }) - await fs.writeFile(path, data) + // One call when the folder is there. When it is not, make it and write again. + await fs.writeFile(path, data).catch(async (error: unknown) => { + if (!isMissing(error)) throw error + await fs.mkdir(nodePath.dirname(path), { recursive: true }) + await fs.writeFile(path, data) + }) }, // Without `recursive`, `fs.rm` refuses a folder. remove: (path) => fs.rm(path), diff --git a/packages/ai-harness/src/first-party/coding/edit.ts b/packages/ai-harness/src/first-party/coding/edit.ts index 1dba9f5180..fdb0da13b9 100644 --- a/packages/ai-harness/src/first-party/coding/edit.ts +++ b/packages/ai-harness/src/first-party/coding/edit.ts @@ -7,11 +7,13 @@ const BOM = '' /** * How one match level changes text before the search. `pieces` cuts the text - * into pieces, and `change` gives the new text of each piece. + * into pieces, and `change` gives the new text of each piece. `changes` is a + * fast test: false when the level leaves the text as it is. */ interface Level { pieces: RegExp change: (piece: RegExpExecArray) => string + changes: (text: string) => boolean } // The match levels after the exact text, from strict to loose. @@ -21,11 +23,13 @@ const LEVELS: ReadonlyArray = [ { pieces: /\P{M}\p{M}*|\p{M}+/gu, change: (piece) => piece[0].normalize('NFC'), + changes: (text) => text.normalize('NFC') !== text, }, // Spaces and tabs at the end of each line removed. { pieces: /([^\S\n]+(?=\n|$))|[\s\S]/g, change: (piece) => (piece[1] ? '' : piece[0]), + changes: (text) => /[^\S\n](?=\n|$)/.test(text), }, ] @@ -96,6 +100,9 @@ export function findMatches(content: string, find: string) { return exact.map((start) => ({ start, end: start + wanted.length })) } for (const level of LEVELS) { + // A level that changes neither side finds only the exact matches, and + // there are none. On a big file the change is slow, so skip it. + if (!level.changes(content) && !level.changes(wanted)) continue const { text, at } = changed(content, level) const needle = changed(wanted, level).text const starts = indexesOf(text, needle) diff --git a/packages/ai-harness/src/first-party/coding/read.ts b/packages/ai-harness/src/first-party/coding/read.ts index 7aaa315901..8bc9e392fa 100644 --- a/packages/ai-harness/src/first-party/coding/read.ts +++ b/packages/ai-harness/src/first-party/coding/read.ts @@ -93,22 +93,25 @@ function numbered(lineNumber: number, line: string) { } /** - * The numbered lines from line `offset`, at most `limit` lines (2000 when - * `limit` is not set) and 50 KiB. When the page shows fewer lines than were - * asked for and the file goes on, a note says how to read the next page. + * The lines of `text` from line `offset`, at most `limit`, and the number of + * lines. A final newline ends the last line. It does not start a new one. + */ +function linesOf(text: string, offset: number, limit: number) { + const all = text.split('\n') + if (text.endsWith('\n')) all.pop() + return { lines: all.slice(offset - 1, offset - 1 + limit), total: all.length } +} + +/** + * The numbered lines of `window`, which starts at line `offset` of a file of + * `total` lines. At most 50 KiB. When the page shows fewer lines than + * `limit` and the file goes on, a note says how to read the next page. */ function page( - text: string, + { lines: window, total }: { lines: Array; total: number }, options: { path: string; offset: number; limit: number | undefined }, ) { const { path, offset, limit } = options - const lines = text.split('\n') - // A final newline ends the last line. It does not start a new one. - if (text.endsWith('\n')) lines.pop() - const window = lines.slice( - offset - 1, - offset - 1 + Math.min(limit ?? MAX_LINES, MAX_LINES), - ) const shown: Array = [] let bytes = 0 for (const [index, line] of window.entries()) { @@ -119,12 +122,12 @@ function page( shown.push(row) } if (shown.length === 0) { - return `[No lines from line ${offset}. ${path} has ${lines.length} lines.]` + return `[No lines from line ${offset}. ${path} has ${total} lines.]` } const last = offset + shown.length - 1 - const isCut = shown.length < (limit ?? Infinity) && last < lines.length + const isCut = shown.length < (limit ?? Infinity) && last < total if (!isCut) return shown.join('\n') - return `${shown.join('\n')}\n\n[Showing lines ${offset}-${last} of ${lines.length}. Read more with offset ${last + 1}.]` + return `${shown.join('\n')}\n\n[Showing lines ${offset}-${last} of ${total}. Read more with offset ${last + 1}.]` } /** @@ -159,19 +162,34 @@ export function readTools(env: ToolEnv) { isRecord(args) && typeof args.limit === 'number' ? args.limit : undefined + const pageLimit = Math.min(limit ?? MAX_LINES, MAX_LINES) return env.lock(full, async () => { const path = env.shown(full) - if (!(await env.backend.stat(full))) throw await notFound(env, full) - const bytes = await env.backend.readFile(full) - const head = String.fromCharCode(...bytes.subarray(0, 12)) - const media = MEDIA.find((item) => item.signature.test(head)) + const { backend } = env + // One call when the file is there. A failed read checks if it is. + const missing = async (error: unknown): Promise => { + if (!(await backend.stat(full))) throw await notFound(env, full) + throw error + } + // `readLines` reads one page. Without it, the whole file is read. + const paged = await backend + .readLines?.(full, { offset, limit: pageLimit }) + .catch(missing) + const whole = paged + ? undefined + : await backend.readFile(full).catch(missing) + const head = paged?.head ?? whole?.subarray(0, BINARY_CHECK_BYTES) + const size = paged?.size ?? whole?.length ?? 0 + const start = String.fromCharCode(...(head ?? []).slice(0, 12)) + const media = MEDIA.find((item) => item.signature.test(start)) if (media) { - if (bytes.length > MAX_MEDIA_BYTES) { - const size = (bytes.length / 1024 / 1024).toFixed(1) + if (size > MAX_MEDIA_BYTES) { + const mb = (size / 1024 / 1024).toFixed(1) throw new Error( - `${path} is ${size} MB. read_file sends images and PDFs up to 5 MB.`, + `${path} is ${mb} MB. read_file sends images and PDFs up to 5 MB.`, ) } + const bytes = whole ?? (await backend.readFile(full)) const part: ContentPart = { type: media.type, source: { @@ -182,12 +200,12 @@ export function readTools(env: ToolEnv) { } return [part] } - if (bytes.subarray(0, BINARY_CHECK_BYTES).includes(0)) { + if (head?.includes(0)) { const type = mimeTypeOf(full) const typeNote = type === undefined ? '' : `, ${type}` - return `${path} is a binary file (${bytes.length} bytes${typeNote}). read_file shows text, images, and PDFs only.` + return `${path} is a binary file (${size} bytes${typeNote}). read_file shows text, images, and PDFs only.` } - const text = decoder.decode(bytes) + const text = paged ?? linesOf(decoder.decode(whole), offset, pageLimit) const parts = [page(text, { path, offset, limit })] const hooks = env.hooks() for (const hook of hooks) { diff --git a/packages/ai-harness/src/first-party/coding/search.ts b/packages/ai-harness/src/first-party/coding/search.ts index 4a79103781..210ac24a87 100644 --- a/packages/ai-harness/src/first-party/coding/search.ts +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -23,6 +23,8 @@ const MAX_SIZE = 1024 * 1024 // ponytail: the JS walk stops after 5000 files, and the tools say so. Raise // it if folders without rg and without git get that big. const MAX_WALK = 5000 +/** The JS walk and grep run this many backend calls at a time. */ +const PARALLEL = 16 const FILES_NOTE = '[Not all files are shown. Use a pattern or a folder to see fewer.]' @@ -164,25 +166,40 @@ function rgHits(out: string) { } /** - * Add the files under `dir` to `out`, as paths that start with `prefix`. - * Stops after MAX_WALK files. It skips links, so it never leaves `dir`. + * The files under `dir`, as paths from `dir` with `/`. Stops after MAX_WALK + * files, and then gives one more, so the caller can tell. It skips links, so + * it never leaves `dir`. It reads up to PARALLEL folders at a time. */ -async function walk( - backend: WorkspaceBackend, - dir: string, - prefix: string, - out: Array, -) { - const entries = await backend.readdir(dir) - for (const entry of entries) { - if (out.length > MAX_WALK) return - if (SKIP.has(entry.name) || entry.type === 'link') continue - const path = prefix + entry.name - if (entry.type === 'dir') { - const next = pathsOf(backend).join(dir, entry.name) - await walk(backend, next, `${path}/`, out) - } else out.push(path) +async function walk(backend: WorkspaceBackend, dir: string) { + const { join } = pathsOf(backend) + const out: Array = [] + let folders = [{ dir, prefix: '' }] + while (folders.length > 0 && out.length <= MAX_WALK) { + const next: typeof folders = [] + for ( + let i = 0; + i < folders.length && out.length <= MAX_WALK; + i += PARALLEL + ) { + const listed = await Promise.all( + folders.slice(i, i + PARALLEL).map(async (folder) => ({ + folder, + entries: await backend.readdir(folder.dir), + })), + ) + for (const { folder, entries } of listed) { + for (const entry of entries) { + if (SKIP.has(entry.name) || entry.type === 'link') continue + const path = folder.prefix + entry.name + if (entry.type === 'dir') { + next.push({ dir: join(folder.dir, entry.name), prefix: `${path}/` }) + } else out.push(path) + } + } + } + folders = next } + return out.slice(0, MAX_WALK + 1) } /** @@ -233,7 +250,33 @@ async function withoutLinks( return kept } -/** grep in JS: read each file through the backend and test each line. */ +/** The hits of `regex` in one file. Too big and binary files have none. */ +async function grepFile( + backend: WorkspaceBackend, + base: string, + path: string, + regex: RegExp, +) { + const hits: Array = [] + const full = pathsOf(backend).resolve(base, path) + const info = await backend.stat(full) + if (info?.type !== 'file' || info.size > MAX_SIZE) return hits + const text = await readText(backend, full) + // rg skips binary files too. + if (text.includes('\0')) return hits + const lines = text.split('\n') + for (const [index, line] of lines.entries()) { + if (!regex.test(line)) continue + hits.push({ path, line: index + 1, text: line }) + if (hits.length > MAX_MATCHES) break + } + return hits +} + +/** + * grep in JS: read the files through the backend and test each line. It reads + * up to PARALLEL files at a time, and keeps the hits in the order of `files`. + */ async function grepFiles( backend: WorkspaceBackend, base: string, @@ -241,17 +284,14 @@ async function grepFiles( regex: RegExp, ) { const hits: Array = [] - for (const path of files) { - const full = pathsOf(backend).resolve(base, path) - const info = await backend.stat(full) - if (info?.type !== 'file' || info.size > MAX_SIZE) continue - const text = await readText(backend, full) - // rg skips binary files too. - if (text.includes('\0')) continue - const lines = text.split('\n') - for (const [index, line] of lines.entries()) { - if (!regex.test(line)) continue - hits.push({ path, line: index + 1, text: line }) + for (let i = 0; i < files.length; i += PARALLEL) { + const found = await Promise.all( + files + .slice(i, i + PARALLEL) + .map((path) => grepFile(backend, base, path, regex)), + ) + for (const hit of found.flat()) { + hits.push(hit) if (hits.length > MAX_MATCHES) return hits } } @@ -305,8 +345,7 @@ export function searchTools(env: ToolEnv) { const files = await withoutLinks(backend, base, listed) return { files, partial: false } } - const walked: Array = [] - await walk(backend, base, '', walked) + const walked = await walk(backend, base) return { files: tidy(walked), partial: walked.length > MAX_WALK } } diff --git a/packages/ai-harness/src/first-party/coding/workspace.ts b/packages/ai-harness/src/first-party/coding/workspace.ts index cb810cee5f..6e0396fa43 100644 --- a/packages/ai-harness/src/first-party/coding/workspace.ts +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -174,15 +174,27 @@ export function createWorkspaceTools( } } + // ponytail: the real path of the root is read once, when the root is there. + // A root that turns into a link later is not seen. + let knownRoot: string | undefined + /** The real path of `root`. */ + const realRoot = async () => { + if (knownRoot !== undefined) return knownRoot + const real = backend.realpath ? await backend.realpath(root) : root + if (real === undefined) return realPath(root) + knownRoot = real + return real + } + /** - * The real path of `full`, as the permission rules see it. In the real - * root, it is a path in `root`, so a root that is a link (like `/tmp` on - * macOS) does not move every path outside. + * `real`, a real path, as the permission rules see it. In the real root, + * it is a path in `root`, so a root that is a link (like `/tmp` on macOS) + * does not move every path outside. */ - const realForRules = async (full: string) => { - const [real, realRoot] = await Promise.all([realPath(full), realPath(root)]) - return within(realRoot, real, paths.sep) - ? paths.join(root, paths.relative(realRoot, real)) + const forRules = async (real: string) => { + const rootReal = await realRoot() + return within(rootReal, real, paths.sep) + ? paths.join(root, paths.relative(rootReal, real)) : real } @@ -219,11 +231,16 @@ export function createWorkspaceTools( * with a stricter rule, like `notes` to `.env`. So the rules check the * real path too, and the strictest decision wins. */ - const checkRealPath = async (path: string, tool: string, full: string) => { + const checkRealPath = async ( + path: string, + tool: string, + full: string, + realFull: string, + ) => { const decide = access?.rules?.() // `bash` touches commands. Its rules do not match paths. if (!access || !decide || !resources[tool]?.paths) return - const real = await realForRules(full) + const real = await forRules(realFull) if (paths.relative(full, real) === '') return // `permissions()` checked `full` before the call. Act only when the // real path makes the decision stricter. @@ -278,10 +295,10 @@ export function createWorkspaceTools( await allowOutside(path, tool, kind, full, full) } const real = await realPath(full) - const folders = [await realPath(root), ...allowed] + const folders = [await realRoot(), ...allowed] const isAllowed = folders.some((dir) => within(dir, real, paths.sep)) if (!isAllowed) await allowOutside(path, tool, kind, full, real) - await checkRealPath(path, tool, full) + await checkRealPath(path, tool, full, real) return full } @@ -299,7 +316,7 @@ export function createWorkspaceTools( if (!decide) return undefined // rg and the other engines do not follow links in `dir`. So only `dir` // itself can be a link, and its real path covers the files in it. - const real = await realForRules(dir) + const real = await forRules(await realPath(dir)) return (path) => decide('read_file', [ paths.resolve(dir, path), diff --git a/packages/ai-harness/tests/coding-backend.test.ts b/packages/ai-harness/tests/coding-backend.test.ts index f71916bbdc..45caada742 100644 --- a/packages/ai-harness/tests/coding-backend.test.ts +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -137,6 +137,48 @@ describe('hostBackend', () => { }) }) +describe('hostBackend.readLines', () => { + /** The page and the line count of a whole-file read, as `read_file` makes them. */ + const expected = (text: string, offset: number, limit: number) => { + const all = text.replace(/^/, '').split('\n') + if (text.endsWith('\n')) all.pop() + return { + lines: all.slice(offset - 1, offset - 1 + limit), + total: all.length, + } + } + // 1 MiB is the read size: `é` is split over the edge of the first read. + const edge = `${'x'.repeat(1024 * 1024 - 1)}é\nlast` + + it.each([ + ['an empty file', '', 1, 5], + ['no final newline', 'a\nb', 1, 5], + ['a final newline', 'a\nb\n', 2, 5], + ['CRLF lines', 'a\r\nb\r\n', 1, 5], + ['a byte order mark', 'a\nb\n', 1, 1], + ['an offset after the end', 'a\nb\n', 9, 5], + ['a character over the read edge', edge, 1, 1], + ['a page after the read edge', edge, 2, 1], + ])('reads %s like a whole-file read', async (_, text, offset, limit) => { + const file = join(dir, 'f.txt') + await writeFile(file, text) + const read = await hostBackend.readLines(file, { offset, limit }) + expect({ lines: read.lines, total: read.total }).toEqual( + expected(text, offset, limit), + ) + expect(read.size).toBe(Buffer.byteLength(text)) + expect(read.head).toEqual(Buffer.from(text).subarray(0, 8 * 1024)) + }) + + it('stops after 64 KiB of lines, and still counts every line', async () => { + const file = join(dir, 'long.txt') + await writeFile(file, `${'y'.repeat(40 * 1024)}\n`.repeat(4)) + const read = await hostBackend.readLines(file, { offset: 1, limit: 10 }) + expect(read.lines).toHaveLength(2) + expect(read.total).toBe(4) + }) +}) + describe('workspace hooks', () => { it('runs afterWrite after each write, and adds afterRead text to read_file', async () => { const written: Array = [] diff --git a/packages/ai-harness/tests/coding-edit.test.ts b/packages/ai-harness/tests/coding-edit.test.ts index 10619c5d80..498f5b3b63 100644 --- a/packages/ai-harness/tests/coding-edit.test.ts +++ b/packages/ai-harness/tests/coding-edit.test.ts @@ -85,6 +85,18 @@ describe('edit_file', () => { expect(text).toBe('x\nc \n') }) + // Only the old text needs the change. The level must still run. + it.each([ + ['another Unicode form', 'café\n', 'café', 'tea\n'], + ['spaces at the end of a line', 'a\nb\n', 'a \nb', 'tea\n'], + ])( + 'matches old text in %s, in a file without it', + async (_, file, old, expected) => { + const { text } = await edit('reverse.txt', file, { old, new: 'tea' }) + expect(text).toBe(expected) + }, + ) + it('prefers an exact match over a looser one', async () => { // With trailing spaces ignored, `foo\n` is in the file twice. const { text } = await edit('prefer.txt', 'foo\nfoo \n', { diff --git a/packages/ai-harness/tests/coding-search.test.ts b/packages/ai-harness/tests/coding-search.test.ts index 3cb99ab0bd..ffeb512c3c 100644 --- a/packages/ai-harness/tests/coding-search.test.ts +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -366,6 +366,37 @@ describe('limits', () => { expect(lines[200]).toBe(MATCHES_NOTE) }) + it('reads up to 16 folders and files at a time in JS, and keeps the order', async () => { + const files = Object.fromEntries( + Array.from({ length: 40 }, (_, index) => [ + `d${String(index).padStart(2, '0')}${sep}a.txt`, + 'two\n', + ]), + ) + const { backend } = fakeBackend(files) + let running = 0 + let most = 0 + /** `fn` with a pause, counting the calls that run at the same time. */ + const counted = + , R>(fn: (...args: A) => Promise) => + async (...args: A) => { + most = Math.max(most, ++running) + await new Promise((done) => setTimeout(done, 1)) + try { + return await fn(...args) + } finally { + running-- + } + } + backend.readdir = counted(backend.readdir) + backend.stat = counted(backend.stat) + const lines = (await searchOn(backend).grep({ pattern: 'two' })).split('\n') + expect(most).toBe(16) + expect(lines).toHaveLength(40) + expect(lines[0]).toBe('d00/a.txt:1: two') + expect(lines[39]).toBe('d39/a.txt:1: two') + }) + it('says so when the JS walk stops before the end', async () => { const fake = fakeBackend({ ...numbered(5001), 'z.md': '' }) expect(await searchOn(fake.backend).listFiles({ pattern: '*.md' })).toBe(