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/.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 dd6239b947..98cb63da0c 100644 --- a/docs/config.json +++ b/docs/config.json @@ -1030,6 +1030,35 @@ "addedAt": "2026-10-13", "updatedAt": "2026-10-16" }, + { + "label": "Give your agent coding tools", + "to": "harness/coding-tools", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-14" + }, + { + "label": "Sandbox and extend the coding tools", + "to": "harness/coding-tools-backends", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-16" + }, + { + "label": "Undo a turn", + "to": "harness/snapshots", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" + }, + { + "label": "Build a coding agent", + "to": "harness/coding-agent", + "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 new file mode 100644 index 0000000000..a8bffe5f11 --- /dev/null +++ b/docs/harness/coding-agent.md @@ -0,0 +1,200 @@ +--- +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, +}) + +// A risky call waits for an answer: once, always, or reject. +// 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 + 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. +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)) + } else { + const turn = session.send(line) + turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { + process.stdout.write(event.delta) + }) + 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() +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. + +## 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. 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 + +- 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-agents.md b/docs/harness/coding-agents.md new file mode 100644 index 0000000000..125577e14c --- /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.send( + '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 new file mode 100644 index 0000000000..1641fe93e3 --- /dev/null +++ b/docs/harness/coding-tools-backends.md @@ -0,0 +1,275 @@ +--- +title: Sandbox and extend the coding tools +id: harness-coding-tools-backends +order: 7 +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 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 + +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. | +| `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. | + +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. +- 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`: + +```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 + +- 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. + +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..577d2a0a99 --- /dev/null +++ b/docs/harness/coding-tools.md @@ -0,0 +1,304 @@ +--- +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, +}) + +// 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 + 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. +let line = await terminal.question('\n> ') +while (line !== 'exit') { + const turn = session.send(line) + turn.onEvent('TEXT_MESSAGE_CONTENT', ({ event }) => { + process.stdout.write(event.delta) + }) + 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() +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. 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 + +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 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 3759862adc..d2a59313c4 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -79,7 +79,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: @@ -184,6 +184,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: @@ -245,6 +288,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 new file mode 100644 index 0000000000..cf8ab2425f --- /dev/null +++ b/docs/harness/snapshots.md @@ -0,0 +1,148 @@ +--- +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.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.` + +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.send('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). + +## 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. +- `/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 dffbd59824..51920da03c 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -49,6 +49,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": { @@ -69,11 +73,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..545ab72801 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/backend.ts @@ -0,0 +1,454 @@ +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 + /** + * 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 + /** + * 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' + +/** `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 + +/** + * 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), + 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) => { + // 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), + 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 the last 10 MiB of a long job, like `exec` keeps the first. + let output = '' + const collect = (chunk: string) => { + output = keepTail(output + chunk) + } + // 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) => { + output += String(error) + done({ exitCode: 1 }) + }) + }) + return { wait: () => exited, kill, output: () => output } + }, +} 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. */ +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..fdb0da13b9 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/edit.ts @@ -0,0 +1,197 @@ +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. `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. +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'), + 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), + }, +] + +/** + * `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) { + // 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) + 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..ad46a78431 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/formatter.ts @@ -0,0 +1,276 @@ +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() + // 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), + ) + 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..8caed0d643 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/patch.ts @@ -0,0 +1,454 @@ +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': + // 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 + 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..8bc9e392fa --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/read.ts @@ -0,0 +1,219 @@ +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 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( + { lines: window, total }: { lines: Array; total: number }, + options: { path: string; offset: number; limit: number | undefined }, +) { + const { path, offset, limit } = options + 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 ${total} lines.]` + } + const last = offset + shown.length - 1 + 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 ${total}. 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 + const pageLimit = Math.min(limit ?? MAX_LINES, MAX_LINES) + return env.lock(full, async () => { + const path = env.shown(full) + 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 (size > MAX_MEDIA_BYTES) { + const mb = (size / 1024 / 1024).toFixed(1) + throw new Error( + `${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: { + type: 'data', + value: toBase64(bytes), + mimeType: media.mimeType, + }, + } + return [part] + } + if (head?.includes(0)) { + const type = mimeTypeOf(full) + const typeNote = type === undefined ? '' : `, ${type}` + return `${path} is a binary file (${size} bytes${typeNote}). read_file shows text, images, and PDFs only.` + } + 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) { + 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..210ac24a87 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/search.ts @@ -0,0 +1,467 @@ +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 +/** 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.]' +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 +} + +/** + * 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) { + 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) +} + +/** + * `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 +} + +/** 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, + files: Array, + regex: RegExp, +) { + const hits: Array = [] + 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 + } + } + 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 = await walk(backend, base) + 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..6e0396fa43 --- /dev/null +++ b/packages/ai-harness/src/first-party/coding/workspace.ts @@ -0,0 +1,471 @@ +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 +} + +/** 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) +} + +/** 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 + } + } + + // 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 + } + + /** + * `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 forRules = async (real: string) => { + const rootReal = await realRoot() + return within(rootReal, real, paths.sep) + ? paths.join(root, paths.relative(rootReal, 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, + 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 forRules(realFull) + 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) + // `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 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, real) + 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 forRules(await realPath(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..45caada742 --- /dev/null +++ b/packages/ai-harness/tests/coding-backend.test.ts @@ -0,0 +1,225 @@ +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('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 = [] + 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.send('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..498f5b3b63 --- /dev/null +++ b/packages/ai-harness/tests/coding-edit.test.ts @@ -0,0 +1,174 @@ +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') + }) + + // 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', { + 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..6dbe9baa9b --- /dev/null +++ b/packages/ai-harness/tests/coding-formatter.test.ts @@ -0,0 +1,465 @@ +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 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) + 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.send('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({ questionId: 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.send('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.send('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.send('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.send('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.send('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..0b8a373cce --- /dev/null +++ b/packages/ai-harness/tests/coding-patch.test.ts @@ -0,0 +1,504 @@ +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`. `dirs` are folders. + */ +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)) + } + 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) => { + 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 } + }, + 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('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( + 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..ffeb512c3c --- /dev/null +++ b/packages/ai-harness/tests/coding-search.test.ts @@ -0,0 +1,533 @@ +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('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( + `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..9696515752 --- /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.send('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.send('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.send('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.send('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.send('Hello') + await session.send('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.send('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.send('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.send('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.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. + 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.send('Edit a.txt') + await session.command('undo') + await session.send('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.send('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.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.') + 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.send('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.send('Hello') + const [message] = await session.transcript() + const running = session.send('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.send('One') + await session.send('Two') + // The user edits a file that no turn changes. + await put('c.txt', 'mine') + await session.send('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.send('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.send('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.send('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.send('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 c3364633f6..c69fcc0549 100644 --- a/packages/ai-harness/tests/edge-safety.test.ts +++ b/packages/ai-harness/tests/edge-safety.test.ts @@ -18,13 +18,15 @@ const srcDir = resolve(here, '..', 'src') /** * Node-only entry points, relative to `src`. A path that ends with `/` allows * the whole folder. `build.ts` and `worker.ts` are the Node-only `./build` - * and `./worker` entries. + * and `./worker` entries. `first-party/coding/` is the Node-only + * `plugins/coding` entry. */ const ALLOWED: Array = [ 'build.ts', 'worker.ts', 'first-party/files.ts', 'first-party/bound-output.ts', + 'first-party/coding/', ] const isAllowed = (file: string) => diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts index 068e587749..e213680260 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,56 @@ 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.send('work') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + expect(question?.message).toContain('write_file') + await session.answer({ + questionId: question!.questionId, + value: { 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.send('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 3e37ad7152..235b3871c0 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.send('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..1ea8f8c01a --- /dev/null +++ b/packages/ai-harness/tests/workspace-tools.test.ts @@ -0,0 +1,1060 @@ +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.send('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.send('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.send('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({ questionId: question.questionId, value: '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.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') + 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.send('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.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}"`) + } + 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.send('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.send('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({ + questionId: question.questionId, + value: { 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.send('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.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', + ) + 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.send('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.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.', + ) + 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 }) + } + }, + ) +}) + +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-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index 97e403ac26..52524e5591 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -35,6 +35,7 @@ export default mergeConfig( './src/build.ts', './src/worker.ts', './src/first-party/index.ts', + './src/first-party/coding/index.ts', ], srcDir: './src', cjs: false, 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..b5e3c51665 --- /dev/null +++ b/packages/ai-sandbox/src/workspace-backend.ts @@ -0,0 +1,160 @@ +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 + * 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 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 + 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. + 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..a799a4167a --- /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.send('first') + await session.send('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.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.send('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.send('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.send('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.send('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..cda25d693d --- /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.send('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 18b1f45f6b..375d0bbadb 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -2654,6 +2654,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.2.1 @@ -3216,6 +3222,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 @@ -6681,6 +6690,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'} @@ -11838,6 +11850,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==} @@ -18019,6 +18094,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==} @@ -22910,6 +22989,8 @@ snapshots: - bufferutil - utf-8-validate + '@mixmark-io/domino@2.2.0': {} + '@modelcontextprotocol/client@2.1.0': dependencies: '@modelcontextprotocol/core': 2.1.0 @@ -29097,6 +29178,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 @@ -37452,6 +37584,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..12d6c61760 --- /dev/null +++ b/testing/e2e/fixtures/tools-test/coding-tools.json @@ -0,0 +1,211 @@ +{ + "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." + } + }, + { + "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/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index d5079c652e..2f42c9fa11 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -77,6 +77,7 @@ import { Route as ApiChatRouteImport } from './routes/api.chat' import { Route as ApiChatCompletionsIncompleteStreamRouteImport } from './routes/api.chat-completions-incomplete-stream' import { Route as ApiClientToolResumeStoreRouteImport } from './routes/api.client-tool-resume-store' 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' @@ -548,6 +549,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', @@ -1202,6 +1208,7 @@ export interface FileRoutesByFullPath { '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/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 @@ -1384,6 +1391,7 @@ export interface FileRoutesByTo { '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/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 @@ -1567,6 +1575,7 @@ export interface FileRoutesById { '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/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 @@ -1751,6 +1760,7 @@ export interface FileRouteTypes { | '/api/chat-completions-incomplete-stream' | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' + | '/api/coding-tools' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -1933,6 +1943,7 @@ export interface FileRouteTypes { | '/api/chat-completions-incomplete-stream' | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' + | '/api/coding-tools' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -2115,6 +2126,7 @@ export interface FileRouteTypes { | '/api/chat-completions-incomplete-stream' | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' + | '/api/coding-tools' | '/api/compaction-usage-wire' | '/api/compaction-wire' | '/api/context-overflow' @@ -2298,6 +2310,7 @@ export interface RootRouteChildren { ApiChatCompletionsIncompleteStreamRoute: typeof ApiChatCompletionsIncompleteStreamRoute ApiClientToolResumeStoreRoute: typeof ApiClientToolResumeStoreRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute + ApiCodingToolsRoute: typeof ApiCodingToolsRoute ApiCompactionUsageWireRoute: typeof ApiCompactionUsageWireRoute ApiCompactionWireRoute: typeof ApiCompactionWireRoute ApiContextOverflowRoute: typeof ApiContextOverflowRoute @@ -2886,6 +2899,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' @@ -3804,6 +3824,7 @@ const rootRouteChildren: RootRouteChildren = { ApiChatCompletionsIncompleteStreamRoute, ApiClientToolResumeStoreRoute: ApiClientToolResumeStoreRoute, 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..c0039be999 --- /dev/null +++ b/testing/e2e/src/routes/api.coding-tools.ts @@ -0,0 +1,183 @@ +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 { 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' + +/** + * 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.send(`[${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({ + questionId: question.questionId, + value: '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.send('[coding-revert] run test').stream() + await writeFile(join(root, 'other.txt'), 'mine\n') + yield* session.send('[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-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.send('[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 + * 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) + : 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 new file mode 100644 index 0000000000..da0ccbcdb4 --- /dev/null +++ b/testing/e2e/tests/tools-test/coding-tools.spec.ts @@ -0,0 +1,207 @@ +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('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, + }) => { + 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') + }) +})