From 1cdc116aa3c94b3b9433c084dd7fef34139d9f35 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:01:15 +0200 Subject: [PATCH 01/22] feat(ai-harness): first-party plugins in @tanstack/ai-harness/plugins Add permissions, agents, question, todos, title, modelPicker, projectInstructions, fileCommands, compact, usage, and boundToolOutput under the ./plugins entry. Export the portable permission names from the package root again, so edge-safe plugins can read the rules. --- .changeset/ai-harness-plugins.md | 11 + docs/config.json | 14 +- docs/harness/agents.md | 229 +++++ docs/harness/permissions.md | 345 ++++++++ docs/harness/plugins.md | 2 + docs/harness/usage.md | 53 +- packages/ai-harness/package.json | 4 + packages/ai-harness/src/first-party/agents.ts | 380 +++++++++ .../src/first-party/bound-output.ts | 247 ++++++ packages/ai-harness/src/first-party/files.ts | 319 +++++++ packages/ai-harness/src/first-party/glob.ts | 14 + packages/ai-harness/src/first-party/index.ts | 31 + .../src/first-party/model-picker.ts | 68 ++ .../ai-harness/src/first-party/permissions.ts | 728 ++++++++++++++++ .../ai-harness/src/first-party/question.ts | 189 +++++ .../src/first-party/session-tools.ts | 272 ++++++ packages/ai-harness/src/first-party/title.ts | 90 ++ packages/ai-harness/src/first-party/todos.ts | 100 +++ .../src/first-party/workspace-hooks.ts | 35 + packages/ai-harness/src/index.ts | 15 + .../ai-harness/tests/agent-middleware.test.ts | 73 +- packages/ai-harness/tests/agents.test.ts | 446 ++++++++++ .../ai-harness/tests/bound-output.test.ts | 314 +++++++ .../ai-harness/tests/discover-tools.test.ts | 109 +++ packages/ai-harness/tests/edge-safety.test.ts | 4 +- packages/ai-harness/tests/first-party.test.ts | 178 ++++ .../ai-harness/tests/instructions.test.ts | 234 ++++++ packages/ai-harness/tests/permissions.test.ts | 780 ++++++++++++++++++ packages/ai-harness/tests/question.test.ts | 190 +++++ .../ai-harness/tests/session-plugins.test.ts | 190 ++++- packages/ai-harness/tests/title-usage.test.ts | 310 +++++++ packages/ai-harness/vite.config.ts | 2 +- .../fixtures/harness-first-party/basic.json | 107 +++ testing/e2e/src/routeTree.gen.ts | 21 + .../e2e/src/routes/api.harness-first-party.ts | 245 ++++++ testing/e2e/tests/harness-first-party.spec.ts | 125 +++ 36 files changed, 6466 insertions(+), 8 deletions(-) create mode 100644 .changeset/ai-harness-plugins.md create mode 100644 docs/harness/agents.md create mode 100644 docs/harness/permissions.md create mode 100644 packages/ai-harness/src/first-party/agents.ts create mode 100644 packages/ai-harness/src/first-party/bound-output.ts create mode 100644 packages/ai-harness/src/first-party/files.ts create mode 100644 packages/ai-harness/src/first-party/glob.ts create mode 100644 packages/ai-harness/src/first-party/index.ts create mode 100644 packages/ai-harness/src/first-party/model-picker.ts create mode 100644 packages/ai-harness/src/first-party/permissions.ts create mode 100644 packages/ai-harness/src/first-party/question.ts create mode 100644 packages/ai-harness/src/first-party/session-tools.ts create mode 100644 packages/ai-harness/src/first-party/title.ts create mode 100644 packages/ai-harness/src/first-party/todos.ts create mode 100644 packages/ai-harness/src/first-party/workspace-hooks.ts create mode 100644 packages/ai-harness/tests/agents.test.ts create mode 100644 packages/ai-harness/tests/bound-output.test.ts create mode 100644 packages/ai-harness/tests/discover-tools.test.ts create mode 100644 packages/ai-harness/tests/first-party.test.ts create mode 100644 packages/ai-harness/tests/instructions.test.ts create mode 100644 packages/ai-harness/tests/permissions.test.ts create mode 100644 packages/ai-harness/tests/question.test.ts create mode 100644 packages/ai-harness/tests/title-usage.test.ts create mode 100644 testing/e2e/fixtures/harness-first-party/basic.json create mode 100644 testing/e2e/src/routes/api.harness-first-party.ts create mode 100644 testing/e2e/tests/harness-first-party.spec.ts diff --git a/.changeset/ai-harness-plugins.md b/.changeset/ai-harness-plugins.md new file mode 100644 index 0000000000..f0812eb769 --- /dev/null +++ b/.changeset/ai-harness-plugins.md @@ -0,0 +1,11 @@ +--- +'@tanstack/ai-harness': minor +--- + +Add first-party plugins in `@tanstack/ai-harness/plugins` (a Node entry): + +- `permissions()`: rules for tools, paths, and commands, the `default`, `plan`, `acceptEdits`, and `bypass` modes, questions answered with once, always, or reject, saved "always" rules with `/permissions` and `/permissions forget `, and `listSavedPermissions` / `deleteSavedPermission` for server code. The rules apply in subagent runs too. +- `agents()`: agent profiles with their own prompt, tools, model, rules, and step limit, in code or in Markdown files, `/agent` to switch the primary agent, and the built-in `build`, `plan`, `general`, and `explore` agents. +- `question()`, `todos()`, `title()`, `modelPicker()`, `projectInstructions()`, `fileCommands()`, `compact()`, `usage()` (with your own prices for calls that have no provider cost), and `boundToolOutput()`. + +`PermissionRules`, `PermissionResources`, `PermissionDecisionCapability`, and `decidePermission` are also exported from the package root. They load no Node module, so an edge-safe plugin can read the permission rules. diff --git a/docs/config.json b/docs/config.json index 0d0de346cd..df3ace69e3 100644 --- a/docs/config.json +++ b/docs/config.json @@ -981,11 +981,23 @@ { "label": "Write a plugin", "to": "harness/plugins", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" }, { "label": "Track usage and cost", "to": "harness/usage", + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" + }, + { + "label": "Ask before risky tool calls", + "to": "harness/permissions", + "addedAt": "2026-10-13" + }, + { + "label": "Switch between agent profiles", + "to": "harness/agents", "addedAt": "2026-10-13" }, { diff --git a/docs/harness/agents.md b/docs/harness/agents.md new file mode 100644 index 0000000000..8d918b519d --- /dev/null +++ b/docs/harness/agents.md @@ -0,0 +1,229 @@ +--- +title: Switch between agent profiles +id: harness-agents +order: 10 +description: "Give one harness several agents, such as a read-only planner and a builder. Each profile has its own prompt, tools, model, rules, and step limit, in code or in Markdown files." +keywords: + - tanstack ai + - harness + - agents + - agent profiles + - plan + - subagents +--- + +Your agent has two jobs. First it must plan a change without touching a file. Then it must build the change with every tool. The main model also needs a fast helper that searches the code. `agents()` gives the harness named agent profiles. Each profile has its own system prompt, tools, model, permission rules, and step limit. + +## 1. Add the plugin + +```ts group=harness-agents +import { readFile, writeFile } from 'node:fs/promises' +import { resolve } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { agents, permissions } from '@tanstack/ai-harness/plugins' +import { anthropicText } from '@tanstack/ai-anthropic' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' + +const root = process.cwd() +const sonnet = anthropicText('claude-sonnet-5-5') + +const readTool = toolDefinition({ + name: 'read_file', + description: 'Read a file in the project', + inputSchema: z.object({ path: z.string() }), +}).server(({ path }) => readFile(resolve(root, path), 'utf8')) + +const writeTool = toolDefinition({ + name: 'write_file', + description: 'Write a file in the project', + inputSchema: z.object({ path: z.string(), content: z.string() }), +}).server(async ({ path, content }) => { + await writeFile(resolve(root, path), content) + return `Wrote ${path}.` +}) + +const coder = defineHarness({ + name: 'acme/coder', + adapter: sonnet, + tools: [readTool, writeTool], + plugins: () => [ + permissions({ root, rules: [{ tool: 'write_file', decision: 'ask', kind: 'edit' }] }), + agents({ adapter: () => sonnet }), + ], +}) +``` + +`adapter` turns a model id into an adapter. Here every agent uses the same model. + +## 2. Plan, then build + +```ts group=harness-agents +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(coder, { threadId: 'dark-mode' }) + +await session.command('agent', 'plan') +const plan = await session.prompt('How do I add a dark mode?') +console.log(plan.text) + +await session.command('agent', 'build') +await session.prompt('Do step 1.') +``` + +1. The `plan` agent reads the code and gives numbered steps. It has no tool that changes a file. +2. The `build` agent changes the files, after your answer to each [permission question](./permissions). + +- `session.command('agent')` with no input shows the current agent and the list: `Agent: build. Agents: build, plan.` +- A switch applies at the next turn. + +The same switch is the `agent` setting. `session.describe().config` lists its options for a menu: + +```ts group=harness-agents +const setting = session.describe().config.find((entry) => entry.key === 'agent') +if (setting?.option.type === 'select') { + console.log(setting.value, setting.option.options) // 'build' ['build', 'plan'] +} + +await session.setConfig('agent', 'plan') +``` + +A client can change the setting only when the harness exposes it. Add `expose: { config: ['agent'] }` to `defineHarness`. + +## The built-in agents + +| Agent | Mode | Tools | What it does | +| --- | --- | --- | --- | +| `build` | `primary`, the default | Every tool | Reads code, changes files, and runs commands. | +| `plan` | `primary` | The read tools and `todo_write` | Reads the code and gives a plan in numbered steps. | +| `general` | `subagent` | Every tool | Does a task in many steps and reports back. | +| `explore` | `subagent` | The read tools | Searches the code and reports the paths that it finds. | + +The read tools are `read_file`, `list_files`, `grep`, `webfetch`, `websearch`, and `question`. An agent gets only the tools that the harness has. + +The mode of a profile says who starts it: + +- `primary`: it answers the user's turns. Pick it with the `agent` command or setting. +- `subagent`: the main model starts it as a child. +- `all`: both. + +The `plan` agent and the `plan` mode of [permissions](./permissions#pick-a-mode) are separate. The agent has only the read tools and `todo_write`. The mode denies edits and commands for every agent. To drop the built-in agents, set `builtIns: false`. + +## Add your own agent + +Give `agents()` a profile. This one reviews the changes with a stronger model: + +```ts group=harness-agents +import type { AnyTextAdapter } from '@tanstack/ai' +import type { AgentProfile } from '@tanstack/ai-harness/plugins' + +const review: AgentProfile = { + name: 'review', + description: 'Reviews the changes and lists the risks. Does not change files.', + mode: 'all', + model: 'claude-opus-5-5', + system: 'You review code. Read the diff, then list each risk with its file and line.', + tools: ['read_file', 'bash'], + permissions: [{ tool: 'bash', resource: 'git diff*', decision: 'allow' }], + steps: 20, +} + +const models: Record = { + 'claude-opus-5-5': anthropicText('claude-opus-5-5'), +} + +export const withReview = agents({ + adapter: (model) => models[model] ?? sonnet, + agents: [review], +}) +``` + +`name`, `description`, and `mode` are required. The other fields: + +| Field | What it does | +| --- | --- | +| `model` | A model id. `adapter` turns it into an adapter. Without it, a primary agent keeps the harness model, and a subagent uses the model of the main turn. | +| `system` | Text for the system prompt of each turn of this agent. | +| `tools` | The tools that this agent can call, by name. A `*` matches any text, for example `todo_*`. Without it, the agent gets every tool. | +| `permissions` | [Permission rules](./permissions#write-rules) that count while this agent is the primary agent. | +| `steps` | The step limit. See [Limit the steps](#limit-the-steps). | +| `hidden` | `true` leaves the agent out of the `agent` command and setting. The model can still start it as a subagent. | + +- A profile replaces an earlier profile with the same name. The order is the built-in agents, then `dirs`, then `agents`. So a profile named `plan` replaces the built-in `plan`. +- `default` names the primary agent of a new session. Without it, the first primary agent is the default, which is `build`. +- In the main turn, a model that the thread stored with `session.configure` wins over the profile `model`. See [What wins](./thread-settings#what-wins). +- A `tools` list also filters the subagent tools. To let the agent start subagents, add their tool names, for example `explore`. + +## Load agents from Markdown files + +Keep your agents next to your code as Markdown files. Save this file as `.agents/agents/review.md`: + +```md +--- +description: Reviews the changes and lists the risks. Does not change files. +mode: all +model: claude-opus-5-5 +tools: [read_file] +steps: 20 +--- + +You review code. Read the diff, then list each risk with its file and line. +``` + +Then give `agents()` the folder: + +```ts group=harness-agents +export const fromFiles = agents({ + adapter: (model) => models[model] ?? sonnet, + dirs: [`${root}/.agents/agents`], +}) +``` + +- The file name is the agent name, here `review`. +- The body is the system prompt. +- The frontmatter sets `description`, `mode`, `model`, `tools`, `steps`, and `hidden: true`. The default `mode` is `all`. +- An unknown `mode`, or a `steps` that is not a whole number, stops `host.open` with an error, for example `Agent file review.md: "mode" is not valid.` +- The plugin reads the folder when a session opens, and again at each `session.reload()`. A folder that it cannot read gives no agents. + +## Limit the steps + +An agent can call tools for a long time. `steps` sets the most model calls with tools in one run. With `steps: 20`: + +1. The agent makes up to 20 model calls with tools. +2. Call 21 gets `toolChoice: 'none'` and a note: answer now, and say what is left to do. +3. Then the run stops. + +The limit counts again in each turn of a primary agent and in each run of a subagent. + +Some providers do not obey `'none'`. On Amazon Bedrock, `'none'` after a tool call goes to the model as `auto`. See [Provider notes](../tools/tools#provider-notes). If the model calls a tool on the last call anyway: + +- The tool does not run. +- The model gets this tool result: `{ "error": "Step limit reached. Answer without calling tools." }` +- The run stops. That turn can end with no text answer. + +## Let the model start a subagent + +A profile with the mode `subagent` or `all` is also a subagent. The main model gets a tool for each one, with the name of the agent, for example `explore`. Ask the build agent: `use explore to find where the theme is set`. + +A subagent run gets these parts of its profile: + +- The tools of the last main turn, filtered by its `tools`. +- Its `system` prompt. +- Its `model`, or the model of the main turn. +- Its step limit. + +A subagent run stops after 50 model calls, or when an answer is cut at the output limit. The harness also limits the depth and the number of children. See [Stay within limits](./subagents#stay-within-limits). + +## Known limits + +- The `permissions` of a profile count only while it is the primary agent. When the profile runs as a subagent, its own rules do not apply. Put rules that must hold in every run in `permissions({ rules })`. See [Rules in subagents](./permissions#rules-in-subagents). +- A Markdown agent cannot set `permissions`. Put those rules in a profile in code. +- On Amazon Bedrock, the model can still call a tool on the last call. The tool does not run, but the turn can end with no text answer. +- A subagent that starts before the first turn of the session has no tools. It also needs its own `model`. + +## What you have now + +- A read-only `plan` agent, and a `build` agent to do the work. +- Your own agents in code or in Markdown files, each with its own prompt, tools, model, and rules. +- A step limit that ends a run with a text answer. +- Subagents that the main model starts to search and to research. diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md new file mode 100644 index 0000000000..15b2a9aa97 --- /dev/null +++ b/docs/harness/permissions.md @@ -0,0 +1,345 @@ +--- +title: Ask before risky tool calls +id: harness-permissions +order: 9 +description: "Make a harness agent ask before it changes a file or runs a command. Pick a mode, write rules for paths and commands, answer the questions, and keep the rules in subagents." +keywords: + - tanstack ai + - harness + - permissions + - approvals + - plan mode + - rules +--- + +Your agent has tools that change files or run commands. It must not do that without a yes from the user, and some files it must never touch. `permissions()` checks each tool call against rules. Each call runs, waits for the user, or does not run. + +## 1. Add the plugin + +```ts group=harness-permissions +import { writeFile } from 'node:fs/promises' +import { resolve } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { permissions } from '@tanstack/ai-harness/plugins' +import { openaiText } from '@tanstack/ai-openai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' + +const root = process.cwd() + +const saveNote = toolDefinition({ + name: 'save_note', + description: 'Write a note file in the project', + inputSchema: z.object({ path: z.string(), text: z.string() }), +}).server(async ({ path, text }) => { + await writeFile(resolve(root, path), text) + return `Saved ${path}.` +}) + +export const notes = defineHarness({ + name: 'acme/notes', + adapter: openaiText('gpt-6.1-sol'), + tools: [saveNote], + plugins: () => [ + permissions({ + root, + rules: [{ tool: 'save_note', decision: 'ask', kind: 'edit' }], + }), + ], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(notes, { threadId: 'notes-1' }) +``` + +The rule makes each `save_note` call wait for the user. + +## 2. Answer the first question + +The question shows up as a `harness.question` event and in `session.snapshot().pendingQuestions`. Answer it with `session.answer`: + +```ts group=harness-permissions +const turn = session.prompt('Save a note with the time of the meeting.') + +for await (const { event } of session.events()) { + if (event.type === 'CUSTOM' && event.name === HARNESS_EVENTS.question) { + const [question] = session.snapshot().pendingQuestions + // For example: Allow save_note {"path":"meeting.md","text":"..."}? + if (question) await session.answer(question.questionId, 'once') + break + } +} + +await turn +``` + +The agent writes the file and continues. + +## Pick a mode + +The `mode` setting decides how much runs without a question: + +| Mode | What happens | +| --- | --- | +| `default` | The rules apply as written. A call that asks waits for the user. | +| `plan` | Read-only. Edits, commands, and calls that ask do not run. The model gets a note to plan and to change no files. | +| `acceptEdits` | Edits run without a question. Commands still ask. | +| `bypass` | Every call runs, also a call that a rule denies. | + +Switch the mode from server code: + +```ts group=harness-permissions +await session.setConfig('mode', 'plan') +// Or run the command. It returns 'Mode: plan.' +await session.command('mode', 'plan') +``` + +`session.command('mode')` with no input shows the current mode. In `acceptEdits` mode, an edit outside `root` or of a `.env` file still asks. + +### Keep the mode on the server + +A client that can set `mode` can set `bypass`. Then every call runs with no question. So by default, a client cannot change the mode: + +- A client `config` input for `mode` gets `{ status: 'rejected', reason: 'not_exposed' }`. The mode stays. +- A client `/mode` command gets the same receipt. +- Server code can still change the mode. + +A client can set only the config keys in `expose.config`. It can run only the commands in `expose.commands`. This harness lets a client pick the model, but not the mode: + +```ts group=harness-permissions +import { modelPicker } from '@tanstack/ai-harness/plugins' + +export const team = defineHarness({ + name: 'acme/team-notes', + adapter: openaiText('gpt-6.1-sol'), + tools: [saveNote], + plugins: () => [ + permissions({ root }), + modelPicker({ + choices: { sol: openaiText('gpt-6.1-sol'), pro: openaiText('gpt-6.1-sol-pro') }, + }), + ], + expose: { config: ['model'], commands: ['model'] }, +}) +``` + +Add `mode` to `expose.config` or `expose.commands` only for clients that you trust. Each of these clients can turn the questions off. + +## Write rules + +Add your own rules with `rules`: + +```ts group=harness-permissions +export const strict = permissions({ + root, + rules: [ + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + { tool: 'edit_file', resource: 'src/**', decision: 'allow' }, + { tool: 'bash', resource: 'pnpm test*', decision: 'allow' }, + { tool: 'deploy_*', decision: 'deny' }, + ], +}) +``` + +A rule has these fields: + +- `tool`: a tool name. A `*` at the end matches every tool that starts with that text. +- `decision`: `allow`, `ask`, or `deny`. +- `resource`: optional. A glob on the path or the command that the call touches. Without it, the rule matches every call of the tool. The tool must say what it touches, see [Rules for your own tool](#rules-for-your-own-tool). +- `kind`: optional. `read`, `edit`, or `execute`. The modes read it: `plan` denies `edit` and `execute`, and `acceptEdits` allows `edit`. + +The plugin reads the rules in this order, and the last rule that matches a call wins: + +1. The rules that other plugins add for their tools. +2. Your `rules`. +3. The saved `always` answers. A saved answer never wins over a `deny`. + +Two more rules apply to each tool: + +- A tool with no rule runs without a question. To ask for each tool that has no rule, set `default: 'ask'`. +- When the last rule of a tool is a `deny` with no `resource`, the model does not get the tool. + +### Paths + +- Inside `root`, a rule matches the path from `root`, for example `src/app.ts`. Outside `root`, it matches the full path, for example `/etc/hosts`. +- `**` matches any part of a path. `*` matches text in one name. `?` matches one character. +- On macOS and Windows, letter case does not count. So a `deny` for `secrets/**` also stops `Secrets/key`. +- A call with more than one path runs only when the rules allow every path. + +These paths ask, unless a rule with a `resource` allows them: + +- A path outside `root`. +- A `.env` file, for example `.env`, `.env.local`, or `prod.env`. +- A path that the plugin cannot resolve without the file system, for example `C:notes.txt` or `notes.txt:hidden`. + +This rule lets the agent read `.env.example` without a question: + +```ts group=harness-permissions +export const exampleEnv = permissions({ + root, + rules: [{ tool: 'read_file', resource: '.env.example', decision: 'allow' }], +}) +``` + +### Commands + +- The plugin splits a command at `&&`, `||`, `;`, `|`, and new lines. The rules must allow every part. So when only `ls*` is allowed, `ls && rm -rf build` asks. +- In a command rule, `*` matches any text, also `/`. So `git *` allows `git add src/app.ts`. +- A part with shell syntax that can hide a command never gets an `allow`. This includes `$(`, a backtick, `&`, `<`, `>`, and a heredoc. Such a part asks, or a matching `deny` stops it. + +## Answer a question + +A call that asks waits for one of three answers: + +| Answer | What happens | +| --- | --- | +| `once` | This call runs. | +| `always` | This call runs, and the plugin saves an `allow` rule for it. | +| `reject` | The call does not run. The model gets a tool error. | + +- Send the answer as a string, `'once'`, or as an object, `{ answer: 'once' }`. +- A reject can carry a `message` for the model: `{ answer: 'reject', message: 'Use pnpm, not npm.' }`. The model gets your message as the tool error. +- Any other answer rejects the call. This includes `y` and `yes`. + +### What `always` saves + +- One `allow` rule for each path or command part of the call, for example `{ tool: 'bash', resource: 'pnpm test', decision: 'allow' }`. The rule names that exact text, not a glob. +- One `allow` for the whole tool, when the tool does not say what it touches. +- With a metadata store (`stores.metadata`), the rules belong to `root`. Every session with that `root` uses them, for every user, also after a restart. +- Without a metadata store, the rules stay in memory until the session closes. +- No rule for a part with shell syntax, `*`, or `?`. The next call like it asks again. + +### See and forget saved answers + +1. Run `session.command('permissions')`. The command lists the saved rules of the project: + +```text +Saved rules: +1. bash pnpm test +2. read_file .env +``` + +2. Run `session.command('permissions', 'forget 2')` to delete rule 2. The open session asks again at the next `read_file` call for `.env`. + +A client can run the command only when `expose.commands` has `permissions`. + +Server code can also read and delete the saved rules without a session. Pass the stores of your host and the same `root` that `permissions()` gets: + +```ts group=harness-permissions +import { deleteSavedPermission, listSavedPermissions } from '@tanstack/ai-harness/plugins' + +// The persistence that you give to createHarnessHost. +const persistence = memoryPersistence() + +const saved = await listSavedPermissions(persistence.stores, root) +const [first] = saved +if (first) await deleteSavedPermission(persistence.stores, root, first) +``` + +- `deleteSavedPermission` finds the rule by its fields. An unknown rule does nothing. +- New sessions stop using a deleted rule. A session that is open now can keep it until it opens again. +- Without a metadata store, both functions find no rules. The command then lists the answers of the open session. + +## Rules in subagents + +The same check runs in every agent run of the session: + +- Each call in a child run gets the same rules and the same mode. This includes the agents that the model calls, the agents that plugins start, and their children. +- A question from a child waits in the same session. Answer it the same way. +- In `plan` mode, a child cannot change a file. Its call gets the tool error `This tool is not allowed in plan mode.` +- The `permissions` of an [agent profile](./agents) count only while that profile is the primary agent. They do not apply when the profile runs as a subagent. + +## Rules for your own tool + +A rule with a `resource` can match the calls of your tool only when the tool says what it touches. Add the tool to `PermissionResources`: + +```ts +import { copyFile } from 'node:fs/promises' +import { resolve } from 'node:path' +import { toolDefinition } from '@tanstack/ai' +import { PermissionResources, PermissionRules, definePlugin } from '@tanstack/ai-harness' +import { z } from 'zod' + +const root = process.cwd() + +const copy = toolDefinition({ + name: 'copy_file', + description: 'Copy a file in the project', + inputSchema: z.object({ from: z.string(), to: z.string() }), +}).server(async ({ from, to }) => { + await copyFile(resolve(root, from), resolve(root, to)) + return `Copied ${from} to ${to}.` +}) + +/** A string field of the raw input. It throws when the field is missing. */ +function field(input: unknown, key: string) { + const value = + typeof input === 'object' && input !== null ? Reflect.get(input, key) : undefined + if (typeof value !== 'string') throw new Error(`No ${key}.`) + return value +} + +export const fileCopy = definePlugin({ + name: 'acme/file-copy', + setup: () => ({ + tools: [copy], + contribute: [ + PermissionRules.item({ tool: 'copy_file', decision: 'ask', kind: 'edit' }), + PermissionResources.item({ + copy_file: { paths: (input) => [field(input, 'from'), field(input, 'to')] }, + }), + ], + }), +}) +``` + +Now `{ tool: 'copy_file', resource: 'docs/**', decision: 'allow' }` copies files in `docs` without a question. + +- Each function gets the raw input of the model, before the schema check. If it throws, the call does not run, and the model gets `Cannot check the permissions of this call: No from.` +- `paths` are relative to `root`, or absolute. `commands` are command parts, one for each simple command. +- `kind: 'edit'` puts the tool under the modes. `plan` denies it, and `acceptEdits` runs it. +- If a tool does not say what it touches, a rule with a `resource` counts for every call. So a `deny` for one folder denies every call of the tool. + +## Ask for a decision from your plugin + +Your plugin must know what `permissions()` decides for a tool before the model calls it. Read the answer with `PermissionDecisionCapability`: + +```ts +import { PermissionDecisionCapability, definePlugin } from '@tanstack/ai-harness' + +export const planTools = definePlugin({ + name: 'acme/plan-tools', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => ({ + // Give the model only the tools that permissions() allows in plan mode. + prepareTools: ({ tools }) => { + const decide = ctx.getOptional(PermissionDecisionCapability) + if (!decide) return tools + return tools.filter((tool) => decide(tool.name, 'plan') === 'allow') + }, + }), +}) +``` + +- The answer is `'allow'`, `'ask'`, or `'deny'`. It uses all rules in their order, your `rules` last, and your `default`. +- Pass what a call touches as a third argument, for example `decide('read_file', 'default', { paths: [file] })`. Then the path rules and `root` apply too. +- Read it when the plugin runs, not in `setup`. Then `permissions()` can come before or after your plugin. +- Saved `always` answers are not part of it. They can only allow, so a check that needs `'allow'` stays strict. +- `PermissionRules`, `PermissionResources`, `PermissionDecisionCapability`, and `decidePermission` come from `@tanstack/ai-harness`. They load no Node module, so a plugin for an edge runtime can use them. The plugins themselves come from `@tanstack/ai-harness/plugins`, which is Node only. + +## Known limits + +- A rule that allows a tool with no `kind` also lets it run in `plan` mode. `plan` denies only edits, commands, and calls that ask. +- A client that can send inputs to the thread can answer the questions, also with `always`. Control who can open a thread. +- `permissions()` does not follow links. Your own tool can check the real path of a link with `PermissionDecisionCapability`. + +## What you have now + +- A question before each risky tool call. +- Modes to plan, to edit fast, or to run everything. +- A mode that a client cannot change, unless you expose it. +- Rules for folders, files, and commands, and saved answers for each project that you can list and forget. +- The same rules in every agent run. + +Next: give the agent a planner and a builder with [agent profiles](./agents). diff --git a/docs/harness/plugins.md b/docs/harness/plugins.md index 05a2143851..2d8398e213 100644 --- a/docs/harness/plugins.md +++ b/docs/harness/plugins.md @@ -395,3 +395,5 @@ If a `setup` throws, the session keeps working with no plugins. The promise reje - A picker that chooses the model of each turn from its message, context, and sender. - Plugins that share services, lists, and events without knowing each other. - Commands, agents, and whole plugins that change while the session runs. + +Next: add the first-party plugins from `@tanstack/ai-harness/plugins`. [Ask before risky tool calls](./permissions) and [Switch between agent profiles](./agents) are two of them. diff --git a/docs/harness/usage.md b/docs/harness/usage.md index 2acdbbbd16..3a936f03d9 100644 --- a/docs/harness/usage.md +++ b/docs/harness/usage.md @@ -2,7 +2,7 @@ title: Track usage and cost id: harness-usage order: 7 -description: "Read the tokens and the cost of a thread, in total, by model, and by sender. Show a live total, and stop a person who is over budget." +description: "Read the tokens and the cost of a thread, in total, by model, and by sender. Show a live total, stop a person who is over budget, and add /usage with your own prices." keywords: - tanstack ai - harness @@ -111,6 +111,56 @@ for (const entry of page.entries) { A background agent can end after the turn. Then the index entry counts its tokens at the end of the next turn. +## Show usage and price the calls + +Many providers report no cost, so `cost` stays empty. The `usage()` plugin adds a `/usage` command, and can price those calls with your own price list: + +```ts group=harness-usage +import { usage } from '@tanstack/ai-harness/plugins' + +// USD per 1M tokens, by model id. +const prices: Record = { + 'gpt-6.1-sol': { input: 2, output: 8, cacheRead: 0.2 }, +} + +export const priced = defineHarness({ + name: 'acme/priced', + adapter: openaiText('gpt-6.1-sol'), + plugins: () => [ + usage({ + model: (id) => { + const cost = prices[id] + return cost === undefined ? undefined : { cost } + }, + }), + ], +}) +``` + +`session.command('usage')` reads `session.usage()` and shows the total. When the thread used more than one model, it also shows a line for each model: + +```text +3 model calls, 41200 input tokens (30000 cache read, 0 cache write), 2100 output tokens, 43300 total. Cost: $0.0432. +openai/gpt-6.1-sol: 2 model calls, 40000 input tokens (30000 cache read, 0 cache write), 2000 output tokens, 42000 total. +openai/gpt-6-luna: 1 model calls, 1200 input tokens, 100 output tokens, 1300 total. +``` + +How the plugin finds the cost: + +- A call with a cost from its provider keeps that cost, as it is. +- `model` prices only the calls with no provider cost. It gets the model id without the provider, like `gpt-6.1-sol`. It returns `{ cost }` with prices in USD per 1M tokens, or `undefined`. +- Without `model`, `/usage` shows `Cost:` only when a provider reported a cost. +- A call with no price adds a note, for example `(cost unknown for 1 call)`. + +The plugin also writes the cost to `usage.cost` of the session index entry, whenever it has one. `session.usage()` keeps only the costs that the providers report. + +The plugin state, `snapshot().plugins['tanstack/usage']`, has a copy of `session.usage().total` and `contextTokens`: the input tokens of the latest call of the lead model. Compare it with the context window of the model to show how full the context is. + +Limits of the plugin: + +- For one model, when some calls have a provider cost and others do not, `model` does not price the others. +- On a host with a log, the cost in the index entry and the state copy can miss the last call of a run. They get it at the end of the next run. + ## What you have now - The tokens and the cost of a thread, in total, by model, and by sender. @@ -118,3 +168,4 @@ A background agent can end after the turn. Then the index entry counts its token - A live total from the `harness.usage` event. - A budget check before each prompt. - The token counts of each session in the session index. +- `/usage` with your own prices, and a cost for each session in the session index. diff --git a/packages/ai-harness/package.json b/packages/ai-harness/package.json index 5d45fa9b9d..1272c358e4 100644 --- a/packages/ai-harness/package.json +++ b/packages/ai-harness/package.json @@ -29,6 +29,10 @@ ".": { "types": "./dist/esm/index.d.ts", "import": "./dist/esm/index.js" + }, + "./plugins": { + "types": "./dist/esm/first-party/index.d.ts", + "import": "./dist/esm/first-party/index.js" } }, "scripts": { diff --git a/packages/ai-harness/src/first-party/agents.ts b/packages/ai-harness/src/first-party/agents.ts new file mode 100644 index 0000000000..7e4a2926e9 --- /dev/null +++ b/packages/ai-harness/src/first-party/agents.ts @@ -0,0 +1,380 @@ +import { + combineStrategies, + defineAgent, + maxIterations, + untilFinishReason, +} from '@tanstack/ai' +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { definePlugin } from '../plugins' +import { readMarkdownFiles } from './files' +import { globToRegExp } from './glob' +import { PermissionRules } from './permissions' +import type { + AnyTextAdapter, + AnyTool, + ChatMiddleware, + ChatMiddlewareConfig, + KeyedAdapter, +} from '@tanstack/ai' +import type { FrontmatterValue } from './files' +import type { PermissionRule } from './permissions' + +/** A named agent: its prompt, tools, model, rules, and step limit. */ +export interface AgentProfile { + name: string + /** What the agent does. The model reads it to pick a subagent. */ + description: string + /** + * - `primary`: it answers the user's turns. Pick it with `/agent`. + * - `subagent`: the main model starts it with a tool named after it. + * - `all`: both. + */ + mode: 'primary' | 'subagent' | 'all' + /** + * A model id. `agents({ adapter })` turns it into an adapter. Without it, + * a primary agent keeps the harness model, and a subagent uses the model + * of the main turn. + */ + model?: string + /** System prompt text added to each turn of this agent. */ + system?: string + /** + * The tools this agent can call, by name. A `*` matches any text, for + * example `todo_*`. Without it, the agent gets every tool. + */ + tools?: ReadonlyArray + /** + * Permission rules that count while this agent is the primary agent. + * `permissions()` reads them. They do not apply in a subagent run. + */ + permissions?: ReadonlyArray + /** + * The most model calls with tools in one run. The next call gets no tool + * calls and a note to answer now, and then the run stops. + */ + steps?: number + /** + * Leave it out of the `/agent` list and the `agent` setting. The model + * can still start it as a subagent. + */ + hidden?: boolean +} + +/** The tools that only read: files, search, the web, and questions. */ +const READ_TOOLS = [ + 'read_file', + 'list_files', + 'grep', + 'webfetch', + 'websearch', + 'question', +] + +/** The built-in agents: `build`, `plan`, `general`, and `explore`. */ +export const builtInAgents: ReadonlyArray = [ + { + name: 'build', + description: 'Does the work: reads code, changes files, and runs commands.', + mode: 'primary', + system: + 'You are the build agent. Do the task the user gives you. Read the code you need, change files, and run commands. Check your work before you say it is done.', + }, + { + name: 'plan', + description: 'Reads the code and makes a plan. Does not change files.', + mode: 'primary', + tools: [...READ_TOOLS, 'todo_write'], + system: + 'You are the plan agent. You can only read and search. Do not change files or run commands. Study the code, then give the user a clear plan in numbered steps.', + }, + { + name: 'general', + description: + 'Does a research or work task in many steps and reports back. Give it the full task.', + mode: 'subagent', + system: + 'You are a general agent. Do the task you get in as many steps as it needs. Use your tools to search, read, and act. At the end, reply with a short report of what you found or did.', + }, + { + name: 'explore', + description: + 'Searches the code fast and reports what it finds. Read-only. Give it a question.', + mode: 'subagent', + tools: READ_TOOLS, + system: + 'You are a fast search agent. Find the files and code that answer the question you get. Only read and search. Reply with the paths and a short note for each.', + }, +] + +const MODES = ['primary', 'subagent', 'all'] as const + +const MAX_STEPS_PROMPT = + 'You reached the step limit. Do not call more tools. Answer now with what you have, and say what is left to do.' + +const MAX_STEPS_RESULT = 'Step limit reached. Answer without calling tools.' + +/** The tools whose name matches one of `patterns`. Without patterns, all of them. */ +function allowedTools( + tools: ReadonlyArray, + patterns: ReadonlyArray | undefined, +) { + if (patterns === undefined) return [...tools] + const matchers = patterns.map((pattern) => globToRegExp(pattern)) + return tools.filter((tool) => + matchers.some((matcher) => matcher.test(tool.name)), + ) +} + +/** + * Apply a profile to each model call of a chat run: keep only its tools, + * and at its step limit, make one last call with no tool calls. + */ +function profileMiddleware(current: () => AgentProfile) { + const isLastStep = (iteration: number) => { + const { steps } = current() + return steps !== undefined && iteration >= steps + } + const middleware: ChatMiddleware = { + name: 'tanstack/agents', + onConfig: (run, config) => { + if (run.phase !== 'init' && run.phase !== 'beforeModel') return + const { tools } = current() + const update: Partial = {} + if (tools !== undefined) update.tools = allowedTools(config.tools, tools) + if (run.phase === 'beforeModel' && isLastStep(run.iteration)) { + update.toolChoice = 'none' + // The config keeps the prompts, so a retried call has it already. + if (!config.systemPrompts.includes(MAX_STEPS_PROMPT)) { + update.systemPrompts = [...config.systemPrompts, MAX_STEPS_PROMPT] + } + } + return update + }, + // Some providers do not follow `toolChoice: 'none'`. Bedrock Converse + // sends it as `auto` after tool calls. So refuse a call of the last step + // here. In the tool phase, `iteration` is still the one of that call. + onBeforeToolCall: (run) => + isLastStep(run.iteration) + ? { type: 'skip', result: { error: MAX_STEPS_RESULT } } + : undefined, + // After the last call, stop. The harness loop limit still applies. + onShouldContinue: (_run, state) => { + const { steps } = current() + return steps === undefined || state.iterationCount <= steps + }, + } + return middleware +} + +/** A profile from a Markdown file. The body is its system prompt. */ +function profileOfFile(file: { + name: string + fields: Record + body: string +}) { + const text = (key: string) => { + const value = file.fields[key] + return typeof value === 'string' && value !== '' ? value : undefined + } + const fail = (key: string) => { + throw new Error(`Agent file ${file.name}.md: "${key}" is not valid.`) + } + const modeText = text('mode') + const mode = + modeText === undefined + ? 'all' + : (MODES.find((candidate) => candidate === modeText) ?? fail('mode')) + const stepsText = text('steps') + const steps = stepsText === undefined ? undefined : Number(stepsText) + if (steps !== undefined && !(Number.isInteger(steps) && steps >= 0)) { + fail('steps') + } + const model = text('model') + const { tools } = file.fields + const system = file.body.trim() + // ponytail: no `permissions` from files, the parser reads no nested values + const profile: AgentProfile = { + name: file.name, + description: text('description') ?? `The ${file.name} agent`, + mode, + ...(model !== undefined && { model }), + ...(system !== '' && { system }), + // `tools: a, b` is a list too, not one tool named `a, b`. + ...(tools !== undefined && { + tools: + typeof tools === 'string' + ? tools + .split(',') + .map((name) => name.trim()) + .filter((name) => name !== '') + : tools, + }), + ...(steps !== undefined && { steps }), + ...(text('hidden') === 'true' && { hidden: true }), + } + return profile +} + +/** + * Named agents, each with its own system prompt, tools, model, permission + * rules, and step limit. + * + * - A primary agent answers the user's turns. Switch it with + * `/agent ` or the `agent` setting. It applies at the next turn. + * - A subagent is a child that the main model starts with a tool named + * after the agent. + * + * The built-in agents are `build` (primary, the default, every tool), + * `plan` (primary, read-only tools), `general` (subagent, every tool), and + * `explore` (subagent, read-only tools). See {@link builtInAgents}. + * + * @param options.adapter - Turns a model id into an adapter: the `model` of + * a profile, or for a subagent without one, the model id of the main turn. + * A `keyedAdapter(...)` is built with the user's key. + * @param options.agents - More profiles. A profile replaces an earlier one + * with the same name: built-ins first, then `dirs`, then `agents`. + * @param options.dirs - Folders of `*.md` agent files. The file name is the + * agent name and the body is its system prompt. The frontmatter sets + * `description`, `mode` (default `all`), `model`, `tools` (a list), + * `steps`, and `hidden`. + * @param options.builtIns - `false` drops the built-in agents. + * @param options.default - The primary agent of a new session. Default: the + * first primary agent, `build` with the built-ins. + * + * @example + * ```ts + * const models: Record = { + * 'claude-sonnet-4-5': anthropicText('claude-sonnet-4-5'), + * } + * agents({ + * adapter: (model) => models[model] ?? openaiText('gpt-5.5'), + * dirs: ['.agents/agents'], + * }) + * ``` + */ +export function agents(options: { + adapter: (model: string) => AnyTextAdapter | KeyedAdapter + agents?: ReadonlyArray + dirs?: ReadonlyArray + builtIns?: boolean + default?: string +}) { + return definePlugin({ + name: 'tanstack/agents', + setup: async (ctx) => { + const dirs = options.dirs ?? [] + const files = (await Promise.all(dirs.map(readMarkdownFiles))).flat() + const all = [ + ...(options.builtIns === false ? [] : builtInAgents), + ...files.map(profileOfFile), + ...(options.agents ?? []), + ] + // A later profile replaces an earlier one and keeps its place. + const profiles = [ + ...new Map(all.map((profile) => [profile.name, profile])).values(), + ] + const visible = profiles.filter( + (profile) => profile.mode !== 'subagent' && !profile.hidden, + ) + const names = visible.map((profile) => profile.name) + const first = visible.find( + (profile) => profile.name === (options.default ?? names[0]), + ) + if (!first) { + throw new Error( + 'agents(): `default` must name a primary agent that is not hidden.', + ) + } + const active = () => { + const name = ctx.config.get('agent') + return visible.find((profile) => profile.name === name) ?? first + } + // The tools and the model id of the last main turn. Subagents use them. + // ponytail: a subagent started before any turn has no tools and needs a `model` + let lead: { tools: ReadonlyArray; model: string } | undefined + + const subagentOf = (profile: AgentProfile) => + defineAgent({ + name: profile.name, + description: profile.description, + run: async (child) => { + const model = profile.model ?? lead?.model + if (model === undefined) { + throw new Error( + `Agent ${profile.name} has no model. Give it a model, or start it from a turn.`, + ) + } + return child.chat({ + adapter: await child.keys.adapter(options.adapter(model)), + tools: allowedTools(lead?.tools ?? [], profile.tools), + systemPrompts: + profile.system === undefined ? [] : [profile.system], + middleware: [profileMiddleware(() => profile)], + // The harness main loop has the same limits. chat() alone stops + // at 5 calls. An answer cut at the output limit ends the run. + agentLoopStrategy: combineStrategies([ + maxIterations(50), + untilFinishReason(['length']), + ]), + }) + }, + }) + + // A profile's rules count only while it is the primary agent. At other + // times the getter gives a tool name that matches no tool. + // ponytail: subagent runs get no profile rules, that needs the run in PermissionRules + const rules = profiles.flatMap((profile) => + (profile.permissions ?? []).map((rule) => + PermissionRules.item({ + ...rule, + get tool() { + return active().name === profile.name ? rule.tool : '' + }, + }), + ), + ) + + return { + config: { + agent: configOption.select({ + options: names, + default: first.name, + description: 'The primary agent', + }), + }, + commands: { + agent: defineCommand({ + description: `Show or switch the primary agent (${names.join(', ')})`, + run: async (input: unknown) => { + const name = typeof input === 'string' ? input.trim() : '' + const list = `Agents: ${names.join(', ')}.` + if (name === '') return `Agent: ${active().name}. ${list}` + const receipt = await ctx.session.setConfig('agent', name) + return receipt.status === 'rejected' + ? `Unknown agent "${name}". ${list}` + : `Agent: ${name}. It applies at the next turn.` + }, + }), + }, + prompts: [ + { id: 'tanstack/agents:system', text: () => active().system ?? '' }, + ], + adapter: () => { + const { model } = active() + return model === undefined ? undefined : options.adapter(model) + }, + prepareTools: ({ tools, model }) => { + lead = { tools, model } + return tools + }, + // In onConfig, so the subagent tools that chat() adds are filtered too. + middleware: [profileMiddleware(active)], + subagents: profiles + .filter((profile) => profile.mode !== 'primary') + .map(subagentOf), + contribute: rules, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/bound-output.ts b/packages/ai-harness/src/first-party/bound-output.ts new file mode 100644 index 0000000000..0cd811144b --- /dev/null +++ b/packages/ai-harness/src/first-party/bound-output.ts @@ -0,0 +1,247 @@ +import { mkdir, readdir, rm, writeFile } from 'node:fs/promises' +import { join } from 'node:path' +import { normalizeToolResult } from '@tanstack/ai' +import { definePlugin } from '../plugins' +import { isRecord } from '../utils' +import type { AfterToolCallInfo, AnyChatMiddleware } from '@tanstack/ai' + +const HOUR = 60 * 60 * 1000 +const DAY = 24 * HOUR +/** Saved outputs start with this. The cleanup removes only these files. */ +const SAVED_PREFIX = 'tool-output-' +/** Core marks subagent tools with this global symbol. */ +const SUBAGENT_TOOL = Symbol.for('tanstack.ai.subagentTool') +const encoder = new TextEncoder() +// ignoreBOM keeps a leading byte order mark in the kept text. +const decoder = new TextDecoder('utf-8', { ignoreBOM: true }) + +/** The number of lines in `text`. A final newline does not start a new line. */ +function countLines(text: string) { + return text.split('\n').length - (text.endsWith('\n') ? 1 : 0) +} + +/** True for a UTF-8 byte inside a character, not at its start. */ +function isContinuationByte(byte: number | undefined) { + return byte !== undefined && (byte & 0xc0) === 0x80 +} + +/** The start of `text` inside both limits. `bytes` is `text` as UTF-8. */ +function firstPart( + text: string, + bytes: Uint8Array, + limits: { maxLines: number; maxBytes: number }, +) { + let end = limits.maxBytes + // Step back to the start of a character, so the cut never splits one. + while (isContinuationByte(bytes[end])) end-- + const firstBytes = decoder.decode(bytes.subarray(0, end)) + const firstLines = text.split('\n', limits.maxLines).join('\n') + return firstLines.length < firstBytes.length ? firstLines : firstBytes +} + +/** The end of `text` inside both limits. `bytes` is `text` as UTF-8. */ +function lastPart( + text: string, + bytes: Uint8Array, + limits: { maxLines: number; maxBytes: number }, +) { + let start = Math.max(0, bytes.length - limits.maxBytes) + // Step forward to the start of a character, so the cut never splits one. + while (isContinuationByte(bytes[start])) start++ + const lastBytes = decoder.decode(bytes.subarray(start)) + const lines = text.split('\n') + // A final newline ends the last line. It does not start a new one. + const finalNewline = text.endsWith('\n') ? 1 : 0 + const from = Math.max(0, lines.length - limits.maxLines - finalNewline) + const lastLines = lines.slice(from).join('\n') + return lastLines.length < lastBytes.length ? lastLines : lastBytes +} + +/** + * Keep the part of `text` that fits in `maxLines` lines and `maxBytes` UTF-8 + * bytes: the first lines (`keep: 'head'`, the default) or the last lines + * (`keep: 'tail'`). When something is cut, `truncated` is `true` and a note + * tells how much is shown: after the text for `head`, before it for `tail`. + * A cut never splits a character. Text inside both limits comes back + * unchanged. + * + * @example + * ```ts + * const { text, truncated } = boundText(output, { + * maxLines: 2000, + * maxBytes: 50 * 1024, + * keep: 'tail', + * }) + * ``` + */ +export function boundText( + text: string, + limits: { maxLines: number; maxBytes: number; keep?: 'head' | 'tail' }, +) { + const lines = countLines(text) + const bytes = encoder.encode(text) + const isOverLimit = lines > limits.maxLines || bytes.length > limits.maxBytes + if (!isOverLimit) return { text, truncated: false } + const isTail = limits.keep === 'tail' + const kept = isTail + ? lastPart(text, bytes, limits) + : firstPart(text, bytes, limits) + const shown = `${countLines(kept)} of ${lines} lines, ${encoder.encode(kept).length} of ${bytes.length} bytes` + return isTail + ? { + text: `[Output cut. Showing the last ${shown}.]\n\n${kept}`, + truncated: true, + } + : { + text: `${kept}\n\n[Output cut. Showing the first ${shown}.]`, + truncated: true, + } +} + +/** + * The time a saved output was written, from its file name, so no stat call + * is needed. `NaN` for a file this plugin did not save. + */ +function savedAt(name: string) { + if (!name.startsWith(SAVED_PREFIX)) return Number.NaN + return Number(name.slice(SAVED_PREFIX.length).split('-')[0]) +} + +/** Remove the saved outputs in `dir` that are older than `maxAgeMs`. */ +async function removeOldOutputs(dir: string, maxAgeMs: number) { + const cutoff = Date.now() - maxAgeMs + const names = await readdir(dir) + const old = names.filter((name) => savedAt(name) < cutoff) + await Promise.all(old.map((name) => rm(join(dir, name), { force: true }))) +} + +/** + * Save the full output of one tool call in `dir`. Gives back the note for + * the model: the file path, or why the save failed. + */ +async function saveFullOutput(dir: string, toolCallId: string, text: string) { + const safeId = toolCallId.replace(/[^\w-]/g, '_') + const path = join(dir, `${SAVED_PREFIX}${Date.now()}-${safeId}.txt`) + try { + await mkdir(dir, { recursive: true }) + await writeFile(path, text) + return `[Full output saved to ${path}.]` + } catch (error) { + // The cut result is still useful, so a failed save does not fail the call. + return `[The full output was not saved: ${String(error)}]` + } +} + +/** + * The result the model sees now. A failed call shows `{ error: message }`, + * unless an earlier middleware replaced it. `undefined` when there is + * nothing to cut. + */ +function currentResult(info: AfterToolCallInfo) { + if (info.ok || info.result !== undefined) return info.result + // A failed subagent call also shows its subagentRunId, which `info` does + // not carry. Leave it alone, so the model can still continue the child. + const isSubagent = info.tool !== undefined && SUBAGENT_TOOL in info.tool + if (isSubagent || !(info.error instanceof Error)) return undefined + return { error: info.error.message } +} + +/** + * The long part of a result, and a function that puts the cut text back. + * An error result keeps its other fields. A subagent result keeps its + * `subagentRunId`, so the model can still continue the child. + */ +function longPart(result: unknown, ok: boolean) { + if (isRecord(result)) { + if (!ok && 'error' in result) { + return { + long: result.error, + put: (cut: string) => ({ ...result, error: cut }), + } + } + if (typeof result.subagentRunId === 'string' && 'result' in result) { + return { + long: result.result, + put: (cut: string) => ({ ...result, result: cut }), + } + } + } + return { long: result, put: (cut: string) => cut } +} + +/** + * Keep every tool result small enough for the model. A result over + * `maxLines` lines (default 2000) or `maxBytes` UTF-8 bytes (default 50 KiB) + * is cut to its first lines, with a note. + * + * - A string result is cut as is. Any other value is cut as JSON text. + * - A failed call has its error text cut. It stays an error. + * - A subagent result has the child's answer cut. Its `subagentRunId` stays. + * - Content parts (images, audio, and so on) are not changed. + * + * It runs for every tool, MCP tools too, in the lead turn and in every agent + * run. With `dir`, the full output is saved to a file in `dir` and the note + * gives its path. Saved files older than `retentionDays` (default 7) are + * removed, at most once an hour. Node only. + * + * @example + * ```ts + * defineHarness({ + * name: 'acme/coder', + * adapter, + * plugins: () => [boundToolOutput({ dir: '.agent/tool-output' })], + * }) + * ``` + */ +export function boundToolOutput( + options: { + maxLines?: number + maxBytes?: number + dir?: string + retentionDays?: number + } = {}, +) { + const { + maxLines = 2000, + maxBytes = 50 * 1024, + dir, + retentionDays = 7, + } = options + return definePlugin({ + name: 'tanstack/bound-tool-output', + setup: () => { + let lastCleanup = 0 + const cleanUp = async (folder: string) => { + const now = Date.now() + if (now - lastCleanup < HOUR) return + lastCleanup = now + // ponytail: best effort. A failed cleanup tries again an hour later. + await removeOldOutputs(folder, retentionDays * DAY).catch(() => {}) + } + const bound = { + name: 'tanstack/bound-tool-output', + onAfterToolCall: async (_ctx, info) => { + const result = currentResult(info) + if (result === undefined) return + const { long, put } = longPart(result, info.ok) + const full = normalizeToolResult(long) + // Content parts (images, audio, and so on) stay as they are. + if (typeof full !== 'string') return + const bounded = boundText(full, { maxLines, maxBytes }) + if (!bounded.truncated) return + if (dir === undefined) { + return { type: 'replaceResult', result: put(bounded.text) } + } + const saved = await saveFullOutput(dir, info.toolCallId, full) + await cleanUp(dir) + return { + type: 'replaceResult', + result: put(`${bounded.text}\n${saved}`), + } + }, + } satisfies AnyChatMiddleware + // The same middleware for the lead turn and every agent run. + return { middleware: [bound], agentMiddleware: [bound] } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/files.ts b/packages/ai-harness/src/first-party/files.ts new file mode 100644 index 0000000000..789d464e06 --- /dev/null +++ b/packages/ai-harness/src/first-party/files.ts @@ -0,0 +1,319 @@ +import { execFile } from 'node:child_process' +import { existsSync, readFileSync, statSync } from 'node:fs' +import { readFile, readdir } from 'node:fs/promises' +import { homedir } from 'node:os' +import { + basename, + dirname, + extname, + isAbsolute, + join, + relative, + resolve, + sep, +} from 'node:path' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import { WorkspaceHooks } from './workspace-hooks' +import type { AnyCommand } from '../commands' + +/** Prompt slots kept free for instruction files that change. */ +const SPARE_SLOTS = 8 + +/** The mtime of the file at `path`, or `undefined` when there is none. */ +function mtimeOf(path: string) { + try { + return statSync(path).mtimeMs + } catch { + return undefined + } +} + +/** The trimmed text of the file at `path`, or `''` when there is none. */ +function textOf(path: string) { + try { + return readFileSync(path, 'utf8').trim() + } catch { + return '' + } +} + +/** + * The folders from the repo root (the nearest folder with `.git`) down to + * `root`, outermost first. Only `root` when no folder has `.git`. + */ +function repoFolders(root: string) { + const folders = [root] + let dir = root + while (!existsSync(join(dir, '.git'))) { + const parent = dirname(dir) + if (parent === dir) return { folders: [root], isRepo: false } + dir = parent + folders.unshift(dir) + } + return { folders, isRepo: true } +} + +/** The date, platform, working folder, and git facts, as a prompt. */ +async function environment(root: string, isRepo: boolean) { + const now = new Date() + // toISOString is UTC. Move by the offset to get the local date. + const date = new Date(now.getTime() - now.getTimezoneOffset() * 60_000) + .toISOString() + .slice(0, 10) + const lines = [ + `Date: ${date}`, + `Platform: ${process.platform}`, + `Working folder: ${root}`, + `Git repository: ${isRepo ? 'yes' : 'no'}`, + ] + if (isRepo) { + const branch = await new Promise((done) => + execFile( + 'git', + ['rev-parse', '--abbrev-ref', 'HEAD'], + { cwd: root, timeout: 5000 }, + (error, stdout) => done(error ? '' : stdout.trim()), + ), + ) + if (branch !== '') lines.push(`Git branch: ${branch}`) + } + return `Environment:\n${lines.map((line) => `- ${line}`).join('\n')}` +} + +/** The note for an instruction file whose text changed from `before`. */ +function changeNote(title: string, before: string, text: string) { + if (text === '') return `${title} were removed.` + if (before === '') return `${title}:\n${text}` + return `${title} changed. Follow this text instead:\n${text}` +} + +/** + * Add instruction files (AGENTS.md, CLAUDE.md) and an environment block to + * the system prompt. + * + * - `global`: files that come first, for example `~/.config/AGENTS.md`. `~` + * is the home folder. + * - Then `files` from each folder between the repo root (the folder with + * `.git`) and `root`, outermost first. With no repo root, only `root`. + * - `env` (default `true`): the date, platform, working folder, and git + * repository and branch, read when the session opens. + * + * At the start of each turn, the plugin checks the mtime of each file. A note + * with the new text of a changed file goes after this plugin's other prompts. + * Put this plugin last: when no prompt comes after its prompts, an adapter + * with mid-conversation changes adds the note to the conversation, and the + * prompt cache holds. Else the change is not additive, and the prompt cache + * starts again. + * + * When `read_file` reads a file below `root`, the instruction files in the + * folders between `root` and that file are added to the result, once per + * file in the session. + * + * @example + * ```ts + * projectInstructions({ root: process.cwd(), global: ['~/.config/AGENTS.md'] }) + * ``` + */ +export function projectInstructions(options: { + root: string + files?: ReadonlyArray + global?: ReadonlyArray + env?: boolean +}) { + const { files: names = ['AGENTS.md', 'CLAUDE.md'], env = true } = options + return definePlugin({ + name: 'tanstack/project-instructions', + setup: async () => { + const root = resolve(options.root) + const { folders, isRepo } = repoFolders(root) + const shownPath = (path: string) => + relative(root, path).split(sep).join('/') + const globals = (options.global ?? []).map((path) => ({ + path: resolve(path.replace(/^~(?=$|[\\/])/, homedir())), + title: `Global instructions from ${path}`, + })) + const projects = folders.flatMap((dir) => + names.map((name) => ({ + path: join(dir, name), + title: `Project instructions from ${shownPath(join(dir, name))}`, + })), + ) + const files = [...globals, ...projects].map((file) => ({ + ...file, + mtimeMs: mtimeOf(file.path), + text: textOf(file.path), + })) + const envBlock = env ? [await environment(root, isRepo)] : [] + /** The prompt with each file as it is now. */ + const current = () => [ + ...files + .filter((file) => file.text !== '') + .map((file) => `${file.title}:\n${file.text}`), + ...envBlock, + ] + const sections = current() + const slots = files.length + envBlock.length + SPARE_SLOTS + + /** Add a note for each file that changed since the last check. */ + const refresh = () => { + const notes: Array = [] + for (const file of files) { + const mtimeMs = mtimeOf(file.path) + if (mtimeMs === file.mtimeMs) continue + const before = file.text + file.mtimeMs = mtimeMs + file.text = textOf(file.path) + if (file.text !== before) + notes.push(changeNote(file.title, before, file.text)) + } + // ponytail: with no spare slot left, the prompt is built again. That + // one change is not added to the conversation. + if (sections.length + notes.length > slots) { + sections.splice(0, sections.length, ...current()) + } else { + sections.push(...notes) + } + } + + /** Nested instruction files the model already got. */ + const given = new Set() + const afterRead = async (path: string) => { + const inside = relative(root, path) + const isBelowRoot = + inside !== '' && !inside.startsWith('..') && !isAbsolute(inside) + if (!isBelowRoot) return undefined + const notes: Array = [] + // The folders below `root`, down to the folder of the file. + const parts = inside.split(sep).slice(0, -1) + let dir = root + for (const part of parts) { + dir = join(dir, part) + for (const name of names) { + const file = join(dir, name) + if (given.has(file)) continue + const text = textOf(file) + if (text === '') continue + given.add(file) + // The model just read this file, so it has the text. + if (file !== path) { + notes.push(`Instructions from ${shownPath(file)}:\n${text}`) + } + } + } + return notes.length > 0 ? notes.join('\n\n') : undefined + } + + return { + // One prompt per slot. Empty slots are left out. The harness reads + // the prompts in order each turn, so slot 0 checks the files first. + prompts: Array.from({ length: slots }, (_, index) => () => { + if (index === 0) refresh() + return sections[index] ?? '' + }), + contribute: [WorkspaceHooks.item({ afterRead })], + } + }, + }) +} + +/** A frontmatter value: text, or a list from `[a, b]` or `- a` lines. */ +export type FrontmatterValue = string | Array + +/** The text without the quotes around it: `"a"` and `'a'` become `a`. */ +function unquote(value: string) { + return value.trim().replace(/^(["'])(.*)\1$/, '$2') +} + +/** + * Split `---` frontmatter from a Markdown file. It reads `key: value` + * lines and lists: `key: [a, b]`, or `key:` with `- a` lines under it. + * Nested values are not read. + */ +export function splitFrontmatter(text: string) { + const fields: Record = {} + const match = /^---\r?\n([\s\S]*?)\r?\n---\r?\n?/.exec(text) + if (!match) return { fields, body: text } + // The list that `- a` lines add to: the last `key:` with no value. + let list: Array | undefined + const lines = (match[1] ?? '').split(/\r?\n/) + for (const line of lines) { + const item = /^\s*-\s+(.*)$/.exec(line) + if (item && list) { + list.push(unquote(item[1] ?? '')) + continue + } + list = undefined + const pair = /^([\w-]+):(.*)$/.exec(line) + if (!pair) continue + const key = pair[1] ?? '' + const value = (pair[2] ?? '').trim() + const inline = /^\[(.*)\]$/.exec(value) + if (value === '') { + list = [] + fields[key] = list + } else if (inline) { + fields[key] = (inline[1] ?? '') + .split(',') + .map(unquote) + .filter((entry) => entry !== '') + } else { + fields[key] = unquote(value) + } + } + return { fields, body: text.slice(match[0].length) } +} + +/** + * The Markdown files in `dir`: the name without `.md`, the frontmatter + * fields, and the body. A folder that cannot be read has no files. + */ +export async function readMarkdownFiles(dir: string) { + let files: Array = [] + try { + files = (await readdir(dir)).filter((file) => extname(file) === '.md') + } catch { + return [] + } + return Promise.all( + files.map(async (file) => ({ + name: basename(file, '.md'), + ...splitFrontmatter(await readFile(resolve(dir, file), 'utf8')), + })), + ) +} + +/** + * One command per Markdown file in `dir` (like `.claude/commands/*.md`). + * Running `/name args` sends the file as a prompt, with `$ARGUMENTS` + * replaced by the args. + */ +export function fileCommands(options: { dir: string }) { + return definePlugin({ + name: 'tanstack/file-commands', + setup: async (ctx) => { + const files = await readMarkdownFiles(options.dir) + const commands: Record = {} + for (const { name, fields, body } of files) { + const { description } = fields + commands[name] = defineCommand({ + description: + typeof description === 'string' && description !== '' + ? description + : `Run ${name}.md`, + run: (input: unknown) => { + const args = + typeof input === 'string' + ? input + : input === undefined + ? '' + : JSON.stringify(input) + ctx.session.prompt(body.split('$ARGUMENTS').join(args).trim()) + return `Sent /${name}.` + }, + }) + } + return { commands } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/glob.ts b/packages/ai-harness/src/first-party/glob.ts new file mode 100644 index 0000000000..8dc2b9e857 --- /dev/null +++ b/packages/ai-harness/src/first-party/glob.ts @@ -0,0 +1,14 @@ +/** A simple glob to RegExp: `**` any path, `*` any name part, `?` one character. */ +export function globToRegExp(glob: string) { + let pattern = '' + for (let index = 0; index < glob.length; index += 1) { + const char = glob.charAt(index) + if (char === '*' && glob[index + 1] === '*') { + pattern += '.*' + index += glob[index + 2] === '/' ? 2 : 1 + } else if (char === '*') pattern += '[^/]*' + else if (char === '?') pattern += '[^/]' + else pattern += char.replace(/[.+^${}()|[\]\\]/g, '\\$&') + } + return new RegExp(`^${pattern}$`) +} diff --git a/packages/ai-harness/src/first-party/index.ts b/packages/ai-harness/src/first-party/index.ts new file mode 100644 index 0000000000..374deb8d06 --- /dev/null +++ b/packages/ai-harness/src/first-party/index.ts @@ -0,0 +1,31 @@ +// First-party plugins. This entry is Node only: some plugins use the file +// system and the shell. +export { modelPicker } from './model-picker' +export { + PERMISSION_MODES, + PermissionDecisionCapability, + PermissionResources, + PermissionRules, + decidePermission, + deleteSavedPermission, + isUnsplittableCommand, + listSavedPermissions, + permissions, +} from './permissions' +export type { + CallResources, + PermissionDecision, + PermissionMode, + PermissionRule, + ToolResources, +} from './permissions' +export { globToRegExp } from './glob' +export { formatTodos, todos } from './todos' +export type { Todo } from './todos' +export { fileCommands, projectInstructions } from './files' +export { compact, usage } from './session-tools' +export { agents, builtInAgents } from './agents' +export type { AgentProfile } from './agents' +export { question } from './question' +export { TitleFailed, title } from './title' +export { boundToolOutput } from './bound-output' diff --git a/packages/ai-harness/src/first-party/model-picker.ts b/packages/ai-harness/src/first-party/model-picker.ts new file mode 100644 index 0000000000..30261ae7f8 --- /dev/null +++ b/packages/ai-harness/src/first-party/model-picker.ts @@ -0,0 +1,68 @@ +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { definePlugin } from '../plugins' +import type { AnyTextAdapter, KeyedAdapter } from '@tanstack/ai' + +/** + * Switch the main model at the next turn with the `model` setting or the + * `/model ` command. A choice can be a `keyedAdapter(...)`: the + * session builds it for each turn with the user's key. + * + * @example + * ```ts + * modelPicker({ + * choices: { + * fast: openaiText('gpt-5.6-luna'), + * claude: keyedAdapter(anthropicByok, (key) => + * createAnthropicChat('claude-sonnet-4-5', key), + * ), + * }, + * default: 'fast', + * }) + * ``` + */ +export function modelPicker(options: { + choices: Record> + default?: string +}) { + const names = Object.keys(options.choices) + const fallback = options.default ?? names[0] + if (fallback === undefined) + throw new Error('modelPicker needs at least one choice.') + return definePlugin({ + name: 'tanstack/model-picker', + setup: (ctx) => ({ + config: { + model: configOption.select({ + options: names, + default: fallback, + category: 'model', + description: 'The main model', + }), + }, + adapter: () => { + const name = ctx.config.get('model') + return typeof name === 'string' ? options.choices[name] : undefined + }, + commands: { + model: defineCommand({ + description: `Show or switch the model (${names.join(', ')})`, + run: async (input: unknown) => { + const name = + typeof input === 'string' + ? input + : typeof input === 'object' && input !== null && 'name' in input + ? String(input.name) + : undefined + if (!name) + return `Model: ${String(ctx.config.get('model'))}. Choices: ${names.join(', ')}.` + if (!names.includes(name)) + return `Unknown model "${name}". Choices: ${names.join(', ')}.` + await ctx.session.setConfig('model', name) + return `Model: ${name}. It applies at the next turn.` + }, + }), + }, + }), + }) +} diff --git a/packages/ai-harness/src/first-party/permissions.ts b/packages/ai-harness/src/first-party/permissions.ts new file mode 100644 index 0000000000..e8dd018975 --- /dev/null +++ b/packages/ai-harness/src/first-party/permissions.ts @@ -0,0 +1,728 @@ +import { createCapability, getMetadata } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { configOption } from '../config' +import { createExtensionPoint } from '../extensions' +import { SessionMetadata, definePlugin } from '../plugins' +import { isRecord } from '../utils' +import { globToRegExp } from './glob' +import type { AnyChatMiddleware, JSONSchema, MetadataStore } from '@tanstack/ai' + +export type PermissionDecision = 'allow' | 'ask' | 'deny' + +/** A rule for one tool. `kind` lets the modes treat edits and commands differently. */ +export interface PermissionRule { + /** A tool name. A trailing `*` matches a prefix, for example `git_*`. */ + tool: string + /** + * A glob on what the call touches: a path for file tools (`src/**`, + * relative to the root), or a command part for shell tools (`git *`). + * Without it, the rule matches every call of the tool. + */ + resource?: string + decision: PermissionDecision + kind?: 'read' | 'edit' | 'execute' +} + +/** + * Tool plugins add their rules here. `permissions()` reads them, and adds its + * own `rules` too, so other readers like `codeMode()` see them. + */ +export const PermissionRules = createExtensionPoint( + 'tanstack/permission-rules', +) + +/** + * What the calls of one tool touch. Each function gets the raw input the + * model sent, before the schema check. A function that throws refuses the + * call. + */ +export interface ToolResources { + /** The files the call reads or writes, relative to the root or absolute. */ + paths?: (input: unknown) => ReadonlyArray + /** The shell commands the call runs, one part for each simple command. */ + commands?: (input: unknown) => ReadonlyArray +} + +/** + * Tool plugins say what each call touches, so rules with a `resource` can + * match it. Each item maps tool names to their resources, for example + * `PermissionResources.item({ read_file: { paths: (input) => [...] } })`. + */ +export const PermissionResources = createExtensionPoint< + Readonly> +>('tanstack/permission-resources') + +export const PERMISSION_MODES = [ + 'default', + 'plan', + 'acceptEdits', + 'bypass', +] as const +export type PermissionMode = (typeof PERMISSION_MODES)[number] + +/** + * What `permissions()` decides for a call of `tool` in `mode`. With + * `resources`, for a call that touches them. Without, for a call that + * declares no `PermissionResources`. It uses all rules in their order, the + * `default`, and the `root`. Other plugins, like `codeMode()` and + * `workspaceTools()`, ask it. Saved answers are not used: they only allow + * more, so an answer here is never looser than a call. + */ +export const PermissionDecisionCapability = createCapability< + ( + tool: string, + mode: PermissionMode, + resources?: CallResources, + ) => PermissionDecision +>()('tanstack/permission-decision') + +/** + * @internal Ask the rules of `permissions()` about `tool`, an action that is + * not a tool call, like `formatter:prettier`. It uses the saved answers. On + * `ask`, it asks the user with `message`, and saves an `always` answer. + * Resolves to `true` when the action can run. The package does not export it. + */ +export const PermissionPrompt = createCapability< + (tool: string, message: string) => Promise +>()('tanstack/permission-prompt') + +/** + * True when a command part still has shell syntax that can run or chain + * more commands (`$(`, a backtick, `;`, `&`, `|`, `<`, `>`, a new line), so + * no allow rule applies to it. + */ +export function isUnsplittableCommand(part: string) { + return /[`;&|<>\r\n]|\$\(/.test(part) +} + +/** What one call touches, as a tool's `ToolResources` returns it. */ +export interface CallResources { + paths?: ReadonlyArray + commands?: ReadonlyArray +} + +/** One thing a call touches, ready to match rule globs against. */ +interface Resource { + type: 'path' | 'command' + /** The text rules match: a path relative to the root, or a command part. */ + match: string + /** Asks even when a rule without a `resource` allows the tool. */ + risky: boolean + /** No allow rule applies. */ + unsplittable: boolean + /** Compare without letter case (a Windows path). */ + foldCase: boolean +} + +function matchesTool(rule: PermissionRule, tool: string) { + return rule.tool.endsWith('*') + ? tool.startsWith(rule.tool.slice(0, -1)) + : rule.tool === tool +} + +/** A path or a drive root in Windows form: `C:\x`, `C:/x`, or `\\server\share`. */ +function isWindowsPath(path: string) { + return /^([a-z]:[\\/]|[\\/]{2})/i.test(path) +} + +/** + * Compare paths without letter case on macOS and Windows, whose file systems + * ignore case by default, and for any Windows-form path. Else `Secrets/key` + * would pass a deny rule for `secrets/**`. + */ +// ponytail: platform default, add a `caseInsensitive` option if a case-sensitive macOS volume needs it +function foldsCase(path: string) { + const platform = typeof process === 'undefined' ? '' : process.platform + return platform === 'darwin' || platform === 'win32' || isWindowsPath(path) +} + +/** `path` with `/` separators, and `.` and `..` resolved. */ +function normalizePath(path: string) { + const slashed = path.replaceAll('\\', '/') + const prefix = + /^[a-z]:\//i.exec(slashed)?.[0] ?? (slashed.startsWith('/') ? '/' : '') + const segments: Array = [] + const parts = slashed.slice(prefix.length).split('/') + for (const segment of parts) { + if (segment === '' || segment === '.') continue + const climbs = + segment === '..' && segments.length > 0 && segments.at(-1) !== '..' + if (climbs) segments.pop() + // `..` above an absolute root stays at the root. + else if (segment !== '..' || prefix === '') segments.push(segment) + } + return prefix + segments.join('/') +} + +/** A `.env` file, or `.env.local` and the like. They often hold secrets. */ +function isEnvFile(path: string) { + const name = path.slice(path.lastIndexOf('/') + 1).toLowerCase() + return /\.env(\.|$)/.test(name) +} + +function pathResource(path: string, root: string | undefined) { + const foldCase = foldsCase(root ?? path) + const slashed = path.replaceAll('\\', '/') + // Forms we cannot resolve without the file system ask: `C:x` (relative to + // the current folder of drive C), a stream (`a.txt:secret`), and a name + // that ends in `.` or a space (Windows drops them). + const cannotResolve = + slashed.replace(/^[a-z]:\//i, '').includes(':') || + slashed + .split('/') + .some((part) => /[. ]$/.test(part) && part !== '.' && part !== '..') + const isAbsolute = /^([a-z]:)?\//i.test(slashed) + const full = normalizePath( + isAbsolute || root === undefined ? slashed : `${root}/${slashed}`, + ) + const resource = { + type: 'path' as const, + match: full, + risky: cannotResolve || isEnvFile(full), + unsplittable: false, + foldCase, + } + if (root === undefined) return resource + const base = normalizePath(root) + const folder = base.endsWith('/') ? base : `${base}/` + const fold = (value: string) => (foldCase ? value.toLowerCase() : value) + const isInside = + fold(full) === fold(base) || fold(full).startsWith(fold(folder)) + // Inside the root, rules match the relative path. Outside, the full path. + return isInside + ? { ...resource, match: full.slice(folder.length) } + : { ...resource, risky: true } +} + +function commandResource(part: string) { + return { + type: 'command' as const, + match: part.trim(), + risky: false, + unsplittable: isUnsplittableCommand(part), + foldCase: false, + } +} + +/** A command glob: `*` matches any text, `/` too. */ +function commandPattern(glob: string) { + const source = glob + .split('*') + .map((part) => part.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')) + .join('.*') + return new RegExp(`^${source}$`) +} + +function matchesResource(glob: string, resource: Resource) { + if (resource.type === 'command') { + return commandPattern(glob).test(resource.match) + } + const pattern = globToRegExp(glob.replaceAll('\\', '/')) + return new RegExp(pattern.source, resource.foldCase ? 'i' : '').test( + resource.match, + ) +} + +function strictest(decisions: ReadonlyArray) { + if (decisions.includes('deny')) return 'deny' + if (decisions.includes('ask')) return 'ask' + return 'allow' +} + +/** + * The decision for a tool call in a mode. The last matching rule wins. + * + * - With `resources`, each path and command part is checked, and every one + * must be allowed. A path outside `root`, a `.env` file, or a path that + * cannot be resolved asks, unless a rule with a `resource` allows it. A + * command part with shell syntax (see `isUnsplittableCommand`) never gets + * an allow. + * - Without `resources`, the call could touch anything, so a later rule + * with a `resource` that asks or denies counts too. + * + * Modes: `bypass` allows all. `plan` denies edits, commands, and every call + * that is not an allow. `acceptEdits` allows edits that would ask. + */ +export function decidePermission( + rules: ReadonlyArray, + tool: string, + mode: PermissionMode, + options: { + fallback?: PermissionDecision + resources?: CallResources + /** The workspace root. Without it, no path counts as outside. */ + root?: string + } = {}, +) { + if (mode === 'bypass') return 'allow' + // An allow rule without `kind` must not hide that a tool edits. + const kind = rules.findLast( + (rule) => matchesTool(rule, tool) && rule.kind !== undefined, + )?.kind + if (mode === 'plan' && (kind === 'edit' || kind === 'execute')) return 'deny' + const decisionOf = (rule: PermissionRule | undefined) => { + const decision = rule?.decision ?? options.fallback ?? 'allow' + const autoEdit = + mode === 'acceptEdits' && + kind === 'edit' && + decision === 'ask' && + rule?.resource === undefined + return autoEdit ? 'allow' : decision + } + const toolRules = rules.filter((rule) => matchesTool(rule, tool)) + const lastToolLevel = toolRules.findLastIndex( + (rule) => rule.resource === undefined, + ) + const toolLevel = lastToolLevel === -1 ? undefined : toolRules[lastToolLevel] + const decideOne = (resource: Resource) => { + const rule = toolRules.findLast( + (candidate) => + candidate.resource === undefined || + matchesResource(candidate.resource, resource), + ) + const decision = decisionOf(rule) + if (decision !== 'allow') return decision + const needsAsk = + resource.unsplittable || (resource.risky && rule?.resource === undefined) + return needsAsk ? 'ask' : 'allow' + } + const resources = + options.resources === undefined + ? undefined + : [ + ...(options.resources.paths ?? []).map((path) => + pathResource(path, options.root), + ), + ...(options.resources.commands ?? []).map(commandResource), + ] + // Unknown resources: the rules with a `resource` after the tool rule could + // match, so the strictest of them counts. + const decisions = + resources === undefined + ? [ + decisionOf(toolLevel), + ...toolRules.slice(lastToolLevel + 1).map((rule) => rule.decision), + ] + : resources.length === 0 + ? [decisionOf(toolLevel)] + : resources.map(decideOne) + const decision = strictest(decisions) + if (mode === 'plan' && decision !== 'allow') return 'deny' + return decision +} + +/** Rules hide a tool when its last rule denies with no `resource`. */ +function isHidden( + rules: ReadonlyArray, + tool: string, + mode: PermissionMode, +) { + if (mode === 'bypass') return false + const last = rules.findLast((rule) => matchesTool(rule, tool)) + return last?.decision === 'deny' && last.resource === undefined +} + +/** The answer to a permission question, for hosts that render a form. */ +const ANSWER_SCHEMA: JSONSchema = { + type: 'object', + properties: { + answer: { + type: 'string', + enum: ['once', 'always', 'reject'], + description: + 'once: allow this call. always: allow it from now on in this project. reject: refuse it.', + }, + message: { + type: 'string', + description: 'With reject: what the model should do instead.', + }, + }, + required: ['answer'], +} + +/** The user's reply. Anything that is not `once` or `always` rejects. */ +function readAnswer(value: unknown) { + const raw = isRecord(value) ? value.answer : value + const choice = typeof raw === 'string' ? raw.trim().toLowerCase() : '' + const text = + isRecord(value) && typeof value.message === 'string' + ? value.message.trim() + : '' + return { + choice: choice === 'once' || choice === 'always' ? choice : 'reject', + message: text === '' ? undefined : text, + } +} + +function isSavedRule(value: unknown): value is PermissionRule { + return ( + isRecord(value) && + typeof value.tool === 'string' && + value.decision === 'allow' && + (value.resource === undefined || typeof value.resource === 'string') + ) +} + +/** A refused call: the model gets a tool error with `error`. */ +function refuse(error: string) { + return { type: 'skip' as const, result: { error } } +} + +const METADATA_NAMESPACE = 'tanstack/permissions' + +/** The root as one string for each folder: `C:\Repo` and `c:/repo/` are the same. */ +function projectId(root: string) { + const path = normalizePath(root) + return foldsCase(root) ? path.toLowerCase() : path +} + +/** Where the metadata store keeps the saved rules of `project`. */ +function savedKeyOf(project: string | undefined) { + return project === undefined + ? 'permissions:saved' + : `permissions:saved:${projectId(project)}` +} + +async function readSaved(metadata: MetadataStore, key: string) { + const stored = await metadata.get(METADATA_NAMESPACE, key) + return Array.isArray(stored) ? stored.filter(isSavedRule) : [] +} + +function isSameRule(a: PermissionRule, b: PermissionRule) { + return ( + a.tool === b.tool && + a.resource === b.resource && + a.decision === b.decision && + a.kind === b.kind + ) +} + +/** + * The rules that `always` answers saved for `project`, in the saved order. + * `project` is the `root` of `permissions()`. `stores` is the persistence + * of the host, for example `persistence.stores`. Without a metadata store, + * it resolves to an empty list. + */ +export async function listSavedPermissions( + stores: { metadata?: MetadataStore }, + project: string | undefined, +) { + return stores.metadata ? readSaved(stores.metadata, savedKeyOf(project)) : [] +} + +/** + * Delete one saved rule of `project`. The rule matches by its fields, as + * `listSavedPermissions` gives them. An unknown rule does nothing. New + * sessions stop using the rule. Open sessions can keep it until they open + * again. + */ +export async function deleteSavedPermission( + stores: { metadata?: MetadataStore }, + project: string | undefined, + rule: PermissionRule, +) { + const { metadata } = stores + if (!metadata) return + const key = savedKeyOf(project) + const rules = await readSaved(metadata, key) + // Remove every copy, so a forgotten rule never stays in force. + const kept = rules.filter((saved) => !isSameRule(saved, rule)) + if (kept.length === rules.length) return + await metadata.set(METADATA_NAMESPACE, key, kept) +} + +/** A saved rule as one line: the tool, then the resource. */ +function describeRule(rule: PermissionRule) { + return rule.resource === undefined + ? rule.tool + : `${rule.tool} ${rule.resource}` +} + +const PLAN_PROMPT = + 'You are in plan mode. Do not change files or run commands. Read what you need, then describe your plan.' + +/** + * Check every tool call against permission rules, with a `mode` setting and a + * `/mode` command: + * + * - `default`: rules apply as written. `ask` asks the user. + * - `plan`: read-only. Edits and commands are denied. + * - `acceptEdits`: edits run without asking. Commands still ask. + * - `bypass`: everything runs. + * + * Rules apply in this order, and the last match wins: rules from tool + * plugins (`PermissionRules`), then `rules`, then saved answers. A saved + * answer never beats a deny. Tool plugins say what a call touches through + * `PermissionResources`. A path outside `root` or a `.env` file asks, unless + * a rule with a `resource` allows it. A tool whose last rule denies with no + * `resource` is removed from the model's tools. + * + * A question takes `once`, `always`, or `reject` with a `message` for the + * model. `always` saves an allow rule for each path or command part of the + * call, for the project (`root`), in the metadata store, else in memory for + * the session. `/permissions` lists the saved rules, and + * `/permissions forget ` deletes one. + */ +export function permissions( + options: { + rules?: ReadonlyArray + default?: PermissionDecision + /** + * The workspace root. A path outside it asks, and saved answers belong + * to it. Without it, no path counts as outside. + */ + root?: string + } = {}, +) { + return definePlugin({ + name: 'tanstack/permissions', + provides: [PermissionDecisionCapability, PermissionPrompt], + setup: (ctx) => { + const contributed = ctx.collect(PermissionRules) + const declared = ctx.collect(PermissionResources) + // The rules of tool plugins, then the user's rules, so the user's win. + const configured = () => [...contributed, ...(options.rules ?? [])] + ctx.provide(PermissionDecisionCapability, (tool, current, resources) => + decidePermission(configured(), tool, current, { + fallback: options.default, + resources, + root: options.root, + }), + ) + const mode = (): PermissionMode => { + const value = ctx.config.get('mode') + return ( + PERMISSION_MODES.find((candidate) => candidate === value) ?? 'default' + ) + } + const root = options.root + const savedKey = savedKeyOf(root) + // Answers saved in this session. Runs without a metadata store use them. + const memory: Array = [] + const loadSaved = async (metadata: MetadataStore | undefined) => { + const rules = metadata ? await readSaved(metadata, savedKey) : [] + return [...rules, ...memory] + } + // ponytail: get + set, a parallel save can drop a rule. It then asks again. + const save = async ( + metadata: MetadataStore | undefined, + rules: ReadonlyArray, + ) => { + memory.push(...rules) + if (!metadata || rules.length === 0) return + const current = await readSaved(metadata, savedKey) + const added = rules.filter( + (rule) => !current.some((saved) => isSameRule(saved, rule)), + ) + if (added.length === 0) return + await metadata.set(METADATA_NAMESPACE, savedKey, [...current, ...added]) + } + // The session's store, for commands. A chat run reads its own. + const sessionMetadata = ctx.getOptional(SessionMetadata) + const savedRules = async () => + sessionMetadata ? readSaved(sessionMetadata, savedKey) : [...memory] + const forget = async (arg: string) => { + const rules = await savedRules() + const rule = /^\d+$/.test(arg) ? rules[Number(arg) - 1] : undefined + if (!rule) { + return `No saved rule ${arg}. Run /permissions to see the list.` + } + await deleteSavedPermission({ metadata: sessionMetadata }, root, rule) + // This session stops using the rule at once. + const kept = memory.filter((saved) => !isSameRule(saved, rule)) + memory.splice(0, memory.length, ...kept) + return `Forgot rule ${arg}: ${describeRule(rule)}.` + } + const resourcesOf = (tool: string, input: unknown) => { + const entry = declared.findLast((item) => Object.hasOwn(item, tool))?.[ + tool + ] + if (!entry) return undefined + const resources = { + paths: entry.paths?.(input) ?? [], + commands: entry.commands?.(input) ?? [], + } + const values = [...resources.paths, ...resources.commands] + if (!values.every((value) => typeof value === 'string')) { + throw new Error('A resource is not a string.') + } + return resources + } + /** The rules an `always` answer saves: one allow for each resource. */ + const rulesToSave = ( + tool: string, + resources: CallResources | undefined, + ) => { + if (resources === undefined) { + return [{ tool, decision: 'allow' as const }] + } + const values = [ + ...(resources.paths ?? []).map( + (path) => pathResource(path, root).match, + ), + ...(resources.commands ?? []) + .filter((part) => !isUnsplittableCommand(part)) + .map((part) => part.trim()), + ] + // A `*` or `?` would turn the saved rule into a wider glob. + return values + .filter((value) => !/[*?]/.test(value)) + .map((resource) => ({ tool, resource, decision: 'allow' as const })) + } + const decide = async ( + tool: string, + input: unknown, + current: PermissionMode, + metadata: MetadataStore | undefined, + ) => { + const resources = resourcesOf(tool, input) + const settings = { fallback: options.default, resources, root } + const rules = configured() + // A saved answer turns an ask into an allow. It never beats a deny, + // so only an ask reads the saved answers. + const first = decidePermission(rules, tool, current, settings) + if (first !== 'ask') return { decision: first, resources } + rules.push(...(await loadSaved(metadata))) + const decision = decidePermission(rules, tool, current, settings) + return { decision, resources } + } + /** Ask the user. An `always` answer saves the rules for the call. */ + const askUser = async ( + message: string, + tool: string, + resources: CallResources | undefined, + metadata: MetadataStore | undefined, + ) => { + // A closed session rejects the question: that is a reject too. + const reply = readAnswer( + await ctx.session + .ask({ + message: `${message} Answer once, always, or reject.`, + schema: ANSWER_SCHEMA, + }) + .catch(() => undefined), + ) + if (reply.choice === 'always') { + // The session keeps the answer in memory, so a failed write does + // not fail the call the user allowed. + await save(metadata, rulesToSave(tool, resources)).catch(() => {}) + } + return reply + } + ctx.provide(PermissionPrompt, async (tool, message) => { + const { decision } = await decide( + tool, + undefined, + mode(), + sessionMetadata, + ) + if (decision !== 'ask') return decision === 'allow' + const reply = await askUser(message, tool, undefined, sessionMetadata) + return reply.choice !== 'reject' + }) + const check = { + name: 'tanstack/permissions', + onBeforeToolCall: async (run, hook) => { + const current = mode() + if (current === 'bypass') return undefined + const metadata = getMetadata(run, { optional: true }) + const checked = await decide( + hook.toolName, + hook.args, + current, + metadata, + ).catch((error: unknown) => + error instanceof Error ? error : new Error(String(error)), + ) + // Fail closed: a call we cannot check does not run. + if (checked instanceof Error) { + return refuse( + `Cannot check the permissions of this call: ${checked.message}`, + ) + } + const { decision, resources } = checked + if (decision === 'allow') return undefined + if (decision === 'deny') { + return refuse( + current === 'plan' + ? 'This tool is not allowed in plan mode.' + : 'A permission rule denies this tool.', + ) + } + const preview = JSON.stringify(hook.args ?? {}).slice(0, 300) + const reply = await askUser( + `Allow ${hook.toolName} ${preview}?`, + hook.toolName, + resources, + metadata, + ) + if (reply.choice === 'reject') { + return refuse(reply.message ?? 'The user denied this tool call.') + } + return undefined + }, + } satisfies AnyChatMiddleware + return { + config: { + mode: configOption.select({ + options: PERMISSION_MODES, + default: 'default', + category: 'mode', + description: 'How tool calls are approved', + }), + }, + prompts: [ + { + id: 'tanstack/permissions:plan', + text: () => (mode() === 'plan' ? PLAN_PROMPT : ''), + }, + ], + commands: { + mode: defineCommand({ + description: `Show or switch the mode (${PERMISSION_MODES.join(', ')})`, + run: async (input: unknown) => { + const next = typeof input === 'string' ? input.trim() : '' + if (!next) return `Mode: ${mode()}.` + const receipt = await ctx.session.setConfig('mode', next) + return receipt.status === 'rejected' + ? `Unknown mode "${next}". Modes: ${PERMISSION_MODES.join(', ')}.` + : `Mode: ${next}.` + }, + }), + permissions: defineCommand({ + description: + 'List the saved "always" rules: /permissions. Delete one: /permissions forget ', + run: async (input: unknown) => { + const arg = typeof input === 'string' ? input.trim() : '' + const forgetArg = /^forget\s+(.+)$/.exec(arg)?.[1] + if (forgetArg !== undefined) return forget(forgetArg.trim()) + if (arg !== '') { + return 'Use /permissions, or /permissions forget .' + } + const rules = await savedRules() + if (rules.length === 0) return 'No saved rules.' + const lines = rules.map( + (rule, index) => `${index + 1}. ${describeRule(rule)}`, + ) + return ['Saved rules:', ...lines].join('\n') + }, + }), + }, + prepareTools: ({ tools }) => { + const rules = configured() + const current = mode() + return tools.filter((tool) => !isHidden(rules, tool.name, current)) + }, + // Other readers, like codeMode(), must see the user's rules too. The + // checks above add them again at the end, so they still win there. + contribute: (options.rules ?? []).map((rule) => + PermissionRules.item(rule), + ), + // The lead turn and every agent run check the same rules. Each run + // gets only one of the two lists, so a call is checked once. + middleware: [check], + agentMiddleware: [check], + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/question.ts b/packages/ai-harness/src/first-party/question.ts new file mode 100644 index 0000000000..7b59088692 --- /dev/null +++ b/packages/ai-harness/src/first-party/question.ts @@ -0,0 +1,189 @@ +import { toolDefinition } from '@tanstack/ai' +import { definePlugin } from '../plugins' +import { isRecord } from '../utils' + +/** One question the model asks, as the tool input has it. */ +interface ToolQuestion { + question: string + header?: string + options: Array<{ label: string; description?: string }> + multiple?: boolean +} + +const NO_ANSWER = 'The user did not answer.' + +function isOption(value: unknown) { + return ( + isRecord(value) && + typeof value.label === 'string' && + (value.description === undefined || typeof value.description === 'string') + ) +} + +function isToolQuestion(value: unknown): value is ToolQuestion { + return ( + isRecord(value) && + typeof value.question === 'string' && + (value.header === undefined || typeof value.header === 'string') && + Array.isArray(value.options) && + value.options.every(isOption) && + (value.multiple === undefined || typeof value.multiple === 'boolean') + ) +} + +/** The text the user sees: the question, the numbered options, and how to answer. */ +function messageOf(item: ToolQuestion) { + const title = item.header ? `${item.header}: ${item.question}` : item.question + const options = item.options.map( + (option, index) => + `${index + 1}. ${option.label}${option.description ? ` - ${option.description}` : ''}`, + ) + const how = + item.multiple === true + ? 'Type one or more numbers or labels, split by commas, or your own answer.' + : 'Type a number or a label, or your own answer.' + return [title, ...options, how].join('\n') +} + +/** + * The answer shape, for hosts that render a form: one label, or a list of + * labels. Hosts that show a text box send the typed text, and the tool reads it. + */ +function answerSchema(item: ToolQuestion) { + const choice = { + type: 'string', + enum: item.options.map((option) => option.label), + } + return item.multiple === true + ? { type: 'array', items: choice, minItems: 1, uniqueItems: true } + : choice +} + +/** The option `text` names, by its label (any letter case) or its number. */ +function labelOf(text: string, options: ToolQuestion['options']) { + const lower = text.toLowerCase() + const byLabel = options.find( + (option) => option.label.trim().toLowerCase() === lower, + ) + if (byLabel) return byLabel.label + return /^\d+$/.test(text) ? options[Number(text) - 1]?.label : undefined +} + +/** + * The answer line for the model: the picked labels, or the user's own + * words. `undefined` when the reply has no text. + */ +function readAnswer(reply: unknown, item: ToolQuestion) { + const isSplit = item.multiple === true && typeof reply === 'string' + const parts = Array.isArray(reply) + ? reply + : isSplit + ? reply.split(',') + : [reply] + const texts = parts + .map((part) => + typeof part === 'string' || typeof part === 'number' + ? String(part).trim() + : '', + ) + .filter((text) => text !== '') + if (texts.length === 0) return undefined + const picked = texts + .map((text) => labelOf(text, item.options)) + .filter((label) => label !== undefined) + const isPicked = + picked.length === texts.length && + (item.multiple === true || picked.length === 1) + if (isPicked) return `Answer: ${[...new Set(picked)].join(', ')}` + const own = typeof reply === 'string' ? reply.trim() : texts.join(', ') + return `Own answer: ${own}` +} + +/** + * Add a `question` tool: the model asks the user one to four questions, + * each with two to six options, and waits for the answers. Each question + * goes to the user through `ctx.session.ask`, one at a time. The user picks + * by number or label, or types their own answer. With no answer, or when + * the session closes, the model gets a tool error. + * + * @example + * ```ts + * plugins: () => [question()] + * ``` + */ +export function question() { + return definePlugin({ + name: 'tanstack/question', + setup: (ctx) => ({ + tools: [ + toolDefinition({ + name: 'question', + description: + 'Ask the user one or more questions, each with options, and wait for the answers. Use it when you need a decision or a preference from the user to continue. The user can also type their own answer.', + inputSchema: { + type: 'object', + properties: { + questions: { + type: 'array', + minItems: 1, + maxItems: 4, + items: { + type: 'object', + properties: { + question: { + type: 'string', + description: 'The full question.', + }, + header: { + type: 'string', + description: 'A short label for the question.', + }, + options: { + type: 'array', + minItems: 2, + maxItems: 6, + items: { + type: 'object', + properties: { + label: { type: 'string' }, + description: { type: 'string' }, + }, + required: ['label'], + }, + }, + multiple: { + type: 'boolean', + description: 'The user can pick more than one option.', + }, + }, + required: ['question', 'options'], + }, + }, + }, + required: ['questions'], + }, + // Asking again after a crash has no side effect. + replay: 'safe', + }).server(async (args: unknown) => { + const items = + isRecord(args) && Array.isArray(args.questions) + ? args.questions + : [] + if (!items.every(isToolQuestion)) + throw new Error('Each question needs a question and options.') + const answers: Array = [] + for (const item of items) { + // A closed session rejects the question: that is no answer too. + const reply = await ctx.session + .ask({ message: messageOf(item), schema: answerSchema(item) }) + .catch(() => undefined) + const answer = readAnswer(reply, item) + if (answer === undefined) throw new Error(NO_ANSWER) + answers.push(`${item.question}\n${answer}`) + } + return answers.join('\n\n') + }), + ], + }), + }) +} diff --git a/packages/ai-harness/src/first-party/session-tools.ts b/packages/ai-harness/src/first-party/session-tools.ts new file mode 100644 index 0000000000..64565ce0e8 --- /dev/null +++ b/packages/ai-harness/src/first-party/session-tools.ts @@ -0,0 +1,272 @@ +import { chat } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import { emptyUsage } from '../usage' +import type { + AnyChatMiddleware, + AnyTextAdapter, + KeyedAdapter, + ModelMessage, +} from '@tanstack/ai' +import type { SessionUsage, UsageCounts } from '../usage' + +const SUMMARY_PROMPT = + 'Summarize the conversation so far for yourself. Keep decisions, open tasks, file names, and facts you still need. Leave out small talk.' + +export function textOf(message: ModelMessage) { + if (typeof message.content === 'string') return message.content + if (Array.isArray(message.content)) { + return message.content + .map((part) => (part.type === 'text' ? part.content : '')) + .join('') + } + return '' +} + +/** One `role: text` line per message. Messages with no text are left out. */ +export function transcriptText(messages: ReadonlyArray) { + return messages + .map((message) => `${message.role}: ${textOf(message)}`) + .filter((line) => !line.endsWith(': ')) + .join('\n') +} + +/** + * `/compact`: replace a long transcript with a summary, so later turns send + * fewer tokens. The summary is written by `adapter`. A `keyedAdapter(...)` + * is built with the user's key when `/compact` runs. + */ +export function compact(options: { + adapter: AnyTextAdapter | KeyedAdapter + keepLast?: number +}) { + return definePlugin({ + name: 'tanstack/compact', + setup: (ctx) => ({ + commands: { + compact: defineCommand({ + description: 'Summarize the conversation to save tokens', + run: async () => { + const messages = await ctx.session.transcript() + const keep = options.keepLast ?? 0 + if (messages.length <= keep + 2) + return 'The conversation is already short.' + // Never keep a tool result without the call before it: the + // provider rejects it. Move the cut back to that call. + let cut = messages.length - keep + while (cut > 0 && messages[cut]?.role === 'tool') cut-- + if (cut < 2) return 'The conversation is already short.' + const older = messages.slice(0, cut) + const transcript = transcriptText(older) + const { text: summary } = await chat({ + adapter: await ctx.keys.adapter(options.adapter), + messages: [ + { role: 'user', content: `${SUMMARY_PROMPT}\n\n${transcript}` }, + ], + stream: false, + }) + await ctx.session.replaceTranscript([ + { + role: 'user', + content: `Summary of our conversation so far:\n${summary}`, + }, + { + role: 'assistant', + content: 'Understood. I will continue from this summary.', + }, + ...messages.slice(cut), + ]) + return `Compacted ${older.length} messages into a summary.` + }, + }), + }, + }), + }) +} + +/** Prices in USD per 1M tokens. */ +interface ModelPrices { + input: number + output: number + cacheRead?: number + cacheWrite?: number +} + +/** The usage of an index entry before the host writes the tokens of a turn. */ +const NO_TOKENS = { + turns: 0, + promptTokens: 0, + completionTokens: 0, + totalTokens: 0, + cachedTokens: 0, + cacheWriteTokens: 0, +} + +/** The cost in USD of some model calls of one model. */ +function priceOf(prices: ModelPrices, counts: UsageCounts) { + const { cachedTokens, cacheWriteTokens } = counts + // `promptTokens` counts the cache parts too, and each part has its own price. + const uncached = counts.promptTokens - cachedTokens - cacheWriteTokens + const part = (price: number | undefined, tokens: number) => + ((price ?? 0) / 1_000_000) * tokens + return ( + part(prices.input, uncached) + + part(prices.output, counts.completionTokens) + + part(prices.cacheRead, cachedTokens) + + part(prices.cacheWrite, cacheWriteTokens) + ) +} + +/** One `/usage` line: the calls and tokens of `counts`. */ +function usageLine(counts: UsageCounts) { + // Only a provider that reports its prompt cache has these counts. + const hasCache = counts.cachedTokens > 0 || counts.cacheWriteTokens > 0 + const cache = hasCache + ? ` (${counts.cachedTokens} cache read, ${counts.cacheWriteTokens} cache write)` + : '' + return `${counts.calls} model calls, ${counts.promptTokens} input tokens${cache}, ${counts.completionTokens} output tokens, ${counts.totalTokens} total.` +} + +/** + * Show the token usage of the session (`session.usage()`): the lead turn and + * every agent run (subagents, background agents, and their children). + * `/usage` shows the totals, with the input tokens read from and written to + * the prompt cache, and a line for each model when there is more than one. + * The plugin state, for a UI, has a copy of `session.usage().total` and + * `contextTokens`, the size of the lead model's context at its latest call. + * + * A model call that has a cost from its provider keeps that cost. With + * `model`, the plugin also prices the calls that have no provider cost. + * `model` gets the model id and returns its prices in USD per 1M tokens, or + * `undefined` when it does not know the model. `/usage` shows the cost when + * a call has one, and says how many calls it could not price. When a call + * has a cost, the plugin writes the cost to `usage.cost` of the session + * index entry. + * + * @example + * ```ts + * const prices: Record = { + * 'claude-sonnet-5-5': { input: 3, output: 15 }, + * } + * const priceOf = (id: string) => (prices[id] ? { cost: prices[id] } : undefined) + * plugins: () => [usage({ model: priceOf })] + * ``` + */ +export function usage( + options: { + model?: (modelId: string) => { cost: ModelPrices } | undefined + } = {}, +) { + return definePlugin({ + name: 'tanstack/usage', + setup: (ctx) => { + /** + * A copy of the session totals, and `contextTokens`: the input tokens + * of the latest model call of the lead turn, so how much of the + * model's context the conversation fills now. + */ + const state = ctx.state({ ...emptyUsage().total, contextTokens: 0 }) + // The cost of the model calls of one model, else `undefined`. + const costOf = (model: string, counts: UsageCounts) => { + if (counts.cost !== undefined) return counts.cost + // The key is `provider/model`. The lookup gets the model id only. + const prices = options.model?.(model.slice(model.indexOf('/') + 1)) + return prices && priceOf(prices.cost, counts) + } + /** + * The cost of the session, and how many calls have no cost. `cost` is + * `undefined` when no call has one. + * ponytail: core sums by model. When a model has a provider cost, its + * calls with no provider cost add nothing. Upgrade: core keeps, by + * model, the counts of the calls with no provider cost. + */ + const sessionCost = (totals: SessionUsage) => { + let cost: number | undefined + let unpriced = 0 + for (const [model, counts] of Object.entries(totals.byModel)) { + const modelCost = costOf(model, counts) + if (modelCost === undefined) unpriced += counts.calls + else cost = (cost ?? 0) + modelCost + } + return { cost, unpriced } + } + const show = () => { + const totals = ctx.session.snapshot().usage + const { cost, unpriced } = sessionCost(totals) + const note = + unpriced > 0 + ? ` (cost unknown for ${unpriced} ${unpriced === 1 ? 'call' : 'calls'})` + : '' + const isPriced = options.model !== undefined || cost !== undefined + const priced = isPriced + ? ` Cost: $${(cost ?? 0).toFixed(4)}${note}.` + : '' + const lines = [`${usageLine(totals.total)}${priced}`] + const models = Object.entries(totals.byModel) + // One model has the same counts as the total. + if (models.length > 1) { + for (const [model, counts] of models) { + lines.push(`${model}: ${usageLine(counts)}`) + } + } + return lines.join('\n') + } + // The host writes the token totals to the index entry at the end of + // each turn. This plugin writes only `cost`, and keeps the tokens of + // the entry. + // ponytail: a host with a log counts a call when its log write lands, + // which can be after the run ends. Then the entry and the state get + // that call at the end of the next run or call. Upgrade: the host + // writes the cost at turn end. + const saveCost = async () => { + const { cost } = sessionCost(ctx.session.snapshot().usage) + if (cost === undefined) return + // A change function: it runs again on the new entry when another + // write landed first, so the tokens of that write stay. Before the + // first turn ends, the entry has no usage. + await ctx.session.updateEntry((entry) => ({ + usage: { ...(entry.usage ?? NO_TOKENS), cost }, + })) + } + // Copy the session totals to the state. Only a call of the lead turn + // sets the context size. + const mirror = (contextTokens?: number) => + state.update((current) => ({ + ...ctx.session.snapshot().usage.total, + contextTokens: contextTokens ?? current.contextTokens, + })) + // At the end of each run, the state gets the totals and the index entry + // gets the cost so far. The index is only a list of sessions, and the + // state only a copy: a failed write does not fail the run. The next run + // writes them again. + const save = async () => { + await Promise.all([mirror(), saveCost()]).catch(() => {}) + } + const atRunEnd = { + name: 'tanstack/usage', + onFinish: save, + onAbort: save, + onError: save, + } satisfies AnyChatMiddleware + return { + middleware: [ + { + ...atRunEnd, + onUsage: async (_run, info) => { + await mirror(info.promptTokens) + }, + } satisfies AnyChatMiddleware, + ], + // An agent's middleware runs before the host counts its call, so its + // run copies the totals only at its end. + agentMiddleware: [atRunEnd], + commands: { + usage: defineCommand({ + description: 'Show token usage for this session', + run: show, + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/title.ts b/packages/ai-harness/src/first-party/title.ts new file mode 100644 index 0000000000..929e97dad1 --- /dev/null +++ b/packages/ai-harness/src/first-party/title.ts @@ -0,0 +1,90 @@ +import { chat } from '@tanstack/ai' +import { createPluginEvent } from '../extensions' +import { definePlugin } from '../plugins' +import { textOf } from './session-tools' +import type { + AnyChatMiddleware, + AnyTextAdapter, + KeyedAdapter, + ModelMessage, +} from '@tanstack/ai' + +const TITLE_PROMPT = + 'Write a short title for a conversation that starts with the message below. Use at most 60 characters on one line, with no quotes. Answer with the title only.' + +/** Sent when the title call fails. The turn does not fail. */ +export const TitleFailed = createPluginEvent<{ message: string }>( + 'tanstack/title:failed', +) + +/** The first line of `answer`, with no quotes around it, at most 60 characters. */ +function cleanTitle(answer: string) { + const line = answer.trim().split('\n')[0] ?? '' + return line + .replace(/^["'`]+|["'`]+$/g, '') + .trim() + .slice(0, 60) +} + +/** + * Give each session a title. At the first turn of a session with no title, + * `adapter` writes a short title from the first user message, and the plugin + * saves it as the `title` of the session index entry. The call runs next to + * the turn: it does not slow the turn, and a failure does not fail it. A + * failure sends a {@link TitleFailed} event, and the next turn tries again. + * Without `stores.sessions`, the plugin does nothing. + * + * @example + * ```ts + * plugins: () => [title({ adapter: openaiText('gpt-5.4-nano') })] + * ``` + */ +export function title(options: { + adapter: AnyTextAdapter | KeyedAdapter +}) { + return definePlugin({ + name: 'tanstack/title', + setup: (ctx) => { + let isNaming = false + const name = async (messages: ReadonlyArray) => { + const first = messages.find((message) => message.role === 'user') + const entry = await ctx.session.entry() + // No sessions store, or the session has a title (also after a restart). + if (!first || !entry || entry.title) return + // ponytail: the first 2000 characters are enough for a title. + const message = textOf(first).slice(0, 2000) + const answer = await chat({ + adapter: await ctx.keys.adapter(options.adapter), + messages: [ + { role: 'user', content: `${TITLE_PROMPT}\n\n${message}` }, + ], + stream: false, + }) + const text = cleanTitle(answer.text) + if (text) await ctx.session.updateEntry({ title: text }) + } + return { + middleware: [ + { + name: 'tanstack/title', + onConfig: (run, config) => { + if (run.phase !== 'init' || isNaming) return + isNaming = true + // Not awaited: the title call never slows or fails the turn. + void name(config.messages) + .catch((error: unknown) => + ctx.emit(TitleFailed, { + message: + error instanceof Error ? error.message : String(error), + }), + ) + .finally(() => { + isNaming = false + }) + }, + } satisfies AnyChatMiddleware, + ], + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/todos.ts b/packages/ai-harness/src/first-party/todos.ts new file mode 100644 index 0000000000..eae2c2cce8 --- /dev/null +++ b/packages/ai-harness/src/first-party/todos.ts @@ -0,0 +1,100 @@ +import { toolDefinition } from '@tanstack/ai' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' + +export interface Todo { + text: string + status: 'pending' | 'in_progress' | 'done' +} + +function isTodo(value: unknown): value is Todo { + return ( + typeof value === 'object' && + value !== null && + 'text' in value && + typeof value.text === 'string' && + 'status' in value && + (value.status === 'pending' || + value.status === 'in_progress' || + value.status === 'done') + ) +} + +const MARK = { pending: '[ ]', in_progress: '[~]', done: '[x]' } as const + +export function formatTodos(items: ReadonlyArray): string { + return items.map((item) => `${MARK[item.status]} ${item.text}`).join('\n') +} + +/** + * A todo list the model keeps for multi-step work (like Claude Code's + * TodoWrite). The list is plugin state, so clients see it change and it + * survives restarts. `/todos` shows it. + */ +export function todos() { + return definePlugin({ + name: 'tanstack/todos', + setup: async (ctx) => { + const state = ctx.state<{ items: Array }>({ items: [] }) + // The prompt runs for each turn, so keep a copy of the list at hand. + let current = (await state.get()).items + return { + tools: [ + toolDefinition({ + name: 'todo_write', + description: + 'Replace the todo list. Use it for work with several steps. Mark one item in_progress at a time.', + inputSchema: { + type: 'object', + properties: { + todos: { + type: 'array', + items: { + type: 'object', + properties: { + text: { type: 'string' }, + status: { + type: 'string', + enum: ['pending', 'in_progress', 'done'], + }, + }, + required: ['text', 'status'], + }, + }, + }, + required: ['todos'], + }, + replay: 'safe', + }).server(async (args: unknown) => { + const list = + typeof args === 'object' && + args !== null && + 'todos' in args && + Array.isArray(args.todos) + ? args.todos + : [] + if (!list.every(isTodo)) + throw new Error('Each todo needs text and a status.') + current = (await state.update(() => ({ items: list }))).items + return formatTodos(current) || 'The todo list is empty.' + }), + ], + prompts: [ + { + id: 'tanstack/todos:list', + text: () => + current.length > 0 + ? `Current todo list:\n${formatTodos(current)}` + : '', + }, + ], + commands: { + todos: defineCommand({ + description: 'Show the todo list', + run: () => formatTodos(current) || 'The todo list is empty.', + }), + }, + } + }, + }) +} diff --git a/packages/ai-harness/src/first-party/workspace-hooks.ts b/packages/ai-harness/src/first-party/workspace-hooks.ts new file mode 100644 index 0000000000..17421b70fc --- /dev/null +++ b/packages/ai-harness/src/first-party/workspace-hooks.ts @@ -0,0 +1,35 @@ +import { createExtensionPoint } from '../extensions' + +/** + * Hooks for the workspace tools. Each runs inside the tool, while the tool + * holds the lock on that path. A hook that throws fails the tool call. + */ +export interface WorkspaceHooks { + /** + * Runs after a tool wrote the file at `path`, an absolute path. Return + * text to add to the result, or nothing to add nothing. + */ + afterWrite?: (path: string) => Promise + /** + * Runs after `read_file` read the file at `path`, an absolute path. Return + * text to add to the result, or `undefined` to add nothing. + */ + afterRead?: (path: string) => Promise +} + +/** + * Plugins add {@link WorkspaceHooks} here. The workspace tools run them. + * + * @example + * ```ts + * definePlugin({ + * name: 'acme/format', + * setup: () => ({ + * contribute: [WorkspaceHooks.item({ afterWrite: async (path) => format(path) })], + * }), + * }) + * ``` + */ +export const WorkspaceHooks = createExtensionPoint( + 'tanstack/workspace-hooks', +) diff --git a/packages/ai-harness/src/index.ts b/packages/ai-harness/src/index.ts index a6de481a21..41c2a5231b 100644 --- a/packages/ai-harness/src/index.ts +++ b/packages/ai-harness/src/index.ts @@ -49,6 +49,21 @@ export type { } from './plugins' export { createExtensionPoint, createPluginEvent } from './extensions' +// Portable (no Node imports), so edge-safe plugins such as code mode can read +// the permission rules and resources without the Node-only `./plugins` entry. +export { + PermissionDecisionCapability, + PermissionResources, + PermissionRules, + decidePermission, +} from './first-party/permissions' +export type { + CallResources, + PermissionDecision, + PermissionMode, + PermissionRule, + ToolResources, +} from './first-party/permissions' export type { ExtensionItem, ExtensionPoint, PluginEvent } from './extensions' export { checkConfigValue, configOption } from './config' diff --git a/packages/ai-harness/tests/agent-middleware.test.ts b/packages/ai-harness/tests/agent-middleware.test.ts index 2a33623009..89c43a0b20 100644 --- a/packages/ai-harness/tests/agent-middleware.test.ts +++ b/packages/ai-harness/tests/agent-middleware.test.ts @@ -1,7 +1,8 @@ import { describe, expect, it } from 'vitest' -import { EventType, defineAgent } from '@tanstack/ai' +import { EventType, defineAgent, toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { permissions, usage } from '../src/first-party' import { mockAdapter, text, toolCall } from './helpers' import type { AnyChatMiddleware, StreamChunk } from '@tanstack/ai' import type { AnyHarness, PluginLifetime } from '../src' @@ -177,4 +178,74 @@ describe('plugin agentMiddleware', () => { ]) await host.close() }) + + it('lets the usage() plugin count the model calls of every agent', async () => { + const researcher = parentAgent( + 'researcher', + textAgent('fetcher', 'fetched', 100), + [10, 20], + ) + const { host, session } = await open( + defineHarness({ + name: 'test/agent-middleware', + adapter: leadCalling('researcher', [1, 2]), + subagents: { agents: [researcher] }, + plugins: () => [usage()], + }), + ) + + await session.prompt('research') + + // The lead (1 + 2), the researcher (10 + 20), and the fetcher (100). + expect(await session.command('usage')).toBe( + '5 model calls, 133 input tokens, 0 output tokens, 133 total.', + ) + // The context size is the input of the lead's latest call, not an agent's. + expect(session.snapshot().plugins['tanstack/usage']).toMatchObject({ + contextTokens: 2, + }) + await host.close() + }) + + it('lets permissions() deny a write a subagent tries in plan mode', async () => { + const writes: Array = [] + const writeFile = toolDefinition({ + name: 'write_file', + description: 'Write a file', + }).server(async (args) => { + writes.push(args) + return 'written' + }) + const writerModel = mockAdapter([ + () => toolCall('write_file', { path: 'x.txt' }), + () => text('writer done'), + ]) + const writer = defineAgent({ + name: 'writer', + description: 'Writes files', + run: (ctx) => + ctx.chat({ adapter: writerModel.adapter, tools: [writeFile] }), + }) + const { host, session } = await open( + defineHarness({ + name: 'test/agent-middleware', + adapter: leadCalling('writer'), + subagents: { agents: [writer] }, + plugins: () => [ + permissions({ + rules: [{ tool: 'write_file', decision: 'ask', kind: 'edit' }], + }), + ], + }), + ) + await session.command('mode', 'plan') + + await session.prompt('write it') + + expect(writes).toEqual([]) + expect(JSON.stringify(writerModel.calls[1].messages)).toContain( + 'This tool is not allowed in plan mode.', + ) + await host.close() + }) }) diff --git a/packages/ai-harness/tests/agents.test.ts b/packages/ai-harness/tests/agents.test.ts new file mode 100644 index 0000000000..c167cbb76b --- /dev/null +++ b/packages/ai-harness/tests/agents.test.ts @@ -0,0 +1,446 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { z } from 'zod' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { permissions } from '../src/first-party' +import { agents } from '../src/first-party/agents' +import { splitFrontmatter } from '../src/first-party/files' +import { messageTexts, mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessPlugin } from '../src' +import type { AgentProfile } from '../src/first-party/agents' + +/** The names of the tools that ran, in order. */ +const ran: Array = [] + +beforeEach(() => { + ran.length = 0 +}) + +const tool = (name: string) => + toolDefinition({ + name, + description: `The ${name} tool`, + inputSchema: z.object({}), + }).server(async () => { + ran.push(name) + return name + }) + +/** + * Open thread `t` with the tools `read_file` and `write_file`. The main + * model starts each subagent with a tool named after it. + */ +async function open(adapter: AnyTextAdapter, plugins: Array) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/agents', + adapter, + tools: [tool('read_file'), tool('write_file')], + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +/** The names of the tools the model got on one call. */ +function toolNames(call: { tools: Array<{ name: string }> }) { + return call.tools.map((entry) => entry.name) +} + +/** The system prompts of one call, as one string. */ +function promptsOf(call: { systemPrompts: unknown }) { + return JSON.stringify(call.systemPrompts) +} + +describe('splitFrontmatter', () => { + it('reads values and both list forms', () => { + const file = [ + '---', + 'description: "Reviews: code"', + 'tools: [read_file, grep]', + 'models:', + ' - fast', + " - 'smart'", + '---', + 'Body', + ].join('\n') + expect(splitFrontmatter(file)).toEqual({ + fields: { + description: 'Reviews: code', + tools: ['read_file', 'grep'], + models: ['fast', 'smart'], + }, + body: 'Body', + }) + }) + + it('keeps a file without frontmatter as the body', () => { + expect(splitFrontmatter('Just text')).toEqual({ + fields: {}, + body: 'Just text', + }) + }) +}) + +describe('agents', () => { + it('offers build and plan as primary agents, and general and explore as subagents', async () => { + const { adapter, calls } = mockAdapter([() => text('hi')]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter }), + ]) + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan.', + ) + await session.prompt('hi') + + expect(toolNames(calls[0])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + expect(promptsOf(calls[0])).toContain('You are the build agent.') + await host.close() + }) + + it('switches to plan with /agent: the next turn gets its prompt and read-only tools', async () => { + const { adapter, calls } = mockAdapter([ + () => text('built'), + () => text('planned'), + ]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter }), + ]) + + await session.prompt('build it') + expect(await session.command('agent', 'plan')).toBe( + 'Agent: plan. It applies at the next turn.', + ) + await session.prompt('plan it') + + expect(toolNames(calls[0])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + expect(toolNames(calls[1])).toEqual(['read_file']) + expect(promptsOf(calls[1])).toContain('You are the plan agent.') + expect(promptsOf(calls[1])).not.toContain('You are the build agent.') + expect(await session.command('agent', 'nope')).toBe( + 'Unknown agent "nope". Agents: build, plan.', + ) + await host.close() + }) + + it('switches the primary agent with the agent setting', async () => { + const { adapter, calls } = mockAdapter([() => text('planned')]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter }), + ]) + + expect((await session.setConfig('agent', 'plan')).status).toBe('accepted') + // A subagent is not a primary agent. + expect((await session.setConfig('agent', 'general')).status).toBe( + 'rejected', + ) + await session.prompt('plan it') + + expect(toolNames(calls[0])).toEqual(['read_file']) + await host.close() + }) + + it('switches the model when the agent has one, and keeps the harness model when not', async () => { + const main = mockAdapter([() => text('from main')]) + const fast = mockAdapter([() => text('from fast')]) + const byModel = (model: string) => { + if (model !== 'fast-model') throw new Error(`Unknown model ${model}`) + return fast.adapter + } + const quick: AgentProfile = { + name: 'quick', + description: 'Answers fast', + mode: 'primary', + model: 'fast-model', + } + const { host, session } = await open(main.adapter, [ + agents({ adapter: byModel, agents: [quick] }), + ]) + + expect(await session.prompt('one')).toEqual({ text: 'from main' }) + await session.command('agent', 'quick') + expect(await session.prompt('two')).toEqual({ text: 'from fast' }) + await host.close() + }) + + it('makes the call after `steps` model calls with no tool calls, then stops', async () => { + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', {}, 'c1'), + () => toolCall('read_file', {}, 'c2'), + () => toolCall('read_file', {}, 'c3'), + () => text('never asked'), + ]) + const short: AgentProfile = { + name: 'short', + description: 'Two steps', + mode: 'primary', + steps: 2, + } + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, agents: [short], default: 'short' }), + ]) + + await session.prompt('go') + + expect(calls.map((call) => call.toolChoice)).toEqual([ + undefined, + undefined, + 'none', + ]) + expect(promptsOf(calls[1])).not.toContain('step limit') + expect(promptsOf(calls[2])).toContain('step limit') + await host.close() + }) + + it('refuses a tool call at the step limit when the provider ignores toolChoice', async () => { + // Bedrock Converse sends `none` as `auto` after tool calls, so the model + // can still call a tool on the last call. + const { adapter, calls } = mockAdapter([ + () => toolCall('read_file', {}, 'c1'), + () => toolCall('read_file', {}, 'c2'), + () => text('Here is what I have.'), + ]) + const short: AgentProfile = { + name: 'short', + description: 'One step', + mode: 'primary', + steps: 1, + } + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, agents: [short], default: 'short' }), + ]) + + await session.prompt('go') + + // The run ends after the last call, and only the first call ran. + expect(calls).toHaveLength(2) + expect(calls[1].toolChoice).toBe('none') + expect(ran).toEqual(['read_file']) + + // The next turn shows the model the refusal of the last call. + await session.prompt('and now?') + const refused = calls[2].messages.find( + (message: { role: string; toolCallId?: string }) => + message.role === 'tool' && message.toolCallId === 'c2', + ) + expect(JSON.stringify(refused?.content)).toContain('Step limit reached') + expect(ran).toEqual(['read_file']) + await host.close() + }) + + it('runs a subagent agent as a child of its tool', async () => { + const lead = mockAdapter([ + () => toolCall('explore', { prompt: 'Find the config' }), + () => text('Done.'), + ]) + const child = mockAdapter(() => text('It is in config.ts')) + const models: Array = [] + const { host, session } = await open(lead.adapter, [ + agents({ + adapter: (model) => { + models.push(model) + return child.adapter + }, + }), + ]) + + await session.prompt('Where is the config?') + + // Without its own model, the child uses the model of the main turn. + expect(models).toEqual(['test-model']) + expect(toolNames(child.calls[0])).toEqual(['read_file']) + expect(promptsOf(child.calls[0])).toContain('You are a fast search agent.') + expect(messageTexts(child.calls[0]).at(-1)).toBe('Find the config') + expect(JSON.stringify(lead.calls[1].messages)).toContain( + 'It is in config.ts', + ) + await host.close() + }) + + it('applies the permission rules of an agent only while it is the primary agent', async () => { + const { adapter, calls } = mockAdapter([ + () => text('one'), + () => text('two'), + ]) + const careful: AgentProfile = { + name: 'careful', + description: 'Never writes', + mode: 'primary', + permissions: [{ tool: 'write_file', decision: 'deny' }], + } + const { host, session } = await open(adapter, [ + permissions(), + agents({ adapter: () => adapter, agents: [careful] }), + ]) + + await session.command('agent', 'careful') + await session.prompt('one') + await session.command('agent', 'build') + await session.prompt('two') + + expect(toolNames(calls[0])).toEqual(['read_file', 'general', 'explore']) + expect(toolNames(calls[1])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + await host.close() + }) + + it('replaces a built-in agent with one of the same name', async () => { + const { adapter, calls } = mockAdapter([() => text('planned')]) + const plan: AgentProfile = { + name: 'plan', + description: 'My plan', + mode: 'primary', + system: 'My own plan prompt.', + } + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, agents: [plan] }), + ]) + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan.', + ) + await session.command('agent', 'plan') + await session.prompt('plan it') + + expect(toolNames(calls[0])).toEqual([ + 'read_file', + 'write_file', + 'general', + 'explore', + ]) + expect(promptsOf(calls[0])).toContain('My own plan prompt.') + await host.close() + }) + + it('drops the built-ins, and leaves hidden agents out of the list', async () => { + const { adapter, calls } = mockAdapter([() => text('hi')]) + const { host, session } = await open(adapter, [ + agents({ + adapter: () => adapter, + builtIns: false, + agents: [ + { name: 'solo', description: 'The only agent', mode: 'primary' }, + { + name: 'secret', + description: 'Not listed', + mode: 'primary', + hidden: true, + }, + ], + }), + ]) + + expect(await session.command('agent')).toBe('Agent: solo. Agents: solo.') + await session.prompt('hi') + + expect(toolNames(calls[0])).toEqual(['read_file', 'write_file']) + await host.close() + }) +}) + +describe('agents from Markdown files', () => { + let root = '' + beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-agents-')) + }) + afterEach(async () => { + await rm(root, { recursive: true, force: true }) + }) + + it('loads an agent with its frontmatter fields', async () => { + await writeFile( + join(root, 'reviewer.md'), + [ + '---', + 'description: Reviews code', + 'mode: all', + 'model: smart-model', + 'tools: [read_file]', + 'steps: 1', + '---', + 'You review code.', + ].join('\n'), + ) + const main = mockAdapter([() => text('built')]) + const smart = mockAdapter([ + () => toolCall('read_file', {}, 'c1'), + () => text('reviewed'), + ]) + const { host, session } = await open(main.adapter, [ + agents({ + adapter: (model) => + model === 'smart-model' ? smart.adapter : main.adapter, + dirs: [root], + }), + ]) + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan, reviewer.', + ) + await session.prompt('build it') + expect( + main.calls[0].tools.find( + (entry: { name: string }) => entry.name === 'reviewer', + )?.description, + ).toBe('Reviews code') + await session.command('agent', 'reviewer') + await session.prompt('review it') + + expect(toolNames(smart.calls[0])).toEqual(['read_file']) + expect(promptsOf(smart.calls[0])).toContain('You review code.') + expect(smart.calls.map((call) => call.toolChoice)).toEqual([ + undefined, + 'none', + ]) + await host.close() + }) + + it('reads the folder again on session.reload()', async () => { + const { adapter } = mockAdapter([]) + const { host, session } = await open(adapter, [ + agents({ adapter: () => adapter, dirs: [root] }), + ]) + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan.', + ) + + await writeFile(join(root, 'writer.md'), 'You write docs.') + await session.reload() + + expect(await session.command('agent')).toBe( + 'Agent: build. Agents: build, plan, writer.', + ) + await host.close() + }) + + it('refuses an agent file with an unknown mode', async () => { + await writeFile(join(root, 'bad.md'), '---\nmode: sometimes\n---\nText') + const { adapter } = mockAdapter([]) + + await expect( + open(adapter, [agents({ adapter: () => adapter, dirs: [root] })]), + ).rejects.toThrow('Agent file bad.md: "mode" is not valid.') + }) +}) diff --git a/packages/ai-harness/tests/bound-output.test.ts b/packages/ai-harness/tests/bound-output.test.ts new file mode 100644 index 0000000000..b4c257e053 --- /dev/null +++ b/packages/ai-harness/tests/bound-output.test.ts @@ -0,0 +1,314 @@ +import { mkdtemp, readFile, readdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it } from 'vitest' +import { defineAgent, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { boundText, boundToolOutput } from '../src/first-party/bound-output' +import { mockAdapter, text, toolCall } from './helpers' +import type { AnyTool, ModelMessage } from '@tanstack/ai' +import type { AnyAgent, AnyHarness } from '../src' + +const DAY = 24 * 60 * 60 * 1000 +/** Five lines, 34 bytes. */ +const FIVE_LINES = 'line 1\nline 2\nline 3\nline 4\nline 5' +/** FIVE_LINES cut to `maxLines: 3`. */ +const FIRST_THREE = + 'line 1\nline 2\nline 3\n\n[Output cut. Showing the first 3 of 5 lines, 20 of 34 bytes.]' + +let dir: string + +beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-bound-output-')) +}) + +afterEach(async () => { + await rm(dir, { recursive: true, force: true }) +}) + +/** A tool named `lookup` that returns `output`. */ +function lookup(output: unknown) { + return toolDefinition({ + name: 'lookup', + description: 'Look something up', + }).server(async () => output) +} + +/** A model that calls `tool` once, then answers. */ +function callsOnce(tool: string) { + return mockAdapter([() => toolCall(tool, {}), () => text('done')]) +} + +/** The tool message in the messages of one model call. */ +function toolMessageIn(call: { messages: Array }) { + return call.messages.find((message) => message.role === 'tool') +} + +/** Open a session on `harness`, send one prompt, and close it. */ +async function promptOnce(harness: AnyHarness) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open(harness, { threadId: 't' }) + await session.prompt('go') + await host.close() +} + +/** The lead model calls `tool`. Gives back the tool message the model saw. */ +async function leadToolMessage( + tool: AnyTool, + options: Parameters[0] = { maxLines: 3 }, +) { + const model = callsOnce(tool.name) + await promptOnce( + defineHarness({ + name: 'test/bound-output', + adapter: model.adapter, + tools: [tool], + plugins: () => [boundToolOutput(options)], + }), + ) + return toolMessageIn(model.calls[1]) +} + +/** The lead model calls `lookup`, which returns `output`. Gives back the result the model saw. */ +async function leadSees( + output: unknown, + options?: Parameters[0], +) { + const message = await leadToolMessage(lookup(output), options) + return message?.content +} + +/** The lead model calls `agent` as a subagent. Gives back the result the model saw, parsed. */ +async function leadSeesAgent(agent: AnyAgent) { + const lead = callsOnce(agent.name) + await promptOnce( + defineHarness({ + name: 'test/bound-output', + adapter: lead.adapter, + subagents: { agents: [agent] }, + plugins: () => [boundToolOutput({ maxLines: 3 })], + }), + ) + return JSON.parse(String(toolMessageIn(lead.calls[1])?.content)) +} + +/** The paths of the outputs saved in `dir`. */ +async function savedOutputs() { + const names = await readdir(dir) + return names + .filter((name) => name.startsWith('tool-output-')) + .map((name) => join(dir, name)) +} + +describe('boundText', () => { + it.each([ + [ + 'too many lines', + 'a\nb\nc\nd', + { maxLines: 2, maxBytes: 100 }, + 'a\nb\n\n[Output cut. Showing the first 2 of 4 lines, 3 of 7 bytes.]', + ], + [ + 'too many bytes', + 'abcdef', + { maxLines: 10, maxBytes: 4 }, + 'abcd\n\n[Output cut. Showing the first 1 of 1 lines, 4 of 6 bytes.]', + ], + [ + 'a two-byte character over the byte limit', + 'aéé', + { maxLines: 10, maxBytes: 4 }, + 'aé\n\n[Output cut. Showing the first 1 of 1 lines, 3 of 5 bytes.]', + ], + [ + 'a surrogate pair over the byte limit', + 'x\u{1F600}', + { maxLines: 10, maxBytes: 4 }, + 'x\n\n[Output cut. Showing the first 1 of 1 lines, 1 of 5 bytes.]', + ], + ])('keeps the head of text with %s', (_case, input, limits, expected) => { + expect(boundText(input, limits)).toEqual({ + text: expected, + truncated: true, + }) + }) + + it.each([ + [ + 'too many lines', + 'a\nb\nc\nd', + { maxLines: 2, maxBytes: 100 }, + '[Output cut. Showing the last 2 of 4 lines, 3 of 7 bytes.]\n\nc\nd', + ], + [ + 'too many lines and a final newline', + 'a\nb\nc\n', + { maxLines: 2, maxBytes: 100 }, + '[Output cut. Showing the last 2 of 3 lines, 4 of 6 bytes.]\n\nb\nc\n', + ], + [ + 'too many bytes', + 'abcdef', + { maxLines: 10, maxBytes: 4 }, + '[Output cut. Showing the last 1 of 1 lines, 4 of 6 bytes.]\n\ncdef', + ], + [ + 'a two-byte character over the byte limit', + 'ééa', + { maxLines: 10, maxBytes: 4 }, + '[Output cut. Showing the last 1 of 1 lines, 3 of 5 bytes.]\n\néa', + ], + [ + 'a surrogate pair over the byte limit', + '\u{1F600}x', + { maxLines: 10, maxBytes: 4 }, + '[Output cut. Showing the last 1 of 1 lines, 1 of 5 bytes.]\n\nx', + ], + ])('keeps the tail of text with %s', (_case, input, limits, expected) => { + expect(boundText(input, { ...limits, keep: 'tail' })).toEqual({ + text: expected, + truncated: true, + }) + }) + + it('leaves text inside both limits, and a final newline is not an extra line', () => { + expect(boundText('a\nb\n', { maxLines: 2, maxBytes: 4 })).toEqual({ + text: 'a\nb\n', + truncated: false, + }) + }) +}) + +describe('boundToolOutput', () => { + it('cuts a long string result to its first lines, with a note', async () => { + expect(await leadSees(FIVE_LINES)).toBe(FIRST_THREE) + }) + + it('cuts any other value as JSON text', async () => { + const seen = await leadSees({ name: 'x'.repeat(200) }, { maxBytes: 20 }) + + expect(seen).toBe( + `{"name":"${'x'.repeat(11)}\n\n[Output cut. Showing the first 1 of 1 lines, 20 of 211 bytes.]`, + ) + }) + + it('leaves a short result alone and saves no file', async () => { + const seen = await leadSees('line 1\nline 2', { maxLines: 3, dir }) + + expect(seen).toBe('line 1\nline 2') + expect(await savedOutputs()).toEqual([]) + }) + + it('leaves content parts alone', async () => { + const image = [ + { + type: 'image', + source: { type: 'data', value: 'A'.repeat(500), mimeType: 'image/png' }, + }, + ] + + expect(await leadSees(image, { maxBytes: 100 })).toEqual(image) + }) + + it('cuts the error text of a failed call, and it stays an error', async () => { + const failing = toolDefinition({ + name: 'lookup', + description: 'Look something up', + }).server(async () => { + throw new Error(FIVE_LINES) + }) + + const message = await leadToolMessage(failing) + + expect(message).toMatchObject({ + content: JSON.stringify({ error: FIRST_THREE }), + error: FIRST_THREE, + }) + }) + + it('saves the full output in dir and gives its path in the note', async () => { + const seen = await leadSees(FIVE_LINES, { maxLines: 3, dir }) + + const [saved = ''] = await savedOutputs() + expect(saved).toMatch(/tool-output-\d+-call-1\.txt$/) + expect(seen).toBe(`${FIRST_THREE}\n[Full output saved to ${saved}.]`) + expect(await readFile(saved, 'utf8')).toBe(FIVE_LINES) + }) + + it('removes saved outputs older than retentionDays and keeps other files', async () => { + const old = `tool-output-${Date.now() - 8 * DAY}-old.txt` + const recent = `tool-output-${Date.now() - DAY}-recent.txt` + const names = [old, recent, 'notes.txt'] + await Promise.all(names.map((name) => writeFile(join(dir, name), 'x'))) + + await leadSees(FIVE_LINES, { maxLines: 3, dir, retentionDays: 7 }) + + const left = await readdir(dir) + expect(left).not.toContain(old) + expect(left).toEqual(expect.arrayContaining([recent, 'notes.txt'])) + expect(left).toHaveLength(3) + }) + + it('still cuts the result when the full output cannot be saved', async () => { + const file = join(dir, 'a-file') + await writeFile(file, 'x') + + const seen = await leadSees(FIVE_LINES, { + maxLines: 3, + dir: join(file, 'out'), + }) + + expect(seen).toEqual( + expect.stringContaining( + `${FIRST_THREE}\n[The full output was not saved: Error:`, + ), + ) + }) + + it('cuts a tool result in a subagent run too', async () => { + const researcherModel = callsOnce('lookup') + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks things up', + run: (ctx) => + ctx.chat({ + adapter: researcherModel.adapter, + tools: [lookup(FIVE_LINES)], + }), + }) + + await leadSeesAgent(researcher) + + expect(toolMessageIn(researcherModel.calls[1])?.content).toBe(FIRST_THREE) + }) + + it('cuts a subagent answer and keeps its subagentRunId', async () => { + const answers = mockAdapter([() => text(FIVE_LINES)]) + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks things up', + run: (ctx) => ctx.chat({ adapter: answers.adapter }), + }) + + expect(await leadSeesAgent(researcher)).toEqual({ + subagentRunId: expect.stringMatching(/.+/), + result: FIRST_THREE, + }) + }) + + it('leaves a failed subagent call alone, so it keeps its subagentRunId', async () => { + const researcher = defineAgent({ + name: 'researcher', + description: 'Looks things up', + run: async () => { + throw new Error(FIVE_LINES) + }, + }) + + expect(await leadSeesAgent(researcher)).toEqual({ + subagentRunId: expect.stringMatching(/.+/), + error: FIVE_LINES, + }) + }) +}) diff --git a/packages/ai-harness/tests/discover-tools.test.ts b/packages/ai-harness/tests/discover-tools.test.ts new file mode 100644 index 0000000000..b5e6b15ec9 --- /dev/null +++ b/packages/ai-harness/tests/discover-tools.test.ts @@ -0,0 +1,109 @@ +import { describe, expect, it } from 'vitest' +import { z } from 'zod' +import { EventType, toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { mockAdapter, text, toolCall } from './helpers' +import type { SessionEvent } from '../src' + +const echoTool = (name: string, reply: string) => + toolDefinition({ + name, + description: `The ${name} tool`, + inputSchema: z.object({}), + }).server(async () => reply) + +describe('discoverTools', () => { + it('adds tools a plugin finds at run time, skips taken names, and warns on failure', async () => { + let online = false + const connector = definePlugin({ + name: 'test/connector', + setup: () => ({ + discoverTools: () => + online + ? [ + echoTool('remote_search', 'remote result'), + echoTool('lookup', 'shadowed'), + ] + : [], + }), + }) + const broken = definePlugin({ + name: 'test/broken', + setup: () => ({ + discoverTools: async () => { + throw new Error('server is down') + }, + }), + }) + const { adapter, calls } = mockAdapter([ + () => text('nothing yet'), + () => toolCall('remote_search', {}), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/discover', + adapter, + tools: [echoTool('lookup', 'local result')], + plugins: () => [connector, broken], + }), + { threadId: 't' }, + ) + const names = (call: number) => + (calls[call].tools ?? []).map((tool: { name: string }) => tool.name) + + const seen: Array = [] + const controller = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: controller.signal })) + seen.push(entry) + })() + + await session.prompt('before sign-in') + expect(names(0)).toEqual(['lookup']) + + online = true + const turn = await session.prompt('after sign-in') + expect(turn.text).toBe('done') + // The static `lookup` wins over the discovered one with the same name. + expect(names(1)).toEqual(['lookup', 'remote_search']) + expect(JSON.stringify(calls[2].messages)).toContain('remote result') + + controller.abort() + await reading + const warnings = seen.filter( + (entry) => + entry.event.type === EventType.CUSTOM && + entry.event.name === 'harness.plugin.warning', + ) + expect(warnings[0]?.event).toMatchObject({ + value: { plugin: 'test/broken', message: 'server is down' }, + }) + await host.close() + }) +}) + +describe('the default agent loop', () => { + it('keeps calling the model past five steps', async () => { + const steps = Array.from( + { length: 7 }, + (_, index) => () => toolCall('lookup', {}, `call-${index}`), + ) + const { adapter, calls } = mockAdapter([...steps, () => text('finally')]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/long-turn', + adapter, + tools: [echoTool('lookup', 'fact')], + }), + { threadId: 't' }, + ) + const turn = await session.prompt('dig deep') + expect(calls).toHaveLength(8) + expect(turn.text).toBe('finally') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/edge-safety.test.ts b/packages/ai-harness/tests/edge-safety.test.ts index e6f443f1e0..f585e5fd8a 100644 --- a/packages/ai-harness/tests/edge-safety.test.ts +++ b/packages/ai-harness/tests/edge-safety.test.ts @@ -17,9 +17,9 @@ const srcDir = resolve(here, '..', 'src') /** * Node-only entry points, relative to `src`. A path that ends with `/` allows - * the whole folder. The core has none. + * the whole folder. */ -const ALLOWED: Array = [] +const ALLOWED = ['first-party/files.ts', 'first-party/bound-output.ts'] const isAllowed = (file: string) => ALLOWED.some((entry) => diff --git a/packages/ai-harness/tests/first-party.test.ts b/packages/ai-harness/tests/first-party.test.ts new file mode 100644 index 0000000000..1b6fc092a1 --- /dev/null +++ b/packages/ai-harness/tests/first-party.test.ts @@ -0,0 +1,178 @@ +import { mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { + compact, + decidePermission, + fileCommands, + globToRegExp, + modelPicker, + projectInstructions, + todos, + usage, +} from '../src/first-party' +import { messageTexts, mockAdapter, text, toolCall } from './helpers' +import type { HarnessPlugin } from '../src' +import type { AnyTextAdapter } from '@tanstack/ai' + +let root = '' +beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), 'harness-')) +}) +afterEach(async () => { + await rm(root, { recursive: true, force: true }) +}) + +async function open(adapter: AnyTextAdapter, plugins: Array) { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/first-party', + adapter, + plugins: () => plugins, + }), + { + threadId: 't', + }, + ) + return { host, session, persistence } +} + +describe('decidePermission', () => { + const rules = [ + { tool: 'read_file', decision: 'allow' as const, kind: 'read' as const }, + { tool: 'write_file', decision: 'ask' as const, kind: 'edit' as const }, + { tool: 'bash', decision: 'ask' as const, kind: 'execute' as const }, + ] + it('applies modes on top of the rules', () => { + expect(decidePermission(rules, 'write_file', 'default')).toBe('ask') + expect(decidePermission(rules, 'write_file', 'acceptEdits')).toBe('allow') + expect(decidePermission(rules, 'bash', 'acceptEdits')).toBe('ask') + expect(decidePermission(rules, 'write_file', 'plan')).toBe('deny') + expect(decidePermission(rules, 'read_file', 'plan')).toBe('allow') + expect(decidePermission(rules, 'bash', 'bypass')).toBe('allow') + }) +}) + +describe('globToRegExp', () => { + it('matches globs like the list tool', () => { + expect(globToRegExp('src/**/*.ts').test('src/a/b.ts')).toBe(true) + expect(globToRegExp('src/**/*.ts').test('src/b.ts')).toBe(true) + expect(globToRegExp('*.md').test('docs/a.md')).toBe(false) + }) +}) + +describe('modelPicker', () => { + it('switches the main model with /model', async () => { + const fast = mockAdapter([() => text('fast answer')]) + const smart = mockAdapter([() => text('smart answer')]) + const { host, session } = await open(smart.adapter, [ + modelPicker({ + choices: { smart: smart.adapter, fast: fast.adapter }, + default: 'smart', + }), + ]) + expect(await session.command('model', 'fast')).toContain('Model: fast') + expect(await session.prompt('hi')).toEqual({ text: 'fast answer' }) + expect(await session.command('model', 'nope')).toContain('Unknown model') + await host.close() + }) +}) + +describe('todos', () => { + it('keeps the list in state and shows it in the next prompt', async () => { + const { adapter, calls } = mockAdapter([ + () => + toolCall('todo_write', { + todos: [{ text: 'write tests', status: 'in_progress' }], + }), + () => text('planned'), + () => text('next'), + ]) + const { host, session, persistence } = await open(adapter, [todos()]) + await session.prompt('plan it') + expect(await session.command('todos')).toBe('[~] write tests') + await session.prompt('continue') + expect(JSON.stringify(calls[2].systemPrompts)).toContain('[~] write tests') + expect( + await persistence.stores.metadata.get('plugin:tanstack/todos', 't'), + ).toEqual({ + items: [{ text: 'write tests', status: 'in_progress' }], + }) + await host.close() + }) +}) + +describe('project files', () => { + it('adds AGENTS.md to the prompt and turns command files into commands', async () => { + await writeFile(join(root, 'AGENTS.md'), 'Use tabs.') + await mkdir(join(root, 'commands')) + await writeFile( + join(root, 'commands', 'review.md'), + '---\ndescription: Review a file\n---\nReview $ARGUMENTS carefully.', + ) + const { adapter, calls } = mockAdapter([() => text('reviewed')]) + const { host, session } = await open(adapter, [ + projectInstructions({ root }), + fileCommands({ dir: join(root, 'commands') }), + ]) + expect(session.commands()).toEqual([ + { + name: 'review', + description: 'Review a file', + owner: 'tanstack/file-commands', + }, + ]) + await session.command('review', 'src/app.ts') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + expect(messageTexts(calls[0])).toEqual(['Review src/app.ts carefully.']) + expect(JSON.stringify(calls[0].systemPrompts)).toContain('Use tabs.') + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await host.close() + }) +}) + +describe('compact and usage', () => { + it('replaces the transcript with a summary and counts tokens', async () => { + const summarizer = mockAdapter([() => text('We planned a trip.')]) + // Real adapters report usage on RUN_FINISHED. + const withUsage = (answer: string) => () => + text(answer).map((chunk) => + chunk.type === 'RUN_FINISHED' + ? { + ...chunk, + usage: [ + { + inputTokens: 10, + outputTokens: 5, + totalTokens: 15, + cachedInputTokens: 6, + cacheWriteInputTokens: 2, + }, + ], + } + : chunk, + ) + const main = mockAdapter([withUsage('one'), withUsage('two')]) + const { host, session, persistence } = await open(main.adapter, [ + compact({ adapter: summarizer.adapter }), + usage(), + ]) + await session.prompt('first') + await session.prompt('second') + expect(await session.command('compact')).toBe( + 'Compacted 4 messages into a summary.', + ) + const transcript = await persistence.stores.messages.loadThread('t') + expect(transcript).toHaveLength(2) + expect(transcript[0]?.content).toContain('We planned a trip.') + expect(await session.command('usage')).toBe( + '2 model calls, 20 input tokens (12 cache read, 4 cache write), 10 output tokens, 30 total.', + ) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/instructions.test.ts b/packages/ai-harness/tests/instructions.test.ts new file mode 100644 index 0000000000..8f8ef03d1b --- /dev/null +++ b/packages/ai-harness/tests/instructions.test.ts @@ -0,0 +1,234 @@ +import { mkdir, mkdtemp, rm, utimes, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { dirname, join } from 'node:path' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness } from '../src' +import { projectInstructions } from '../src/first-party' +import { mockAdapter, text } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessPlugin } from '../src' + +// `~` in a global path is the home folder. Each test points it into its own +// folder. +const home = vi.hoisted(() => ({ dir: '' })) +vi.mock('node:os', async (importOriginal) => ({ + ...(await importOriginal()), + homedir: () => home.dir, +})) + +// A fake git: real git can take seconds to start on a busy machine, more +// than the plugin waits. The plugin runs git with `execFile`. The fake names +// `git.branch` for `git rev-parse` in `git.cwd`, and fails for anything else. +const git = vi.hoisted(() => { + const state: { cwd: string; branch: string | undefined } = { + cwd: '', + branch: undefined, + } + return state +}) +vi.mock('node:child_process', async (importOriginal) => ({ + ...(await importOriginal()), + execFile: ( + file: string, + args: Array, + options: { cwd?: string }, + callback: (error: Error | null, stdout: string) => void, + ) => { + const isBranchQuery = + file === 'git' && + args.join(' ') === 'rev-parse --abbrev-ref HEAD' && + options.cwd === git.cwd && + git.branch !== undefined + if (isBranchQuery) callback(null, `${git.branch}\n`) + else callback(new Error('git failed'), '') + }, +})) + +let dir = '' +/** The `root` of the plugin, in the test folder. */ +let root = '' +beforeEach(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-instructions-')) + home.dir = join(dir, 'home') + root = join(dir, 'project') + await mkdir(root) +}) +afterEach(async () => { + vi.useRealTimers() + await rm(dir, { recursive: true, force: true }) +}) + +/** Write `content` to `path` in the test folder, and make its folders. */ +async function put(path: string, content: string) { + const full = join(dir, path) + await mkdir(dirname(full), { recursive: true }) + await writeFile(full, content) +} + +/** Write `content` to `path`, with an mtime `minutes` in the future. */ +async function change(path: string, content: string, minutes: number) { + await put(path, content) + const time = Date.now() / 1000 + minutes * 60 + await utimes(join(dir, path), time, time) +} + +/** A model that answers 'ok', on a model with mid-conversation changes. */ +function okModel() { + const { adapter, calls } = mockAdapter(() => text('ok')) + const withChannels: AnyTextAdapter = { + ...adapter, + midConversationChannels: { tools: true, systemPrompts: true }, + } + return { adapter: withChannels, calls } +} + +/** Open a session on `adapter` with `plugins`. */ +async function open(adapter: AnyTextAdapter, plugins: Array) { + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/instructions', + adapter, + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +/** The system prompts of one turn with `projectInstructions(options)`. */ +async function promptsOfOneTurn( + options: Parameters[0], +) { + const { adapter, calls } = okModel() + const { host, session } = await open(adapter, [projectInstructions(options)]) + await session.prompt('hi') + await host.close() + return calls[0].systemPrompts +} + +describe('projectInstructions', () => { + it('puts global files first, then the files from the repo root down to root', async () => { + await put('home/.config/AGENTS.md', 'Global rule.') + await mkdir(join(dir, 'repo', '.git'), { recursive: true }) + await put('repo/AGENTS.md', 'Repo rule.') + await put('repo/pkg/AGENTS.md', 'Package rule.') + await put('repo/pkg/CLAUDE.md', 'Claude rule.') + + expect( + await promptsOfOneTurn({ + root: join(dir, 'repo', 'pkg'), + global: ['~/.config/AGENTS.md'], + env: false, + }), + ).toEqual([ + 'Global instructions from ~/.config/AGENTS.md:\nGlobal rule.', + 'Project instructions from ../AGENTS.md:\nRepo rule.', + 'Project instructions from AGENTS.md:\nPackage rule.', + 'Project instructions from CLAUDE.md:\nClaude rule.', + ]) + }) + + it('reads only root when no folder has .git', async () => { + await put('AGENTS.md', 'Outer rule.') + await put('project/AGENTS.md', 'Inner rule.') + + expect(await promptsOfOneTurn({ root, env: false })).toEqual([ + 'Project instructions from AGENTS.md:\nInner rule.', + ]) + }) + + it('adds a changed file at the end, as a mid-conversation change', async () => { + await put('project/AGENTS.md', 'Use tabs.') + const { adapter, calls } = okModel() + const { host, session } = await open(adapter, [ + projectInstructions({ root, env: false }), + ]) + + await session.prompt('one') + await change('project/AGENTS.md', 'Use spaces.', 1) + await session.prompt('two') + await host.close() + + expect(calls[1].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nUse tabs.', + 'Project instructions from AGENTS.md changed. Follow this text instead:\nUse spaces.', + ]) + expect(calls[1].midConversationChanges).toEqual({ + start: { tools: [], systemPrompts: 1 }, + changes: [{ before: calls[1].messages.length, systemPrompts: 1 }], + }) + }) + + it('notes a removed file and a new file', async () => { + await put('project/AGENTS.md', 'Use tabs.') + const { adapter, calls } = okModel() + const { host, session } = await open(adapter, [ + projectInstructions({ root, env: false }), + ]) + + await session.prompt('one') + await rm(join(root, 'AGENTS.md')) + await put('project/CLAUDE.md', 'Be brief.') + await session.prompt('two') + await host.close() + + expect(calls[1].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nUse tabs.', + 'Project instructions from AGENTS.md were removed.', + 'Project instructions from CLAUDE.md:\nBe brief.', + ]) + }) + + it('builds the prompt again when no spare slot is left', async () => { + await put('project/AGENTS.md', 'Version 0.') + const { adapter, calls } = okModel() + // One file and 8 spare slots: 9 slots in all. + const { host, session } = await open(adapter, [ + projectInstructions({ root, files: ['AGENTS.md'], env: false }), + ]) + + await session.prompt('turn 0') + for (let version = 1; version <= 9; version++) { + await change('project/AGENTS.md', `Version ${version}.`, version) + await session.prompt(`turn ${version}`) + } + await host.close() + + expect(calls[8].systemPrompts.at(-1)).toBe( + 'Project instructions from AGENTS.md changed. Follow this text instead:\nVersion 8.', + ) + expect(calls[9].systemPrompts).toEqual([ + 'Project instructions from AGENTS.md:\nVersion 9.', + ]) + }) + + it('adds an environment block with the date, platform, and folder', async () => { + vi.setSystemTime(new Date('2026-03-04T12:00:00Z')) + + expect(await promptsOfOneTurn({ root })).toEqual([ + `Environment:\n- Date: 2026-03-04\n- Platform: ${process.platform}\n- Working folder: ${root}\n- Git repository: no`, + ]) + }) + + it('names the git branch in a git repository', async () => { + await mkdir(join(root, '.git')) + git.cwd = root + git.branch = 'instructions-test' + + const [environment] = await promptsOfOneTurn({ root }) + expect(environment).toMatch( + /- Git repository: yes\n- Git branch: instructions-test$/, + ) + }) + + it('leaves out the branch when git cannot name it', async () => { + await mkdir(join(root, '.git')) + git.cwd = root + git.branch = undefined + + const [environment] = await promptsOfOneTurn({ root }) + expect(environment).toMatch(/- Git repository: yes$/) + }) +}) diff --git a/packages/ai-harness/tests/permissions.test.ts b/packages/ai-harness/tests/permissions.test.ts new file mode 100644 index 0000000000..d6a64ed6c3 --- /dev/null +++ b/packages/ai-harness/tests/permissions.test.ts @@ -0,0 +1,780 @@ +import { describe, expect, it, vi } from 'vitest' +import { toolDefinition } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { + PermissionDecisionCapability, + PermissionPrompt, + PermissionResources, + PermissionRules, + decidePermission, + deleteSavedPermission, + listSavedPermissions, + permissions, +} from '../src/first-party/permissions' +import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessPersistence, HarnessPlugin, HarnessSession } from '../src' +import type { + CallResources, + PermissionDecision, + PermissionMode, + PermissionRule, +} from '../src/first-party/permissions' +import type { Reply } from './helpers' + +const ROOT = '/work/repo' +const WINDOWS_ROOT = 'C:\\work\\repo' +const readAllowed: Array = [ + { tool: 'read_file', decision: 'allow', kind: 'read' }, +] + +function read( + path: string, + options: { root?: string; rules?: Array } = {}, +) { + return decidePermission( + options.rules ?? readAllowed, + 'read_file', + 'default', + { + resources: { paths: [path] }, + root: options.root ?? ROOT, + }, + ) +} + +/** Run `fn` as if the host were `platform`. Letter case rules depend on it. */ +function onPlatform(platform: NodeJS.Platform, fn: () => T) { + const real = process.platform + Object.defineProperty(process, 'platform', { value: platform }) + try { + return fn() + } finally { + Object.defineProperty(process, 'platform', { value: real }) + } +} + +function run(commands: Array, rules: Array) { + return decidePermission(rules, 'bash', 'default', { + resources: { commands }, + }) +} + +describe('decidePermission with resources', () => { + it('lets the last matching rule win, and needs every resource allowed', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + { tool: 'read_file', resource: 'secrets/public/**', decision: 'allow' }, + ] + const decide = (paths: Array) => + decidePermission(rules, 'read_file', 'default', { + resources: { paths }, + root: ROOT, + }) + + expect(decide(['src/a.ts'])).toBe('allow') + expect(decide(['secrets/key'])).toBe('deny') + expect(decide(['secrets/public/logo.png'])).toBe('allow') + expect(decide(['src/a.ts', 'secrets/key'])).toBe('deny') + }) + + it.each([ + ['darwin', 'deny'], + ['win32', 'deny'], + ['linux', 'allow'], + ] as const)( + 'on %s, a deny rule for a posix path matches another letter case: %s', + (platform, expected) => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ] + expect(onPlatform(platform, () => read('Secrets/key', { rules }))).toBe( + expected, + ) + }, + ) + + it('fails closed when the resources of a call are not known', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'ask' }, + ] + + expect(decidePermission(rules, 'read_file', 'default')).toBe('ask') + }) + + it('matches command rules against each part, and * spans slashes', () => { + const rules: Array = [ + { tool: 'bash', decision: 'ask', kind: 'execute' }, + { tool: 'bash', resource: 'ls*', decision: 'allow' }, + { tool: 'bash', resource: 'git *', decision: 'allow' }, + ] + + expect(run(['ls -la'], rules)).toBe('allow') + expect(run(['git add src/a.ts'], rules)).toBe('allow') + expect(run(['ls', 'rm x'], rules)).toBe('ask') + }) + + it.each([ + ['a command that was not split', 'ls && rm x'], + ['command substitution', 'echo $(rm x)'], + ['a backtick', 'echo `rm x`'], + ['process substitution', 'diff <(rm x) y'], + ['a heredoc', 'cat < ~/.bashrc'], + ])('asks for %s even when * is allowed', (_name, part) => { + expect( + run([part], [{ tool: 'bash', resource: '*', decision: 'allow' }]), + ).toBe('ask') + }) + + it('still denies a part with shell syntax when a deny matches', () => { + const rules: Array = [ + { tool: 'bash', resource: '*', decision: 'allow' }, + { tool: 'bash', resource: 'echo *', decision: 'deny' }, + ] + + expect(run(['echo $(rm x)'], rules)).toBe('deny') + }) + + it.each([ + ['..', '../secret.txt'], + ['.. in the middle', 'src/../../secret.txt'], + ['an absolute path', '/etc/passwd'], + ['a folder that only starts like the root', '/work/repo-evil/a.ts'], + ])('asks for a path that leaves the workspace: %s', (_name, path) => { + expect(read(path)).toBe('ask') + }) + + it('asks for another letter case of the root on a case-sensitive system', () => { + expect(onPlatform('linux', () => read('/work/Repo/a.ts'))).toBe('ask') + }) + + it.each([ + ['another drive', 'D:\\other\\x.txt'], + ['a sibling folder in another letter case', 'C:/Work/Other/x.txt'], + ['a folder that only starts like the root', 'C:\\WORK\\REPO-evil\\x.txt'], + ['a path relative to the current folder of a drive', 'C:secret.txt'], + ['a stream of a file', 'notes.txt:hidden'], + ['a name that ends in a dot', 'notes.txt.'], + ])('asks for a Windows path that leaves the workspace: %s', (_name, path) => { + expect(read(path, { root: WINDOWS_ROOT })).toBe('ask') + }) + + it('compares Windows paths without letter case, for allows and denies', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ] + + expect(read('c:\\WORK\\Repo\\src\\a.ts', { root: WINDOWS_ROOT })).toBe( + 'allow', + ) + expect(read('SECRETS\\key', { root: WINDOWS_ROOT, rules })).toBe('deny') + }) + + it('lets a rule with a resource allow a path outside the workspace', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: '/shared/**', decision: 'allow' }, + ] + + expect(read('/shared/notes.md', { rules })).toBe('allow') + }) + + it.each(['.env', 'config/prod.env', '.env.local', 'Prod.ENV'])( + 'asks before reading %s', + (path) => { + expect(read(path)).toBe('ask') + }, + ) + + it('reads other files, and a .env file a rule names', () => { + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: '.env.example', decision: 'allow' }, + ] + + expect(read('environment.ts')).toBe('allow') + expect(read('.env.example', { rules })).toBe('allow') + }) + + it('keeps the edit kind in plan mode when a later rule allows the tool', () => { + const rules: Array = [ + { tool: 'write_file', decision: 'ask', kind: 'edit' }, + { tool: 'write_file', decision: 'allow' }, + ] + + expect(decidePermission(rules, 'write_file', 'plan')).toBe('deny') + }) + + it('in acceptEdits mode, still asks for an edit outside the workspace', () => { + const rules: Array = [ + { tool: 'write_file', decision: 'ask', kind: 'edit' }, + ] + const write = (path: string, mode: PermissionMode) => + decidePermission(rules, 'write_file', mode, { + resources: { paths: [path] }, + root: ROOT, + }) + + expect(write('src/a.ts', 'acceptEdits')).toBe('allow') + expect(write('../a.ts', 'acceptEdits')).toBe('ask') + }) +}) + +/** The string `key` of a call input. Throws when it is missing. */ +function field(input: unknown, key: string) { + const value = + typeof input === 'object' && input !== null + ? Reflect.get(input, key) + : undefined + if (typeof value !== 'string') throw new Error(`No ${key}.`) + return value +} + +/** + * `read_file` and `bash` stand-ins that say what each call touches. `ran` + * lists the tools that ran. The `bash` splitter splits on ` && ` only. + */ +function fakeTools() { + const ran: Array = [] + const tool = (name: string) => + toolDefinition({ name, description: name }).server(async () => { + ran.push(name) + return `${name} done` + }) + const plugin = definePlugin({ + name: 'test/tools', + setup: () => ({ + tools: [tool('read_file'), tool('bash'), tool('deploy')], + contribute: [ + PermissionRules.item({ + tool: 'read_file', + decision: 'allow', + kind: 'read', + }), + PermissionRules.item({ + tool: 'bash', + decision: 'ask', + kind: 'execute', + }), + PermissionResources.item({ + read_file: { paths: (input) => [field(input, 'path')] }, + bash: { commands: (input) => field(input, 'command').split(' && ') }, + }), + ], + }), + }) + return { plugin, ran } +} + +async function open( + persistence: HarnessPersistence, + plugins: Array, + replies: Array, + threadId = 't', +) { + const { adapter, calls } = mockAdapter([...replies, () => text('done')]) + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/permissions', + adapter, + plugins: () => plugins, + }), + { threadId }, + ) + return { host, session, calls } +} + +/** Wait for the one open question, then answer it. Returns its message. */ +async function answer(session: HarnessSession, value: unknown) { + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + if (!question) throw new Error('No question.') + await session.answer(question.questionId, value) + return question.message +} + +describe('permissions()', () => { + it('saves an always answer for the project, for the next call and a new session', async () => { + const persistence = memoryPersistence() + const tools = fakeTools() + const plugins = [permissions({ root: ROOT }), tools.plugin] + const first = await open(persistence, plugins, [ + () => toolCall('read_file', { path: '.env' }, 'c1'), + () => toolCall('read_file', { path: '.env' }, 'c2'), + ]) + + const turn = first.session.prompt('read it twice') + expect(await answer(first.session, { answer: 'always' })).toContain('.env') + await turn + + expect(tools.ran).toEqual(['read_file', 'read_file']) + expect( + await persistence.stores.metadata.get( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + ), + ).toEqual([{ tool: 'read_file', resource: '.env', decision: 'allow' }]) + await first.host.close() + + const second = await open( + persistence, + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('read_file', { path: '.env' }, 'c3'), + () => toolCall('read_file', { path: 'other.env' }, 'c4'), + ], + 't2', + ) + const next = second.session.prompt('read again') + // `.env` runs at once. `other.env` has no saved answer, so it asks. + expect(await answer(second.session, { answer: 'once' })).toContain( + 'other.env', + ) + await next + + expect(tools.ran).toHaveLength(4) + await second.host.close() + }) + + it('keeps an always answer for the session when there is no metadata store', async () => { + const { stores } = memoryPersistence() + const tools = fakeTools() + const { host, session } = await open( + { stores: { messages: stores.messages, runs: stores.runs } }, + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('read_file', { path: '../shared.txt' }, 'c1'), + () => toolCall('read_file', { path: '../shared.txt' }, 'c2'), + ], + ) + + const turn = session.prompt('read it twice') + await answer(session, { answer: 'always' }) + await turn + + expect(tools.ran).toEqual(['read_file', 'read_file']) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) + + it('gives the model a tool error with the message of a reject', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [() => toolCall('read_file', { path: '../secret.txt' }, 'c1')], + ) + + const turn = session.prompt('read outside') + await answer(session, { + answer: 'reject', + message: 'Stay inside the repo.', + }) + await turn + + expect(tools.ran).toEqual([]) + expect(JSON.stringify(calls[1].messages)).toContain( + '{\\"error\\":\\"Stay inside the repo.\\"}', + ) + await host.close() + }) + + it('asks for ls && rm x when only ls* is allowed', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [ + permissions({ + rules: [{ tool: 'bash', resource: 'ls*', decision: 'allow' }], + }), + tools.plugin, + ], + [ + () => toolCall('bash', { command: 'ls -la' }, 'c1'), + () => toolCall('bash', { command: 'ls && rm x' }, 'c2'), + ], + ) + + const turn = session.prompt('list, then clean') + expect(await answer(session, { answer: 'reject' })).toContain('rm x') + await turn + + expect(tools.ran).toEqual(['bash']) + expect(JSON.stringify(calls[2].messages)).toContain( + 'The user denied this tool call.', + ) + await host.close() + }) + + it('hides a tool that its last rule denies, and keeps one a later rule allows in part', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [ + permissions({ + rules: [ + { tool: 'deploy', decision: 'deny' }, + { tool: 'bash', decision: 'deny' }, + { tool: 'bash', resource: 'ls*', decision: 'allow' }, + ], + }), + tools.plugin, + ], + [], + ) + + await session.prompt('hi') + + const names = calls[0].tools.map((tool: { name: string }) => tool.name) + expect(names.sort()).toEqual(['bash', 'read_file']) + await host.close() + }) + + it('refuses a call whose resources cannot be read', async () => { + const tools = fakeTools() + const { host, session, calls } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [() => toolCall('read_file', {}, 'c1')], + ) + + await session.prompt('read nothing') + + expect(tools.ran).toEqual([]) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + expect(JSON.stringify(calls[1].messages)).toContain( + 'Cannot check the permissions of this call: No path.', + ) + await host.close() + }) + + it('shares its rules with other plugins, and they still beat the rules of tool plugins', async () => { + const tools = fakeTools() + let seen: Array = [] + // Another plugin that reads the rules, like codeMode() does. + const reader = definePlugin({ + name: 'test/reader', + setup: (ctx) => { + const rules = ctx.collect(PermissionRules) + return { + prepareTools: ({ tools: list }) => { + seen = [...rules] + return list + }, + } + }, + }) + const rule: PermissionRule = { tool: 'read_file', decision: 'ask' } + const { host, session } = await open( + memoryPersistence(), + // The tool plugin comes after permissions(), and it allows read_file. + [permissions({ root: ROOT, rules: [rule] }), tools.plugin, reader], + [() => toolCall('read_file', { path: 'src/a.ts' }, 'c1')], + ) + + const turn = session.prompt('read a file') + expect(await answer(session, { answer: 'reject' })).toContain('src/a.ts') + await turn + + expect(tools.ran).toEqual([]) + expect(seen).toContainEqual(rule) + await host.close() + }) + + it('tells other plugins the decision a call gets, for a rule, the default, and plan mode', async () => { + const ran: Array = [] + const tool = (name: string) => + toolDefinition({ name, description: name }).server(async () => { + ran.push(name) + return `${name} done` + }) + let decide: + | ((tool: string, mode: PermissionMode) => PermissionDecision) + | undefined + // Another plugin that asks permissions(), like codeMode() does. + const reader = definePlugin({ + name: 'test/reader', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => { + decide = ctx.getOptional(PermissionDecisionCapability) + return { tools: [tool('lookup'), tool('deploy')] } + }, + }) + const { host, session, calls } = await open( + memoryPersistence(), + [ + permissions({ + default: 'ask', + rules: [{ tool: 'lookup', decision: 'allow' }], + }), + reader, + ], + [ + () => toolCall('lookup', {}, 'c1'), + () => toolCall('deploy', {}, 'c2'), + () => text('asked'), + () => toolCall('lookup', {}, 'c3'), + () => toolCall('deploy', {}, 'c4'), + ], + ) + + // The rule allows lookup, and deploy has no rule, so it asks. + const turn = session.prompt('look up, then deploy') + expect(await answer(session, { answer: 'reject' })).toContain('deploy') + await turn + expect(ran).toEqual(['lookup']) + expect(decide?.('lookup', 'default')).toBe('allow') + expect(decide?.('deploy', 'default')).toBe('ask') + + // Plan mode denies what would ask. + await session.command('mode', 'plan') + await session.prompt('again, in plan mode') + expect(ran).toEqual(['lookup', 'lookup']) + expect(JSON.stringify(calls.at(-1).messages)).toContain( + 'This tool is not allowed in plan mode.', + ) + expect(decide?.('lookup', 'plan')).toBe('allow') + expect(decide?.('deploy', 'plan')).toBe('deny') + await host.close() + }) + + it('tells other plugins the decision for a call that touches resources, from the root', async () => { + let decide: + | (( + tool: string, + mode: PermissionMode, + resources?: CallResources, + ) => PermissionDecision) + | undefined + const reader = definePlugin({ + name: 'test/reader', + optionalRequires: [PermissionDecisionCapability], + setup: (ctx) => { + decide = ctx.getOptional(PermissionDecisionCapability) + return {} + }, + }) + const rules: Array = [ + ...readAllowed, + { tool: 'read_file', resource: 'secrets/**', decision: 'deny' }, + ] + const { host } = await open( + memoryPersistence(), + [permissions({ root: ROOT, rules }), reader], + [], + ) + const read = (mode: PermissionMode, ...paths: Array) => + decide?.('read_file', mode, { paths }) + + expect(read('default', `${ROOT}/src/a.ts`)).toBe('allow') + expect(read('default', `${ROOT}/.env`)).toBe('ask') + expect(read('plan', `${ROOT}/.env`)).toBe('deny') + expect(read('default', `${ROOT}/src/a.ts`, `${ROOT}/secrets/key`)).toBe( + 'deny', + ) + await host.close() + }) + + it('never lets a saved answer beat a deny', async () => { + const persistence = memoryPersistence() + await persistence.stores.metadata.set( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + [{ tool: 'read_file', resource: 'secret.txt', decision: 'allow' }], + ) + const tools = fakeTools() + const { host, session } = await open( + persistence, + [ + permissions({ + root: ROOT, + rules: [ + { tool: 'read_file', resource: 'secret.txt', decision: 'deny' }, + ], + }), + tools.plugin, + ], + [() => toolCall('read_file', { path: 'secret.txt' }, 'c1')], + ) + + await session.prompt('read the secret') + + expect(tools.ran).toEqual([]) + await host.close() + }) + + it('lists the saved rules of a project, and deletes one by its fields', async () => { + const persistence = memoryPersistence() + const tools = fakeTools() + const { host, session } = await open( + persistence, + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('read_file', { path: '.env' }, 'c1'), + () => toolCall('bash', { command: 'git status' }, 'c2'), + ], + ) + + const turn = session.prompt('read, then run') + await answer(session, { answer: 'always' }) + await answer(session, { answer: 'always' }) + await turn + await host.close() + + expect(await listSavedPermissions(persistence.stores, ROOT)).toEqual([ + { tool: 'read_file', resource: '.env', decision: 'allow' }, + { tool: 'bash', resource: 'git status', decision: 'allow' }, + ]) + expect(await listSavedPermissions(persistence.stores, '/other')).toEqual([]) + + await deleteSavedPermission(persistence.stores, ROOT, { + tool: 'bash', + resource: 'git *', + decision: 'allow', + }) + await deleteSavedPermission(persistence.stores, ROOT, { + tool: 'read_file', + resource: '.env', + decision: 'allow', + }) + expect(await listSavedPermissions(persistence.stores, ROOT)).toEqual([ + { tool: 'bash', resource: 'git status', decision: 'allow' }, + ]) + }) + + it('lists saved rules with /permissions, and forgets one with /permissions forget', async () => { + const persistence = memoryPersistence() + await persistence.stores.metadata.set( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + [ + { tool: 'bash', resource: 'git status', decision: 'allow' }, + { tool: 'deploy', decision: 'allow' }, + ], + ) + const tools = fakeTools() + const { host, session } = await open( + persistence, + [permissions({ root: ROOT }), tools.plugin], + [], + ) + + expect(await session.command('permissions')).toBe( + 'Saved rules:\n1. bash git status\n2. deploy', + ) + expect(await session.command('permissions', 'forget 3')).toBe( + 'No saved rule 3. Run /permissions to see the list.', + ) + expect(await session.command('permissions', 'forget 1')).toBe( + 'Forgot rule 1: bash git status.', + ) + expect(await listSavedPermissions(persistence.stores, ROOT)).toEqual([ + { tool: 'deploy', decision: 'allow' }, + ]) + await session.command('permissions', 'forget 1') + expect(await session.command('permissions')).toBe('No saved rules.') + await host.close() + }) + + it('asks again in the open session after /permissions forget', async () => { + const tools = fakeTools() + const { host, session } = await open( + memoryPersistence(), + [permissions({ root: ROOT }), tools.plugin], + [ + () => toolCall('bash', { command: 'git status' }, 'c1'), + () => text('ran'), + () => toolCall('bash', { command: 'git status' }, 'c2'), + ], + ) + + const first = session.prompt('run it') + await answer(session, { answer: 'always' }) + await first + expect(await session.command('permissions', 'forget 1')).toBe( + 'Forgot rule 1: bash git status.', + ) + + const second = session.prompt('run it again') + expect(await answer(session, { answer: 'once' })).toContain('git status') + await second + expect(tools.ran).toEqual(['bash', 'bash']) + await host.close() + }) + + it('asks about an action that is not a tool call, like a formatter, by the mode', async () => { + const persistence = memoryPersistence() + let prompt: + | ((tool: string, message: string) => Promise) + | undefined + // A plugin like formatter(): it adds an ask rule, and asks before it runs. + const runner = definePlugin({ + name: 'test/runner', + optionalRequires: [PermissionPrompt], + setup: (ctx) => { + prompt = ctx.getOptional(PermissionPrompt) + return { + contribute: [ + PermissionRules.item({ + tool: 'formatter:*', + decision: 'ask', + kind: 'execute', + }), + ], + } + }, + }) + const { host, session } = await open( + persistence, + [ + permissions({ + root: ROOT, + rules: [{ tool: 'formatter:biome', decision: 'allow' }], + }), + runner, + ], + [], + ) + if (!prompt) throw new Error('No PermissionPrompt.') + + // A user rule allows biome. prettier asks, also in acceptEdits mode. + expect(await prompt('formatter:biome', 'Run biome?')).toBe(true) + await session.command('mode', 'acceptEdits') + const first = prompt('formatter:prettier', 'Run prettier?') + expect(await answer(session, { answer: 'reject' })).toBe( + 'Run prettier? Answer once, always, or reject.', + ) + expect(await first).toBe(false) + + // always saves the rule for the project, so it asks no more. + const second = prompt('formatter:prettier', 'Run prettier?') + await answer(session, { answer: 'always' }) + expect(await second).toBe(true) + expect(await prompt('formatter:prettier', 'Run prettier?')).toBe(true) + expect( + await persistence.stores.metadata.get( + 'tanstack/permissions', + 'permissions:saved:/work/repo', + ), + ).toEqual([{ tool: 'formatter:prettier', decision: 'allow' }]) + + // bypass runs everything. plan runs no command. + await session.command('mode', 'bypass') + expect(await prompt('formatter:oxfmt', 'Run oxfmt?')).toBe(true) + await session.command('mode', 'plan') + expect(await prompt('formatter:prettier', 'Run prettier?')).toBe(false) + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/question.test.ts b/packages/ai-harness/tests/question.test.ts new file mode 100644 index 0000000000..dafa3de4b5 --- /dev/null +++ b/packages/ai-harness/tests/question.test.ts @@ -0,0 +1,190 @@ +import { describe, expect, it, vi } from 'vitest' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { EventType } from '@tanstack/ai' +import { + HARNESS_EVENTS, + createHarnessHost, + defineCommand, + defineHarness, + definePlugin, +} from '../src' +import { question } from '../src/first-party/question' +import { mockAdapter, text, toolCall } from './helpers' +import type { HarnessSession } from '../src' + +const database = { + question: 'Which database?', + header: 'Database', + options: [ + { label: 'Postgres', description: 'Relational, the default' }, + { label: 'SQLite' }, + ], +} + +const features = { + question: 'Which features?', + options: [{ label: 'Auth' }, { label: 'Search' }, { label: 'Billing' }], + multiple: true, +} + +/** The model calls `question` with `questions`, then ends the turn. */ +async function open(questions: Array) { + const { adapter, calls } = mockAdapter([ + () => toolCall('question', { questions }, 'c1'), + () => text('done'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/question', + adapter, + plugins: () => [question()], + }), + { threadId: 't' }, + ) + const turn = session.prompt('ask me') + return { host, session, calls, turn } +} + +/** Wait for the one open question, then answer it. Returns the question. */ +async function reply(session: HarnessSession, value: unknown) { + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [pending] = session.snapshot().pendingQuestions + if (!pending) throw new Error('No question.') + await session.answer(pending.questionId, value) + return pending +} + +/** The tool result the model got on its second call. */ +function toolResult( + calls: Array<{ messages: Array<{ role: string; content: unknown }> }>, +) { + return calls[1]?.messages.find((message) => message.role === 'tool')?.content +} + +describe('question()', () => { + it('asks each question in turn and gives the model the picked labels', async () => { + const { host, session, calls, turn } = await open([database, features]) + + const first = await reply(session, '2') + const second = await reply(session, ['billing', 'Auth']) + await turn + + expect(first.message).toBe( + 'Database: Which database?\n1. Postgres - Relational, the default\n2. SQLite\nType a number or a label, or your own answer.', + ) + expect(first.schema).toEqual({ + type: 'string', + enum: ['Postgres', 'SQLite'], + }) + expect(second.message).toContain('3. Billing') + expect(second.schema).toEqual({ + type: 'array', + items: { type: 'string', enum: ['Auth', 'Search', 'Billing'] }, + minItems: 1, + uniqueItems: true, + }) + expect(toolResult(calls)).toBe( + 'Which database?\nAnswer: SQLite\n\nWhich features?\nAnswer: Billing, Auth', + ) + await host.close() + }) + + it.each([ + ['a label in another case', database, ' postgres ', 'Answer: Postgres'], + [ + 'numbers and labels split by commas', + features, + 'auth, 3', + 'Answer: Auth, Billing', + ], + [ + 'text that names no option', + database, + 'We use Mongo', + 'Own answer: We use Mongo', + ], + ])('reads %s', async (_name, asked, value, expected) => { + const { host, session, calls, turn } = await open([asked]) + + await reply(session, value) + await turn + + expect(toolResult(calls)).toBe(`${asked.question}\n${expected}`) + await host.close() + }) + + it('gives the model a tool error when the reply is empty', async () => { + const { host, session, calls, turn } = await open([database, features]) + + await reply(session, ' ') + await turn + + expect(toolResult(calls)).toContain('The user did not answer.') + // The second question is not asked. + expect(session.snapshot().pendingQuestions).toHaveLength(0) + await host.close() + }) +}) + +describe('ask with a url', () => { + it('sends the url in the event and the snapshot', async () => { + const asker = definePlugin({ + name: 'test/url-asker', + setup: (ctx) => ({ + commands: { + verify: defineCommand({ + description: 'Asks the user to open a page', + run: () => + ctx.session.ask({ + message: 'Open the page, then type done.', + url: 'https://example.com/verify', + }), + }), + }, + }), + }) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/url-question', + adapter: mockAdapter([]).adapter, + plugins: () => [asker], + }), + { threadId: 't' }, + ) + const running = session.command('verify') + + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + expect(session.snapshot().pendingQuestions[0]?.url).toBe( + 'https://example.com/verify', + ) + const pending = await reply(session, 'done') + expect(pending.url).toBe('https://example.com/verify') + expect(await running).toBe('done') + + const controller = new AbortController() + let url: unknown + for await (const { event } of session.events({ + from: '0', + signal: controller.signal, + })) { + if ( + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.question && + event.value !== null && + typeof event.value === 'object' && + 'url' in event.value + ) { + url = event.value.url + controller.abort() + } + } + expect(url).toBe('https://example.com/verify') + await host.close() + }) +}) diff --git a/packages/ai-harness/tests/session-plugins.test.ts b/packages/ai-harness/tests/session-plugins.test.ts index 3f4f6f6e26..d0c56f6112 100644 --- a/packages/ai-harness/tests/session-plugins.test.ts +++ b/packages/ai-harness/tests/session-plugins.test.ts @@ -1,4 +1,7 @@ -import { describe, expect, it, vi } from 'vitest' +import { mkdtemp, rm, writeFile, mkdir } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { EventType, toolDefinition } from '@tanstack/ai' import { memoryPersistence } from '@tanstack/ai-persistence' @@ -11,7 +14,15 @@ import { defineHarness, definePlugin, } from '../src' -import { gate, mockAdapter, text } from './helpers' +import { + compact, + fileCommands, + modelPicker, + projectInstructions, + todos, + usage, +} from '../src/first-party' +import { gate, mockAdapter, text, toolCall } from './helpers' import type { SessionEvent } from '../src' const settings = definePlugin({ @@ -374,6 +385,181 @@ describe('the plugin session API and inspect', () => { }) }) +describe('first-party plugin edges', () => { + let dir: string + beforeAll(async () => { + dir = await mkdtemp(join(tmpdir(), 'harness-files-')) + await mkdir(join(dir, 'commands')) + await writeFile( + join(dir, 'commands', 'review.md'), + '---\ndescription: Review a file\n---\nReview $ARGUMENTS carefully.', + ) + await writeFile(join(dir, 'commands', 'plain.md'), 'Say hi to $ARGUMENTS.') + await writeFile(join(dir, 'commands', 'skip.txt'), 'not a command') + await writeFile(join(dir, 'RULES.md'), 'Use tabs.') + await writeFile(join(dir, 'EMPTY.md'), ' ') + }) + afterAll(async () => { + await rm(dir, { recursive: true, force: true }) + }) + + it('turns markdown files into commands and reads chosen instruction files', async () => { + const { adapter, calls } = mockAdapter([ + () => text('a'), + () => text('b'), + () => text('c'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/files', + adapter, + plugins: () => [ + fileCommands({ dir: join(dir, 'commands') }), + projectInstructions({ + root: dir, + files: ['RULES.md', 'EMPTY.md', 'NONE.md'], + }), + ], + }), + { threadId: 't' }, + ) + const described = session.commands() + expect( + described.map((command) => [command.name, command.description]).sort(), + ).toEqual([ + ['plain', 'Run plain.md'], + ['review', 'Review a file'], + ]) + expect(await session.command('review', 'src/a.ts')).toBe('Sent /review.') + await vi.waitFor(() => expect(calls).toHaveLength(1)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await session.command('plain', { who: 'Ada' }) + await vi.waitFor(() => expect(calls).toHaveLength(2)) + await vi.waitFor(() => expect(session.snapshot().status).toBe('idle')) + await session.command('plain') + await vi.waitFor(() => expect(calls).toHaveLength(3)) + + // A missing folder gives no commands. + const empty = await host.open( + defineHarness({ + name: 'test/files-missing', + adapter, + plugins: () => [fileCommands({ dir: join(dir, 'missing') })], + }), + { threadId: 'empty' }, + ) + expect(empty.commands()).toEqual([]) + + const said = calls.map((call) => JSON.stringify(call.messages.at(-1))) + expect(said[0]).toContain('Review src/a.ts carefully.') + expect(said[1]).toContain('Say hi to {\\"who\\":\\"Ada\\"}.') + expect(said[2]).toContain('Say hi to .') + const prompts = JSON.stringify(calls[0].systemPrompts) + expect(prompts).toContain('Project instructions from RULES.md') + expect(prompts).not.toContain('EMPTY.md') + await host.close() + }) + + it('checks todo input and shows an empty list', async () => { + const { adapter, calls } = mockAdapter([ + () => toolCall('todo_write', { todos: [{ text: 'x', status: 'later' }] }), + () => text('bad input'), + () => toolCall('todo_write', { todos: [] }, 'call-2'), + () => text('cleared'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ name: 'test/todos', adapter, plugins: () => [todos()] }), + { threadId: 't' }, + ) + const before = session.snapshot().plugins + await session.prompt('plan') + // The tool has a plain JSON Schema, which chat() does not check, so the + // tool checks its input itself. + expect(JSON.stringify(calls[1].messages)).toContain( + 'Each todo needs text and a status.', + ) + expect(session.snapshot().plugins).toEqual(before) + expect(await session.command('todos')).toBe('The todo list is empty.') + await session.prompt('clear') + expect(JSON.stringify(calls[3].messages)).toContain( + 'The todo list is empty.', + ) + expect(await session.command('todos')).toBe('The todo list is empty.') + await host.close() + }) + + it('needs a model choice and answers every /model form', async () => { + expect(() => modelPicker({ choices: {} })).toThrow('at least one choice') + const fast = mockAdapter([]).adapter + const smart = mockAdapter([]).adapter + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/models', + adapter: smart, + plugins: () => [modelPicker({ choices: { smart, fast } })], + }), + { threadId: 't' }, + ) + expect(await session.command('model')).toBe( + 'Model: smart. Choices: smart, fast.', + ) + expect(await session.command('model', 'slow')).toBe( + 'Unknown model "slow". Choices: smart, fast.', + ) + expect(await session.command('model', { name: 'fast' })).toBe( + 'Model: fast. It applies at the next turn.', + ) + await host.close() + }) + + it('keeps short conversations, keeps the last messages, and counts usage', async () => { + const summarizer = mockAdapter([() => text('the summary')]) + const withUsage = (content: string) => + text(content).map((chunk) => + chunk.type === EventType.RUN_FINISHED + ? { + ...chunk, + usage: { promptTokens: 3, completionTokens: 2, totalTokens: 5 }, + } + : chunk, + ) + const { adapter } = mockAdapter([ + () => withUsage('one'), + () => withUsage('two'), + ]) + const host = createHarnessHost({ persistence: memoryPersistence() }) + const session = await host.open( + defineHarness({ + name: 'test/compact', + adapter, + plugins: () => [ + compact({ adapter: summarizer.adapter, keepLast: 1 }), + usage(), + ], + }), + { threadId: 't' }, + ) + expect(await session.command('compact')).toBe( + 'The conversation is already short.', + ) + await session.prompt([{ type: 'text', content: 'first question' }] as never) + await session.prompt('second question') + expect(await session.command('compact')).toBe( + 'Compacted 3 messages into a summary.', + ) + expect(JSON.stringify(summarizer.calls[0].messages)).toContain( + 'first question', + ) + expect(await session.command('usage')).toBe( + '2 model calls, 6 input tokens, 4 output tokens, 10 total.', + ) + await host.close() + }) +}) + describe('session events', () => { it('publishes config changes', async () => { const host = createHarnessHost({ persistence: memoryPersistence() }) diff --git a/packages/ai-harness/tests/title-usage.test.ts b/packages/ai-harness/tests/title-usage.test.ts new file mode 100644 index 0000000000..9e2f6ca709 --- /dev/null +++ b/packages/ai-harness/tests/title-usage.test.ts @@ -0,0 +1,310 @@ +import { describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { createHarnessHost, defineHarness, definePlugin } from '../src' +import { usage } from '../src/first-party/session-tools' +import { TitleFailed, title } from '../src/first-party/title' +import { messageTexts, mockAdapter, text } from './helpers' +import type { StreamChunk, TokenUsage } from '@tanstack/ai' +import type { AnyAgent, HarnessPersistence, HarnessPlugin } from '../src' +import type { Reply } from './helpers' + +/** + * Open thread `t`. The main model answers `answer` unless `reply` is given. + * `agents` are the background agents of the session. + */ +async function open( + plugins: Array, + { + persistence = memoryPersistence(), + reply = () => text('answer'), + agents = [], + }: { + persistence?: HarnessPersistence + reply?: Reply + agents?: Array + } = {}, +) { + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/title-usage', + adapter: mockAdapter(reply).adapter, + agents, + plugins: () => plugins, + }), + { threadId: 't' }, + ) + return { host, session } +} + +describe('title()', () => { + it('writes a title from the first user message after the first turn', async () => { + const titler = mockAdapter(() => text('"Trip to Rome"\nA second line')) + const { host, session } = await open([title({ adapter: titler.adapter })]) + + const turn = await session.prompt('Plan a trip to Rome for me') + + expect(turn.text).toBe('answer') + await vi.waitFor(async () => + expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome'), + ) + expect(messageTexts(titler.calls[0])[0]).toContain( + 'Plan a trip to Rome for me', + ) + await host.close() + }) + + it('does not ask for a title again once the session has one', async () => { + const titler = mockAdapter([ + () => text('Trip to Rome'), + () => text('Another title'), + ]) + const { host, session } = await open([title({ adapter: titler.adapter })]) + await session.prompt('Plan a trip to Rome for me') + await vi.waitFor(async () => + expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome'), + ) + + await session.prompt('Add a day in Florence') + + expect(titler.calls).toHaveLength(1) + expect((await host.sessions.get('t'))?.title).toBe('Trip to Rome') + await host.close() + }) + + it('keeps the turn when the title call fails', async () => { + const failures: Array<{ message: string }> = [] + const listener = definePlugin({ + name: 'test/listener', + setup: (ctx) => { + ctx.on(TitleFailed, (failure) => failures.push(failure)) + }, + }) + const titler = mockAdapter(() => { + throw new Error('title model down') + }) + const { host, session } = await open([ + listener, + title({ adapter: titler.adapter }), + ]) + + const turn = await session.prompt('Plan a trip to Rome for me') + + expect(turn.text).toBe('answer') + await vi.waitFor(() => + expect(failures).toEqual([{ message: 'title model down' }]), + ) + expect((await host.sessions.get('t'))?.title).toBeUndefined() + await host.close() + }) + + it('does nothing without a sessions store', async () => { + const { messages, runs, metadata } = memoryPersistence().stores + const titler = mockAdapter(() => text('Trip to Rome')) + const { host, session } = await open([title({ adapter: titler.adapter })], { + persistence: { stores: { messages, runs, metadata } }, + }) + + await session.prompt('Plan a trip to Rome for me') + + expect(titler.calls).toHaveLength(0) + await host.close() + }) +}) + +describe('usage()', () => { + /** Prices in USD per 1M tokens. */ + const PRICES = { input: 2, output: 10, cacheRead: 1, cacheWrite: 4 } + /** + * One model call. With PRICES it costs $3.25: 500k uncached input ($1), + * 100k output ($1), 250k cache read ($0.25), 250k cache write ($1). + */ + const CALL: TokenUsage = { + promptTokens: 1_000_000, + completionTokens: 100_000, + totalTokens: 1_100_000, + promptTokensDetails: { cachedTokens: 250_000, cacheWriteTokens: 250_000 }, + } + /** The `/usage` line of one CALL. */ + const CALL_LINE = + '1 model calls, 1000000 input tokens (250000 cache read, 250000 cache write), 100000 output tokens, 1100000 total.' + /** A model call that answers and reports `usage`. */ + const replyWith = (usage: TokenUsage) => () => + text('answer').map( + (chunk): StreamChunk => + chunk.type === EventType.RUN_FINISHED ? { ...chunk, usage } : chunk, + ) + const withUsage = replyWith(CALL) + /** The prices of the mock adapter's model only. */ + const pricesOf = (modelId: string) => + modelId === 'test-model' ? { cost: PRICES } : undefined + /** A background agent. Its model, `worker-model`, reports `usage`. */ + const workerWith = (usage: TokenUsage) => { + const { adapter } = mockAdapter(replyWith(usage)) + return defineAgent({ + name: 'worker', + description: 'Works', + run: (ctx) => + ctx.chat({ adapter: { ...adapter, model: 'worker-model' } }), + }) + } + + it('uses a provider cost as is, and prices only the calls with no provider cost', async () => { + // The lookup knows every model: it would price the worker call at $3.25. + const { host, session } = await open( + [usage({ model: () => ({ cost: PRICES }) })], + { reply: withUsage, agents: [workerWith({ ...CALL, cost: 0.5 })] }, + ) + + await session.prompt('one') + await session.agent('worker')?.start() + + // The lead call costs $3.25 from the lookup, the worker call $0.50 from + // its provider. + expect(await session.command('usage')).toBe( + [ + '2 model calls, 2000000 input tokens (500000 cache read, 500000 cache write), 200000 output tokens, 2200000 total. Cost: $3.7500.', + `mock/test-model: ${CALL_LINE}`, + `mock/worker-model: ${CALL_LINE}`, + ].join('\n'), + ) + // `/usage` shows the totals of the session. + expect(session.usage().total).toEqual({ + calls: 2, + promptTokens: 2_000_000, + completionTokens: 200_000, + totalTokens: 2_200_000, + cachedTokens: 500_000, + cacheWriteTokens: 500_000, + cost: 0.5, + }) + expect((await host.sessions.get('t'))?.usage?.cost).toBeCloseTo(3.75, 10) + await host.close() + }) + + it('shows a provider cost without the model option, and writes it to the index entry', async () => { + const { host, session } = await open([usage()], { + reply: replyWith({ ...CALL, cost: 0.5 }), + }) + + await session.prompt('one') + + expect(await session.command('usage')).toBe(`${CALL_LINE} Cost: $0.5000.`) + expect((await host.sessions.get('t'))?.usage?.cost).toBe(0.5) + await host.close() + }) + + it('counts a provider cost once in the index entry', async () => { + const { host, session } = await open([usage({ model: pricesOf })], { + reply: replyWith({ ...CALL, cost: 0.5 }), + }) + + await session.prompt('one') + await session.prompt('two') + + // Two calls at $0.50. The host writes no cost of its own. + expect((await host.sessions.get('t'))?.usage?.cost).toBeCloseTo(1, 10) + await host.close() + }) + + it('writes the session totals to the index entry, with the calls of a background agent', async () => { + const { host, session } = await open([usage()], { + reply: withUsage, + agents: [workerWith(CALL)], + }) + + await session.prompt('one') + await session.agent('worker')?.start() + await session.prompt('two') + + // Two lead calls and one worker call. + expect((await host.sessions.get('t'))?.usage).toEqual({ + turns: 3, + promptTokens: 3_000_000, + completionTokens: 300_000, + totalTokens: 3_300_000, + cachedTokens: 750_000, + cacheWriteTokens: 750_000, + }) + await host.close() + }) + + it('keeps a copy of the session totals in its state', async () => { + const { host, session } = await open([usage()], { + reply: withUsage, + agents: [workerWith({ ...CALL, promptTokens: 2_000_000 })], + }) + + await session.prompt('one') + await session.agent('worker')?.start() + + // The worker's call is in the copy too. + expect(session.snapshot().plugins['tanstack/usage']).toEqual({ + calls: 2, + promptTokens: 3_000_000, + completionTokens: 200_000, + totalTokens: 2_200_000, + cachedTokens: 500_000, + cacheWriteTokens: 500_000, + // The input of the lead's latest call, not the worker's. + contextTokens: 1_000_000, + }) + await host.close() + }) + + it('prices each model call and shows the cost in /usage', async () => { + const { host, session } = await open([usage({ model: pricesOf })], { + reply: withUsage, + }) + + await session.prompt('one') + await session.prompt('two') + + expect(await session.command('usage')).toBe( + '2 model calls, 2000000 input tokens (500000 cache read, 500000 cache write), 200000 output tokens, 2200000 total. Cost: $6.5000.', + ) + await host.close() + }) + + it('writes the cost to the index entry and keeps the token totals of the host', async () => { + const { host, session } = await open([usage({ model: pricesOf })], { + reply: withUsage, + }) + + await session.prompt('one') + await session.prompt('two') + + expect((await host.sessions.get('t'))?.usage).toEqual({ + turns: 2, + promptTokens: 2_000_000, + completionTokens: 200_000, + totalTokens: 2_200_000, + cachedTokens: 500_000, + cacheWriteTokens: 500_000, + cost: expect.closeTo(6.5, 10), + }) + await host.close() + }) + + it('adds no cost for a model with no prices, and says so in /usage', async () => { + const { host, session } = await open([usage({ model: () => undefined })], { + reply: withUsage, + }) + + await session.prompt('one') + + expect(await session.command('usage')).toBe( + `${CALL_LINE} Cost: $0.0000 (cost unknown for 1 call).`, + ) + expect((await host.sessions.get('t'))?.usage).toEqual({ + turns: 1, + promptTokens: 1_000_000, + completionTokens: 100_000, + totalTokens: 1_100_000, + cachedTokens: 250_000, + cacheWriteTokens: 250_000, + }) + await host.close() + }) +}) diff --git a/packages/ai-harness/vite.config.ts b/packages/ai-harness/vite.config.ts index 1f3542380f..86b2a7a465 100644 --- a/packages/ai-harness/vite.config.ts +++ b/packages/ai-harness/vite.config.ts @@ -28,7 +28,7 @@ const config = defineConfig({ export default mergeConfig( config, tanstackViteConfig({ - entry: ['./src/index.ts'], + entry: ['./src/index.ts', './src/first-party/index.ts'], srcDir: './src', cjs: false, }), diff --git a/testing/e2e/fixtures/harness-first-party/basic.json b/testing/e2e/fixtures/harness-first-party/basic.json new file mode 100644 index 0000000000..77195c081f --- /dev/null +++ b/testing/e2e/fixtures/harness-first-party/basic.json @@ -0,0 +1,107 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[first-party-plan] write x.txt", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "writer", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[first-party-plan] write x.txt", + "sequenceIndex": 1 + }, + "response": { + "content": "The writer could not write in plan mode." + } + }, + { + "match": { + "userMessage": "[first-party-plan-writer] write x.txt", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "write_file", + "arguments": "{\"path\":\"x.txt\"}" + } + ] + } + }, + { + "match": { + "userMessage": "[first-party-plan-writer] write x.txt", + "sequenceIndex": 1 + }, + "response": { + "content": "Plan mode blocked the write." + } + }, + { + "match": { + "userMessage": "[first-party-agents] first" + }, + "response": { + "content": "Built it." + } + }, + { + "match": { + "userMessage": "[first-party-agents] second", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "lookup", + "arguments": "{}" + } + ] + } + }, + { + "match": { + "userMessage": "[first-party-agents] second", + "sequenceIndex": 1 + }, + "response": { + "content": "Stopped at the step limit." + } + }, + { + "match": { + "userMessage": "[first-party-env] where am i" + }, + "response": { + "content": "You are in a new repository." + } + }, + { + "match": { + "userMessage": "[first-party-title] plan a trip to Rome", + "model": "gpt-5.4-nano" + }, + "response": { + "content": "\"Trip to Rome\"" + } + }, + { + "match": { + "userMessage": "[first-party-title] plan a trip to Rome", + "model": "gpt-5.5" + }, + "response": { + "content": "Here is a plan for Rome." + } + } + ] +} diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index d7fe95be71..e604264202 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -86,6 +86,7 @@ import { Route as ApiEmbeddingRouteImport } from './routes/api.embedding' import { Route as ApiFakeTextRouteImport } from './routes/api.fake-text' import { Route as ApiHarnessContinueRouteImport } from './routes/api.harness-continue' import { Route as ApiHarnessRetryAfterRouteImport } from './routes/api.harness-retry-after' +import { Route as ApiHarnessFirstPartyRouteImport } from './routes/api.harness-first-party' import { Route as ApiFileSourceWireRouteImport } from './routes/api.file-source-wire' import { Route as ApiForeignChunkEventsRouteImport } from './routes/api.foreign-chunk-events' import { Route as ApiForeignClientToolRouteImport } from './routes/api.foreign-client-tool' @@ -582,6 +583,11 @@ const ApiHarnessRetryAfterRoute = ApiHarnessRetryAfterRouteImport.update({ path: '/api/harness-retry-after', getParentRoute: () => rootRouteImport, } as any) +const ApiHarnessFirstPartyRoute = ApiHarnessFirstPartyRouteImport.update({ + id: '/api/harness-first-party', + path: '/api/harness-first-party', + getParentRoute: () => rootRouteImport, +} as any) const ApiFileSourceWireRoute = ApiFileSourceWireRouteImport.update({ id: '/api/file-source-wire', path: '/api/file-source-wire', @@ -1149,6 +1155,7 @@ export interface FileRoutesByFullPath { '/api/fake-text': typeof ApiFakeTextRoute '/api/harness-continue': typeof ApiHarnessContinueRoute '/api/harness-retry-after': typeof ApiHarnessRetryAfterRoute + '/api/harness-first-party': typeof ApiHarnessFirstPartyRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute '/api/foreign-client-tool': typeof ApiForeignClientToolRoute @@ -1321,6 +1328,7 @@ export interface FileRoutesByTo { '/api/fake-text': typeof ApiFakeTextRoute '/api/harness-continue': typeof ApiHarnessContinueRoute '/api/harness-retry-after': typeof ApiHarnessRetryAfterRoute + '/api/harness-first-party': typeof ApiHarnessFirstPartyRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute '/api/foreign-client-tool': typeof ApiForeignClientToolRoute @@ -1494,6 +1502,7 @@ export interface FileRoutesById { '/api/fake-text': typeof ApiFakeTextRoute '/api/harness-continue': typeof ApiHarnessContinueRoute '/api/harness-retry-after': typeof ApiHarnessRetryAfterRoute + '/api/harness-first-party': typeof ApiHarnessFirstPartyRoute '/api/file-source-wire': typeof ApiFileSourceWireRoute '/api/foreign-chunk-events': typeof ApiForeignChunkEventsRoute '/api/foreign-client-tool': typeof ApiForeignClientToolRoute @@ -1668,6 +1677,7 @@ export interface FileRouteTypes { | '/api/fake-text' | '/api/harness-continue' | '/api/harness-retry-after' + | '/api/harness-first-party' | '/api/file-source-wire' | '/api/foreign-chunk-events' | '/api/foreign-client-tool' @@ -1840,6 +1850,7 @@ export interface FileRouteTypes { | '/api/fake-text' | '/api/harness-continue' | '/api/harness-retry-after' + | '/api/harness-first-party' | '/api/file-source-wire' | '/api/foreign-chunk-events' | '/api/foreign-client-tool' @@ -2012,6 +2023,7 @@ export interface FileRouteTypes { | '/api/fake-text' | '/api/harness-continue' | '/api/harness-retry-after' + | '/api/harness-first-party' | '/api/file-source-wire' | '/api/foreign-chunk-events' | '/api/foreign-client-tool' @@ -2185,6 +2197,7 @@ export interface RootRouteChildren { ApiFakeTextRoute: typeof ApiFakeTextRoute ApiHarnessContinueRoute: typeof ApiHarnessContinueRoute ApiHarnessRetryAfterRoute: typeof ApiHarnessRetryAfterRoute + ApiHarnessFirstPartyRoute: typeof ApiHarnessFirstPartyRoute ApiFileSourceWireRoute: typeof ApiFileSourceWireRoute ApiForeignChunkEventsRoute: typeof ApiForeignChunkEventsRoute ApiForeignClientToolRoute: typeof ApiForeignClientToolRoute @@ -2817,6 +2830,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiHarnessRetryAfterRouteImport parentRoute: typeof rootRouteImport } + '/api/harness-first-party': { + id: '/api/harness-first-party' + path: '/api/harness-first-party' + fullPath: '/api/harness-first-party' + preLoaderRoute: typeof ApiHarnessFirstPartyRouteImport + parentRoute: typeof rootRouteImport + } '/api/file-source-wire': { id: '/api/file-source-wire' path: '/api/file-source-wire' @@ -3611,6 +3631,7 @@ const rootRouteChildren: RootRouteChildren = { ApiFakeTextRoute: ApiFakeTextRoute, ApiHarnessContinueRoute: ApiHarnessContinueRoute, ApiHarnessRetryAfterRoute: ApiHarnessRetryAfterRoute, + ApiHarnessFirstPartyRoute: ApiHarnessFirstPartyRoute, ApiFileSourceWireRoute: ApiFileSourceWireRoute, ApiForeignChunkEventsRoute: ApiForeignChunkEventsRoute, ApiForeignClientToolRoute: ApiForeignClientToolRoute, diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts new file mode 100644 index 0000000000..8dc25c4ef9 --- /dev/null +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -0,0 +1,245 @@ +import { execFile } from 'node:child_process' +import { mkdtemp, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' +import { promisify } from 'node:util' +import { createFileRoute } from '@tanstack/react-router' +import { defineAgent, toolDefinition } from '@tanstack/ai' +import { + createHarnessHost, + defineHarness, + definePlugin, +} from '@tanstack/ai-harness' +import { + TitleFailed, + agents, + permissions, + projectInstructions, + title, +} from '@tanstack/ai-harness/plugins' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { z } from 'zod' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { HarnessHost } from '@tanstack/ai-harness' +import { createTextAdapter } from '@/lib/providers' + +const run = promisify(execFile) + +/** + * `agents()`: the first turn runs with the default `build` agent. `/agent + * brief` switches to an agent with `steps: 1`. Its turn calls `lookup` once, + * then the step limit makes the last call without tools. The spec reads the + * model calls in aimock's journal. + */ +async function agentProfiles(host: HarnessHost, model: () => AnyTextAdapter) { + const lookup = toolDefinition({ + name: 'lookup', + description: 'Look up a fact', + inputSchema: z.object({}), + }).server(async () => 'A fact.') + const session = await host.open( + defineHarness({ + name: 'e2e/harness-agents', + adapter: model(), + tools: [lookup], + plugins: () => [ + agents({ + adapter: () => model(), + agents: [ + { + name: 'brief', + description: 'Answers after one tool step', + mode: 'primary', + system: 'You are the brief agent.', + steps: 1, + }, + ], + }), + ], + }), + { threadId: 'e2e-agents' }, + ) + const first = (await session.prompt('[first-party-agents] first')).text + const switched = await session.command('agent', 'brief') + const second = (await session.prompt('[first-party-agents] second')).text + return { first, switched, second } +} + +/** + * `permissions()` in `plan` mode: the main model starts the `writer` + * subagent, and the subagent's `write_file` call does not run. + */ +async function planSubagent(host: HarnessHost, model: () => AnyTextAdapter) { + let writes = 0 + const results: Array = [] + // A fake write_file, so a failing test changes no files. + const writeFile = toolDefinition({ + name: 'write_file', + description: 'Write a file', + inputSchema: z.object({ path: z.string() }), + }).server(async () => { + writes += 1 + return { written: true } + }) + const writer = defineAgent({ + name: 'writer', + description: 'Writes files', + run: (ctx) => + ctx.chat({ + adapter: model(), + messages: [ + { role: 'user', content: '[first-party-plan-writer] write x.txt' }, + ], + tools: [writeFile], + middleware: [ + { + name: 'e2e/tool-results', + onAfterToolCall: (_ctx, info) => { + results.push(info.result) + }, + }, + ], + stream: false, + }), + }) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-plan', + adapter: model(), + subagents: { agents: [writer] }, + plugins: () => [ + permissions({ + rules: [{ tool: 'write_file', decision: 'ask', kind: 'edit' }], + }), + ], + }), + { threadId: 'e2e-plan' }, + ) + await session.command('mode', 'plan') + const turn = await session.prompt('[first-party-plan] write x.txt') + return { writes, results, text: turn.text } +} + +/** + * `projectInstructions()` in a new git repository on the branch `e2e-env`. + * Returns the folder and the platform, so the spec can find the environment + * block in the system prompt of the model call. + */ +async function environmentPrompt( + host: HarnessHost, + model: () => AnyTextAdapter, +) { + const root = await mkdtemp(join(tmpdir(), 'e2e-env-')) + try { + // The branch line needs a commit. + await run('git', ['init', '-q', '-b', 'e2e-env'], { cwd: root }) + await run( + 'git', + [ + '-c', + 'user.name=e2e', + '-c', + 'user.email=e2e@example.com', + 'commit', + '-q', + '--allow-empty', + '-m', + 'init', + ], + { cwd: root }, + ) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-env', + adapter: model(), + plugins: () => [projectInstructions({ root })], + }), + { threadId: 'e2e-env' }, + ) + const turn = await session.prompt('[first-party-env] where am i') + return { root, platform: process.platform, text: turn.text } + } finally { + await rm(root, { recursive: true, force: true }) + } +} + +/** `title()`: the first turn names the session in the session index. */ +async function sessionTitle( + host: HarnessHost, + model: () => AnyTextAdapter, + titleModel: AnyTextAdapter, +) { + const failures: Array = [] + const listener = definePlugin({ + name: 'e2e/title-failures', + setup: (ctx) => { + ctx.on(TitleFailed, (failure) => failures.push(failure.message)) + }, + }) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-title', + adapter: model(), + plugins: () => [listener, title({ adapter: titleModel })], + }), + { threadId: 'e2e-title' }, + ) + const turn = await session.prompt('[first-party-title] plan a trip to Rome') + const deadline = Date.now() + 5_000 + let entry = await host.sessions.get('e2e-title') + while (!entry?.title && failures.length === 0 && Date.now() < deadline) { + await new Promise((resolve) => setTimeout(resolve, 50)) + entry = await host.sessions.get('e2e-title') + } + return { text: turn.text, title: entry?.title, failures } +} + +/** + * First-party harness plugins, one scenario per request: + * `?scenario=agents|plan-subagent|env|title&testId=...`. The models are + * OpenAI adapters against aimock. + */ +export const Route = createFileRoute('/api/harness-first-party')({ + server: { + handlers: { + POST: async ({ request }) => { + const params = new URL(request.url).searchParams + const scenario = params.get('scenario') + const testId = params.get('testId') ?? undefined + const model = () => + createTextAdapter('openai', 'gpt-5.5', undefined, testId).adapter + const host = createHarnessHost({ persistence: memoryPersistence() }) + try { + if (scenario === 'agents') { + return Response.json(await agentProfiles(host, model)) + } + if (scenario === 'plan-subagent') { + return Response.json(await planSubagent(host, model)) + } + if (scenario === 'env') { + return Response.json(await environmentPrompt(host, model)) + } + if (scenario === 'title') { + const titleModel = createTextAdapter( + 'openai', + 'gpt-5.4-nano', + undefined, + testId, + ).adapter + return Response.json(await sessionTitle(host, model, titleModel)) + } + return Response.json( + { error: `Unknown scenario ${scenario}` }, + { status: 400 }, + ) + } catch (error) { + return Response.json({ + error: error instanceof Error ? error.message : String(error), + }) + } finally { + await host.close() + } + }, + }, + }, +}) diff --git a/testing/e2e/tests/harness-first-party.spec.ts b/testing/e2e/tests/harness-first-party.spec.ts new file mode 100644 index 0000000000..058cc9419d --- /dev/null +++ b/testing/e2e/tests/harness-first-party.spec.ts @@ -0,0 +1,125 @@ +import { test, expect } from './fixtures' +import type { APIRequestContext } from '@playwright/test' + +/** + * First-party harness plugins in a session: `agents()`, `permissions()`, + * `projectInstructions()`, and `title()`. Each test runs one scenario of + * `/api/harness-first-party`. The model is the OpenAI adapter against + * aimock, and aimock's journal shows what each model call sent. + */ + +type JournalEntry = { + headers?: Record + body: { + messages?: Array<{ role: string; content: unknown }> + tool_choice?: unknown + } | null +} + +/** This test's model calls, in order, from aimock's journal. */ +async function modelCalls( + request: APIRequestContext, + aimockPort: number, + testId: string, +) { + const journal = await request.get( + `http://127.0.0.1:${aimockPort}/v1/_requests`, + ) + const entries: Array = await journal.json() + return entries.filter((entry) => + Object.entries(entry.headers ?? {}).some( + ([key, value]) => key.toLowerCase() === 'x-test-id' && value === testId, + ), + ) +} + +/** The system prompt text of one model call. */ +const systemOf = (call: JournalEntry | undefined) => + (call?.body?.messages ?? []) + .filter((message) => message.role === 'system') + .map((message) => + typeof message.content === 'string' ? message.content : '', + ) + .join('\n') + +/** Run one scenario of the route. */ +async function runScenario( + request: APIRequestContext, + scenario: string, + testId: string, +) { + const response = await request.post( + `/api/harness-first-party?scenario=${scenario}&testId=${encodeURIComponent(testId)}`, + ) + expect(response.ok()).toBe(true) + return response.json() +} + +test.describe('harness first-party plugins', () => { + test('agents(): /agent switches the profile, and the step limit makes a last call without tools', async ({ + request, + testId, + aimockPort, + }) => { + expect(await runScenario(request, 'agents', testId)).toEqual({ + first: 'Built it.', + switched: 'Agent: brief. It applies at the next turn.', + second: 'Stopped at the step limit.', + }) + + const calls = await modelCalls(request, aimockPort, testId) + expect(calls.map((call) => call.body?.tool_choice)).toEqual([ + undefined, + undefined, + 'none', + ]) + expect(systemOf(calls[0])).toContain('You are the build agent.') + expect(systemOf(calls[1])).toContain('You are the brief agent.') + expect(systemOf(calls[1])).not.toContain('You are the build agent.') + expect(systemOf(calls[1])).not.toContain('step limit') + expect(systemOf(calls[2])).toContain('You reached the step limit.') + }) + + test('permissions(): plan mode stops a write in a subagent', async ({ + request, + testId, + }) => { + expect(await runScenario(request, 'plan-subagent', testId)).toEqual({ + writes: 0, + results: [{ error: 'This tool is not allowed in plan mode.' }], + text: 'The writer could not write in plan mode.', + }) + }) + + test('projectInstructions(): the environment block reaches the model', async ({ + request, + testId, + aimockPort, + }) => { + const body = await runScenario(request, 'env', testId) + expect(body.text).toBe('You are in a new repository.') + + const [call] = await modelCalls(request, aimockPort, testId) + const system = systemOf(call) + expect(system).toMatch(/Environment:\n- Date: \d{4}-\d{2}-\d{2}\n/) + expect(system).toContain( + [ + `- Platform: ${body.platform}`, + `- Working folder: ${body.root}`, + '- Git repository: yes', + '- Git branch: e2e-env', + ].join('\n'), + ) + }) + + test('title(): the first turn names the session in the session index', async ({ + request, + testId, + }) => { + expect(await runScenario(request, 'title', testId)).toEqual({ + text: 'Here is a plan for Rome.', + title: 'Trip to Rome', + failures: [], + }) + }) +}) From 84c615396cfc7f6c69aa1160e790aa0675b2a005 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:08:25 +0200 Subject: [PATCH 02/22] docs(ai-harness): say that a tool with no rule runs in plan mode --- docs/harness/permissions.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/harness/permissions.md b/docs/harness/permissions.md index 15b2a9aa97..9062b3f9bf 100644 --- a/docs/harness/permissions.md +++ b/docs/harness/permissions.md @@ -330,7 +330,7 @@ export const planTools = definePlugin({ ## Known limits -- A rule that allows a tool with no `kind` also lets it run in `plan` mode. `plan` denies only edits, commands, and calls that ask. +- `plan` denies only edits, commands, and calls that ask. A tool with no rule, or with an `allow` rule and no `kind`, still runs in `plan` mode. Give each tool that changes something a rule with a `kind`. - A client that can send inputs to the thread can answer the questions, also with `always`. Control who can open a thread. - `permissions()` does not follow links. Your own tool can check the real path of a link with `PermissionDecisionCapability`. From 397f1837f93f09f188330282f4b198a4bf650ae0 Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Tue, 13 Oct 2026 23:18:13 +0200 Subject: [PATCH 03/22] feat(ai-harness): providerKeys() plugin lets each user connect their own model provider key --- .changeset/ai-harness-provider-keys.md | 5 + docs/config.json | 8 +- docs/harness/overview.md | 2 + docs/harness/provider-keys.md | 204 +++++++++ packages/ai-harness/src/first-party/index.ts | 6 + .../src/first-party/provider-keys.ts | 190 ++++++++ .../ai-harness/tests/provider-keys.test.ts | 406 ++++++++++++++++++ .../fixtures/harness-first-party/basic.json | 8 + .../e2e/src/routes/api.harness-first-party.ts | 57 ++- testing/e2e/tests/harness-first-party.spec.ts | 20 +- 10 files changed, 901 insertions(+), 5 deletions(-) create mode 100644 .changeset/ai-harness-provider-keys.md create mode 100644 docs/harness/provider-keys.md create mode 100644 packages/ai-harness/src/first-party/provider-keys.ts create mode 100644 packages/ai-harness/tests/provider-keys.test.ts diff --git a/.changeset/ai-harness-provider-keys.md b/.changeset/ai-harness-provider-keys.md new file mode 100644 index 0000000000..574005d03f --- /dev/null +++ b/.changeset/ai-harness-provider-keys.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-harness': minor +--- + +Add the `providerKeys({ providers })` plugin to `@tanstack/ai-harness/plugins`. Users of a shipped harness connect their own model provider in the app, with no `.env` file. It adds `/connect:`, `/disconnect:`, and `/keys`. `/connect:` runs the provider's `signIn` when it has one, else it opens `keyUrl` when set and asks the user to paste the key as a secret question. The key is saved in the credential store of the user, and a `keyedAdapter(...)` for that provider uses it. `/keys` and the plugin state show where each key comes from: a saved key (the last 4 characters only), the env var, or no key. diff --git a/docs/config.json b/docs/config.json index df3ace69e3..daf3e14b8a 100644 --- a/docs/config.json +++ b/docs/config.json @@ -926,7 +926,8 @@ { "label": "Build your first harness", "to": "harness/overview", - "addedAt": "2026-10-13" + "addedAt": "2026-10-13", + "updatedAt": "2026-10-13" }, { "label": "Send and show media", @@ -1000,6 +1001,11 @@ "to": "harness/agents", "addedAt": "2026-10-13" }, + { + "label": "Connect model providers", + "to": "harness/provider-keys", + "addedAt": "2026-10-13" + }, { "label": "Run agents from a harness", "to": "harness/subagents", diff --git a/docs/harness/overview.md b/docs/harness/overview.md index 9a73e71e1f..bc7b6caad5 100644 --- a/docs/harness/overview.md +++ b/docs/harness/overview.md @@ -58,6 +58,8 @@ export const assistant = defineHarness({ `agents` are agents you run from code. Put an agent in `subagents` instead when the model must call it as a tool, the same as in `chat()`. +To ship the harness to users with no `.env` file, let each user connect their own key. See [Connect model providers](./provider-keys). + ## 2. Open a session A host runs sessions. Give it the stores from `@tanstack/ai-persistence`, then open a session for a conversation id. diff --git a/docs/harness/provider-keys.md b/docs/harness/provider-keys.md new file mode 100644 index 0000000000..ac9da022f6 --- /dev/null +++ b/docs/harness/provider-keys.md @@ -0,0 +1,204 @@ +--- +title: Connect model providers +id: harness-provider-keys +order: 11 +description: "Let each user connect their own OpenAI, Anthropic, xAI, fal, or OpenRouter key inside your app with /connect. A harness that you ship needs no .env file." +keywords: + - tanstack ai + - harness + - byok + - api key + - providerKeys + - keyedAdapter +--- + +You ship your harness to users: a CLI that they install, a desktop app, or a hosted app with accounts. `openaiText('gpt-6.1-sol')` reads `OPENAI_API_KEY` from the environment. Your users do not have your `.env` file, and they do not start your app from a terminal with a key set. + +With `providerKeys`, each user runs `/connect:openai` in your app and pastes their own key. The harness saves the key for that user and uses it for every turn. + +## 1. Define keyed adapters + +A keyed adapter waits for a key. `keyedAdapter` takes the provider and a function that makes the adapter from the key: + +```ts group=harness-provider-keys +import { keyedAdapter } from '@tanstack/ai' +import { createAnthropicChat } from '@tanstack/ai-anthropic' +import { anthropicByok } from '@tanstack/ai-anthropic/byok' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { openaiByok } from '@tanstack/ai-openai/byok' + +const gpt = keyedAdapter(openaiByok, (key) => createOpenaiChat('gpt-6.1-sol', key)) +const claude = keyedAdapter(anthropicByok, (key) => + createAnthropicChat('claude-sonnet-5-5', key), +) +``` + +- `openaiByok` describes the provider: the id `openai`, the label `OpenAI`, and the env var `OPENAI_API_KEY`. +- Each provider package exports one from its `/byok` subpath, for example `anthropicByok`, `grokByok`, `falByok`, and `openrouterByok`. +- `keyedAdapter` works for every adapter kind: text, image, speech, audio, and video. + +Before each turn, the session calls your function with the key of the user of that session. + +## 2. Add the providerKeys plugin + +Give the harness a keyed adapter as its `adapter`. Then list in `providerKeys` every provider that users can connect: + +```ts group=harness-provider-keys +import { createHarnessHost, defineHarness } from '@tanstack/ai-harness' +import { modelPicker, providerKeys } from '@tanstack/ai-harness/plugins' +import { falByok } from '@tanstack/ai-fal/byok' +import { grokByok } from '@tanstack/ai-grok/byok' +import { openrouterByok } from '@tanstack/ai-openrouter/byok' +import { openrouterSignIn } from '@tanstack/ai-openrouter/pkce' +import { memoryPersistence } from '@tanstack/ai-persistence' + +export const assistant = defineHarness({ + name: 'acme/assistant', + adapter: gpt, + plugins: () => [ + providerKeys({ + providers: [ + { ...openaiByok, keyUrl: 'https://platform.openai.com/api-keys' }, + anthropicByok, + grokByok, + falByok, + { ...openrouterByok, signIn: openrouterSignIn() }, + ], + }), + modelPicker({ choices: { gpt, claude }, default: 'gpt' }), + ], +}) + +const host = createHarnessHost({ persistence: memoryPersistence() }) +const session = await host.open(assistant, { + threadId: 'ada-chat', + principal: { id: 'ada' }, +}) +``` + +- For each provider, `providerKeys` adds `/connect:` and `/disconnect:`. It also adds `/keys`. +- `modelPicker` and `compact` also take keyed adapters. After `/model claude`, the next turn uses the Anthropic key of the user. +- `signIn` is optional. OpenRouter has one. See [Sign in with OpenRouter](#sign-in-with-openrouter). +- `keyUrl` is optional. It is the page where the user makes a key. `/connect:` opens it in the browser, then asks for the key. + +## 3. Connect a provider + +`/connect:openai` asks the user for the key. The question has `secret: true`, so show a password field for it: + +```ts group=harness-provider-keys +const connecting = session.command('connect:openai') + +for await (const { event } of session.events()) { + if (event.type === 'CUSTOM' && event.name === 'harness.question') { + const [question] = session.snapshot().pendingQuestions + // question.secret is true. Read the key without showing it. + if (question) await session.answer(question.questionId, 'sk-proj-...') + break + } +} + +console.log(await connecting) // Connected to OpenAI (key ...abcd). +``` + +- With `keyUrl`, the page to make a key opens first. The session sends a `harness.auth_required` event with the URL, and your app opens it. +- `/disconnect:openai` deletes the key. +- `/keys` shows where each key comes from: + +```text +OpenAI: connected (key ...abcd) +Anthropic: from the ANTHROPIC_API_KEY env var +xAI Grok: missing. Run /connect grok. +fal.ai: missing. Run /connect fal. +OpenRouter: missing. Run /connect openrouter. +``` + +In your own UI, read the same list from `session.snapshot().plugins['tanstack/provider-keys'].providers`. Each entry has an `id`, a `label`, and a `state`: `connected`, `env`, or `missing`. + +### Sign in with OpenRouter + +OpenRouter has a browser sign-in, so users do not copy a key. `openrouterSignIn()` in the providers list turns it on: + +1. Run `/connect:openrouter`. +2. The session sends a `harness.auth_required` event with the OpenRouter sign-in page. Your app opens it in the browser. Sign in and approve. +3. The browser shows `You are signed in to OpenRouter. You can close this tab.` +4. The command returns `Connected to OpenRouter (key ...wxyz).` + +The sign-in uses PKCE. It listens on `127.0.0.1` on a random port for one callback, and it checks a random state. After 10 minutes it stops. To change the time, pass `openrouterSignIn({ timeoutMs })`. + +The listener runs on the computer of the harness. So the sign-in works when the harness and the browser are on the same computer: a CLI or a desktop app. In a hosted app, remove `signIn`, and users paste their OpenRouter key. + +OpenAI and Anthropic have no sign-in for other apps, so users paste a key. + +## When a key is missing + +If a turn needs a key that the user did not connect, the turn stops before the model call: + +- Clients get a `harness.auth_required` event with the provider id. +- The turn fails with `Sign in to openai first. Run /connect openai.` + +After `/connect:openai`, send the message again. + +## Where the keys are stored + +`/connect:` saves the key in the credential store of the host (`stores.credentials`), as an `api_key` credential with the provider id. The store keeps it for the principal of the session: + +- CLI or desktop app: a file in the home folder of the user. +- Hosted app: your database. Each signed-in user is a principal, so each user has their own keys. Write the store with `defineCredentialStore` from `@tanstack/ai-persistence`. +- `memoryPersistence()`: the keys stay in memory until the process stops. + +Encrypt the keys at rest in a real store. + +## Use an Anthropic token + +The key that a user pastes can be a Claude OAuth token (`sk-ant-oat...`). `createAnthropicChat` picks the auth mode from the token, so the keyed adapter above works for it too. See [Bearer and OAuth tokens](../adapters/anthropic#bearer-and-oauth-tokens). + +## The env var still works + +If the user has no saved key, the session reads the env var of the provider, for example `OPENAI_API_KEY`. A saved key wins over the env var. So you can develop with a `.env` file, and your users connect in the app. + +## Use the key of the user in an agent + +Your agent makes images with fal, so it needs the fal key of the user. Every agent gets `ctx.keys`. In a harness, `ctx.keys` reads the same keys as the turn: + +```ts group=harness-provider-keys +import { defineAgent } from '@tanstack/ai' +import { falImage } from '@tanstack/ai-fal' +import { z } from 'zod' + +const flux = keyedAdapter(falByok, (key) => falImage('fal-ai/flux/dev', { apiKey: key })) + +export const illustrator = defineAgent({ + name: 'illustrator', + description: 'Draws an image from a prompt', + inputSchema: z.object({ prompt: z.string() }), + run: async (ctx) => + ctx.generateImage({ + adapter: await ctx.keys.adapter(flux), + prompt: ctx.input.prompt, + }), +}) +``` + +Put the agent in `subagents.agents` of the harness, and the model calls it as a tool. + +- `ctx.keys.adapter(adapter)`: makes a keyed adapter with the key of the user. A plain adapter comes back unchanged. +- `ctx.keys.get(provider)`: the key, or `null`. +- `ctx.keys.require(provider)`: the key. If it is missing, the run stops with `harness.auth_required`, the same as a turn. +- The children of the agent get the same keys. +- Plugins get the same keys as `ctx.keys` in `setup`. + +Outside a harness, `ctx.keys` reads the env var. See [Keys inside agents](../advanced/byok#keys-inside-agents). + +## Keep keys safe + +- A key shows masked only: `...` and the last 4 characters. +- The plugin never puts a key in an event, the transcript, or a command result. +- The answer to the key question does not go into the inbox of the session. +- Do not log a key. For error messages, use `maskKey` from `@tanstack/ai/byok`. + +## What you have now + +- `/connect:` for OpenAI, Anthropic, xAI, fal, and OpenRouter in the harness that you ship, with no `.env` file. +- One key for each user, in your credential store. +- The env var as a fallback while you develop. +- Agents and plugins that use the key of the user through `ctx.keys`. diff --git a/packages/ai-harness/src/first-party/index.ts b/packages/ai-harness/src/first-party/index.ts index 374deb8d06..aa29bdd810 100644 --- a/packages/ai-harness/src/first-party/index.ts +++ b/packages/ai-harness/src/first-party/index.ts @@ -24,6 +24,12 @@ export { formatTodos, todos } from './todos' export type { Todo } from './todos' export { fileCommands, projectInstructions } from './files' export { compact, usage } from './session-tools' +export { providerKeys } from './provider-keys' +export type { + KeyProvider, + ProviderKeyStatus, + ProviderKeysState, +} from './provider-keys' export { agents, builtInAgents } from './agents' export type { AgentProfile } from './agents' export { question } from './question' diff --git a/packages/ai-harness/src/first-party/provider-keys.ts b/packages/ai-harness/src/first-party/provider-keys.ts new file mode 100644 index 0000000000..20b2aaed36 --- /dev/null +++ b/packages/ai-harness/src/first-party/provider-keys.ts @@ -0,0 +1,190 @@ +import { maskKey } from '@tanstack/ai/byok' +import { envKeyOf } from '../auth' +import { defineCommand } from '../commands' +import { definePlugin } from '../plugins' +import type { ByokProvider } from '@tanstack/ai/byok' +import type { CredentialsAccess } from '../auth' +import type { AnyCommand, CommandContext } from '../commands' + +/** Signs the user in and resolves to an API key, for example `openrouterSignIn()`. */ +type ProviderSignIn = (ctx: { + open: (url: string) => void + signal?: AbortSignal +}) => Promise + +/** + * A provider the user can connect: its BYOK descriptor, with a sign-in when + * it has one. `keyUrl` is the provider's page to make a key: `/connect` opens + * it in the browser before it asks for the key. + */ +export type KeyProvider = ByokProvider & { + signIn?: ProviderSignIn + keyUrl?: string +} + +export interface ProviderKeyStatus { + id: string + label: string + /** `connected`: a saved key. `env`: the env var. `missing`: neither. */ + state: 'connected' | 'env' | 'missing' +} + +/** What UIs read from `state.plugins['tanstack/provider-keys']`. */ +export interface ProviderKeysState { + providers: Array +} + +/** + * `promise`, or a rejection once `signal` aborts. A cancelled `/connect` + * stops waiting for the answer, and a late answer saves nothing. + */ +function untilAborted(promise: Promise, signal: AbortSignal) { + signal.throwIfAborted() + return Promise.race([ + promise, + new Promise((_, reject) => + signal.addEventListener('abort', () => reject(signal.reason), { + once: true, + }), + ), + ]) +} + +/** The state of one provider, and its `/keys` line. The key shows masked only. */ +async function describeKey( + credentials: CredentialsAccess, + provider: ByokProvider, +) { + const { id, label } = provider + const status = (state: ProviderKeyStatus['state']) => ({ id, label, state }) + const saved = await credentials.get(id) + if (saved?.type === 'api_key') { + return { + status: status('connected'), + line: `${label}: connected (key ...${maskKey(saved.value)})`, + } + } + const env = envKeyOf(provider) + if (env) { + return { + status: status('env'), + line: `${label}: from the ${env.name} env var`, + } + } + return { + status: status('missing'), + line: `${label}: missing. Run /connect ${id}.`, + } +} + +/** + * Let each user connect their own model providers inside the app, so a + * shipped harness needs no `.env` file. For each provider it adds + * `/connect ` and `/disconnect `, and `/keys` lists where each key + * comes from. `/connect` runs the provider's `signIn` when it has one (the + * host opens the browser), else it asks the user to paste the key. The key + * is saved in the credential store of the session's principal. A + * `keyedAdapter(...)` for that provider then uses it. The env var still works + * when no key is saved. Keys show masked only (the last 4 characters). + * + * @param options.providers - BYOK descriptors, for example `openaiByok`. Add + * `signIn` for a provider with a browser sign-in, for example + * `{ ...openrouterByok, signIn: openrouterSignIn() }`, or `keyUrl` to open + * the provider's key page before the question, for example + * `{ ...openaiByok, keyUrl: 'https://platform.openai.com/api-keys' }`. + * + * @example + * ```ts + * plugins: () => [ + * providerKeys({ + * providers: [ + * openaiByok, + * anthropicByok, + * { ...openrouterByok, signIn: openrouterSignIn() }, + * ], + * }), + * ] + * ``` + */ +export function providerKeys(options: { + providers: ReadonlyArray +}) { + const { providers } = options + return definePlugin({ + name: 'tanstack/provider-keys', + setup: async (ctx) => { + const state = ctx.state({ providers: [] }) + /** + * Check every provider, keep the result for UIs, and return the `/keys` + * lines. A command passes the credentials of the user who runs it. + */ + const refresh = async (credentials = ctx.credentials) => { + const described = await Promise.all( + providers.map((provider) => describeKey(credentials, provider)), + ) + await state.update(() => ({ + providers: described.map((entry) => entry.status), + })) + return described.map((entry) => entry.line).join('\n') + } + await refresh() + + /** + * The key from the provider's sign-in, else the one the user pastes. + * It asks the user who runs the command. + */ + const readKey = async ( + provider: KeyProvider, + { signal, session }: CommandContext, + ) => { + const open = (url: string) => + session.authRequired({ connector: provider.id, url }) + if (provider.signIn) return provider.signIn({ open, signal }) + // No sign-in for other apps: open the page to make a key, then ask. + if (provider.keyUrl) open(provider.keyUrl) + return session.ask({ + message: provider.keyUrl + ? `Paste your ${provider.label} API key. The page to make one is open in your browser.` + : `Paste your ${provider.label} API key`, + secret: true, + }) + } + + const commands: Record = { + keys: defineCommand({ + description: + 'Show the model providers and where their keys come from', + run: (_input, { credentials }) => refresh(credentials), + }), + } + for (const provider of providers) { + const { id, label } = provider + commands[`connect:${id}`] = defineCommand({ + description: `Connect ${label} with your own key`, + // The key belongs to the user who runs the command. + run: async (_input, command) => { + const { signal, credentials } = command + const key = ( + await untilAborted(readKey(provider, command), signal) + ).trim() + if (key === '') { + throw new Error(`No ${label} key was given. Nothing was saved.`) + } + await credentials.set(id, { type: 'api_key', value: key }) + await refresh(credentials) + return `Connected to ${label} (key ...${maskKey(key)}).` + }, + }) + commands[`disconnect:${id}`] = defineCommand({ + description: `Remove your ${label} key`, + run: async (_input, { credentials }) => { + await credentials.delete(id) + await refresh(credentials) + return `Disconnected from ${label}.` + }, + }) + } + return { commands } + }, + }) +} diff --git a/packages/ai-harness/tests/provider-keys.test.ts b/packages/ai-harness/tests/provider-keys.test.ts new file mode 100644 index 0000000000..bc0217bd7b --- /dev/null +++ b/packages/ai-harness/tests/provider-keys.test.ts @@ -0,0 +1,406 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { EventType, defineAgent } from '@tanstack/ai' +import { defineByokProvider, keyedAdapter } from '@tanstack/ai/byok' +import { memoryPersistence } from '@tanstack/ai-persistence' +import { HARNESS_EVENTS, createHarnessHost, defineHarness } from '../src' +import { compact, modelPicker, providerKeys } from '../src/first-party' +import { mockAdapter, text, toolCall } from './helpers' +import type { AnyTextAdapter } from '@tanstack/ai' +import type { + AnyAgent, + HarnessPlugin, + HarnessSession, + SessionEvent, +} from '../src' + +const acme = defineByokProvider({ + id: 'acme', + label: 'Acme', + env: 'ACME_TEST_KEY', +}) +const beta = defineByokProvider({ + id: 'beta', + label: 'Beta', + env: 'BETA_TEST_KEY', +}) +const gamma = defineByokProvider({ id: 'gamma', label: 'Gamma' }) + +/** The credential scope of the test user. */ +const USER = { threadId: 't', userId: 'user-1' } +const SAVED_KEY = 'sk-acme-saved-1234' +const PASTED_KEY = 'sk-acme-pasted-5678' + +afterEach(() => { + vi.unstubAllEnvs() +}) + +/** A keyed adapter for `acme` whose turns are scripted. Records each key it got. */ +function keyedMock(replies: Parameters[0]) { + const mock = mockAdapter(replies) + const keys: Array = [] + const adapter = keyedAdapter(acme, (key) => { + keys.push(key) + return mock.adapter + }) + return { adapter, keys, calls: mock.calls } +} + +async function open(options: { + adapter: AnyTextAdapter | ReturnType['adapter'] + plugins?: Array + agents?: Array + subagents?: Array +}) { + const persistence = memoryPersistence() + const host = createHarnessHost({ persistence }) + const session = await host.open( + defineHarness({ + name: 'test/provider-keys', + adapter: options.adapter, + agents: options.agents ?? [], + subagents: { agents: options.subagents ?? [] }, + plugins: () => options.plugins ?? [], + }), + { threadId: 't', principal: { id: 'user-1' } }, + ) + const events: Array = [] + const reader = new AbortController() + const reading = (async () => { + for await (const entry of session.events({ signal: reader.signal })) + events.push(entry) + })() + const saveKey = (value: string) => + persistence.stores.credentials.set(USER, 'acme', { + type: 'api_key', + value, + }) + const close = async () => { + reader.abort() + await reading + await host.close() + } + return { session, persistence, events, saveKey, close } +} + +/** The connectors named by the `harness.auth_required` events. */ +function authRequired(events: ReadonlyArray) { + return events.flatMap(({ event }) => + event.type === EventType.CUSTOM && + event.name === HARNESS_EVENTS.authRequired + ? [event.value] + : [], + ) +} + +/** Start `/connect acme`, and answer its question with `answer`. */ +async function connectWith(session: HarnessSession, answer: string) { + const result = session.command('connect:acme') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + const receipt = await session.answer(question?.questionId ?? '', answer) + return { question, receipt, result: await result } +} + +describe('a keyed main adapter', () => { + it('uses the key saved in the credential store before the env var', async () => { + vi.stubEnv('ACME_TEST_KEY', 'sk-acme-env-0000') + const keyed = keyedMock([() => text('hello from acme')]) + const { session, saveKey, close } = await open({ adapter: keyed.adapter }) + await saveKey(SAVED_KEY) + expect(await session.prompt('hi')).toEqual({ text: 'hello from acme' }) + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) + + it('reads the env var when no key is saved', async () => { + vi.stubEnv('ACME_TEST_KEY', 'sk-acme-env-0000') + const keyed = keyedMock([() => text('hello')]) + const { session, close } = await open({ adapter: keyed.adapter }) + await session.prompt('hi') + expect(keyed.keys).toEqual(['sk-acme-env-0000']) + await close() + }) + + it('stops the turn and asks for a sign-in when there is no key', async () => { + vi.stubEnv('ACME_TEST_KEY', '') + const keyed = keyedMock([() => text('never')]) + const { session, events, close } = await open({ adapter: keyed.adapter }) + await expect(session.prompt('hi')).rejects.toThrow( + 'Sign in to acme first. Run /connect acme.', + ) + expect(keyed.calls).toHaveLength(0) + expect(authRequired(events)).toEqual([{ connector: 'acme' }]) + await close() + }) + + it('builds a keyed model picker choice after /model', async () => { + const plain = mockAdapter([() => text('plain answer')]) + const keyed = keyedMock([() => text('keyed answer')]) + const { session, saveKey, close } = await open({ + adapter: plain.adapter, + plugins: [ + modelPicker({ + choices: { plain: plain.adapter, acme: keyed.adapter }, + default: 'plain', + }), + ], + }) + await saveKey(SAVED_KEY) + await session.command('model', 'acme') + expect(await session.prompt('hi')).toEqual({ text: 'keyed answer' }) + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) +}) + +describe('ctx.keys in agents', () => { + /** An agent that builds a keyed adapter with `ctx.keys` and answers with it. */ + function writerWith(adapter: ReturnType['adapter']) { + return defineAgent({ + name: 'writer', + description: 'Writes with the Acme model', + run: async (ctx) => + ctx.chat({ adapter: await ctx.keys.adapter(adapter), stream: false }), + }) + } + + it('gives a background agent the session keys', async () => { + const keyed = keyedMock([() => text('agent answer')]) + const { session, saveKey, close } = await open({ + adapter: mockAdapter([]).adapter, + agents: [writerWith(keyed.adapter)], + }) + await saveKey(SAVED_KEY) + expect(await session.agent('writer')?.run()).toBe('agent answer') + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) + + it('gives a subagent the model calls the session keys', async () => { + const keyed = keyedMock([() => text('draft')]) + const lead = mockAdapter([ + () => toolCall('writer', {}), + () => text('The writer is done.'), + ]) + const { session, saveKey, close } = await open({ + adapter: lead.adapter, + subagents: [writerWith(keyed.adapter)], + }) + await saveKey(SAVED_KEY) + expect(await session.prompt('write it')).toEqual({ + text: 'The writer is done.', + }) + expect(keyed.keys).toEqual([SAVED_KEY]) + await close() + }) +}) + +describe('compact with a keyed adapter', () => { + it('compacts with the saved key', async () => { + const summarizer = keyedMock([() => text('We talked.')]) + const main = mockAdapter([() => text('one'), () => text('two')]) + const { session, saveKey, close } = await open({ + adapter: main.adapter, + plugins: [compact({ adapter: summarizer.adapter })], + }) + await saveKey(SAVED_KEY) + await session.prompt('first') + await session.prompt('second') + expect(await session.command('compact')).toBe( + 'Compacted 4 messages into a summary.', + ) + expect(summarizer.keys).toEqual([SAVED_KEY]) + await close() + }) +}) + +describe('providerKeys', () => { + function plugin() { + return providerKeys({ providers: [acme, beta, gamma] }) + } + + it('saves a pasted key and shows it masked', async () => { + const { session, persistence, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + const { question, result } = await connectWith(session, ` ${PASTED_KEY} `) + expect(question).toMatchObject({ + message: 'Paste your Acme API key', + secret: true, + }) + expect(result).toBe('Connected to Acme (key ...5678).') + expect(await persistence.stores.credentials.get(USER, 'acme')).toEqual({ + type: 'api_key', + value: PASTED_KEY, + }) + expect(await session.command('keys')).toContain( + 'Acme: connected (key ...5678)', + ) + await close() + }) + + it('refuses an empty key', async () => { + const { session, persistence, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + await expect(connectWith(session, ' ')).rejects.toThrow( + 'No Acme key was given. Nothing was saved.', + ) + expect(await persistence.stores.credentials.get(USER, 'acme')).toBeNull() + await close() + }) + + it('stops a cancelled /connect, and a late answer saves nothing', async () => { + const { session, persistence, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + const connecting = session.command('connect:acme') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toHaveLength(1), + ) + const [question] = session.snapshot().pendingQuestions + await connecting.cancel() + await expect(connecting).rejects.toBeDefined() + if (question) await session.answer(question.questionId, PASTED_KEY) + expect(await persistence.stores.credentials.get(USER, 'acme')).toBeNull() + await close() + }) + + it('deletes the key on disconnect', async () => { + const { session, persistence, saveKey, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + await saveKey(SAVED_KEY) + expect(await session.command('disconnect:acme')).toBe( + 'Disconnected from Acme.', + ) + expect(await persistence.stores.credentials.get(USER, 'acme')).toBeNull() + expect(await session.command('keys')).toContain( + 'Acme: missing. Run /connect acme.', + ) + await close() + }) + + it('opens the key page of a provider with keyUrl before it asks', async () => { + const { session, persistence, events, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [ + providerKeys({ + providers: [{ ...acme, keyUrl: 'https://acme.example/keys' }], + }), + ], + }) + const { question, result } = await connectWith(session, PASTED_KEY) + expect(authRequired(events)).toEqual([ + { connector: 'acme', url: 'https://acme.example/keys' }, + ]) + expect(question).toMatchObject({ + message: + 'Paste your Acme API key. The page to make one is open in your browser.', + secret: true, + }) + expect(result).toBe('Connected to Acme (key ...5678).') + expect(await persistence.stores.credentials.get(USER, 'acme')).toEqual({ + type: 'api_key', + value: PASTED_KEY, + }) + await close() + }) + + it('runs signIn with open through authRequired and saves its key', async () => { + const signIn = async (ctx: { open: (url: string) => void }) => { + ctx.open('https://acme.example/sign-in') + return 'sk-acme-signed-9876' + } + const { session, persistence, events, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [providerKeys({ providers: [{ ...acme, signIn }] })], + }) + expect(await session.command('connect:acme')).toBe( + 'Connected to Acme (key ...9876).', + ) + expect(authRequired(events)).toEqual([ + { connector: 'acme', url: 'https://acme.example/sign-in' }, + ]) + expect(await persistence.stores.credentials.get(USER, 'acme')).toEqual({ + type: 'api_key', + value: 'sk-acme-signed-9876', + }) + await close() + }) + + it('lists every provider with its state in the plugin state', async () => { + vi.stubEnv('BETA_TEST_KEY', 'sk-beta-env-4321') + const { session, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + expect(session.snapshot().plugins['tanstack/provider-keys']).toEqual({ + providers: [ + { id: 'acme', label: 'Acme', state: 'missing' }, + { id: 'beta', label: 'Beta', state: 'env' }, + { id: 'gamma', label: 'Gamma', state: 'missing' }, + ], + }) + expect(await session.command('keys')).toBe( + [ + 'Acme: missing. Run /connect acme.', + 'Beta: from the BETA_TEST_KEY env var', + 'Gamma: missing. Run /connect gamma.', + ].join('\n'), + ) + await connectWith(session, PASTED_KEY) + expect(session.snapshot().plugins['tanstack/provider-keys']).toEqual({ + providers: [ + { id: 'acme', label: 'Acme', state: 'connected' }, + { id: 'beta', label: 'Beta', state: 'env' }, + { id: 'gamma', label: 'Gamma', state: 'missing' }, + ], + }) + await close() + }) + + it('marks the key question as secret in the snapshot', async () => { + const { session, close } = await open({ + adapter: mockAdapter([]).adapter, + plugins: [plugin()], + }) + const connecting = session.command('connect:acme') + await vi.waitFor(() => + expect(session.snapshot().pendingQuestions).toMatchObject([ + { message: 'Paste your Acme API key', secret: true }, + ]), + ) + const [question] = session.snapshot().pendingQuestions + if (question) await session.answer(question.questionId, PASTED_KEY) + await connecting + await close() + }) + + it('keeps the key out of events, results, the inbox, and the transcript', async () => { + const keyed = keyedMock([() => text('hello')]) + const { session, persistence, events, close } = await open({ + adapter: keyed.adapter, + plugins: [plugin()], + }) + const { receipt, result } = await connectWith(session, PASTED_KEY) + const results = [ + result, + await session.command('keys'), + await session.prompt('hi'), + ] + expect(keyed.keys).toEqual([PASTED_KEY]) + const inbox = await persistence.stores.inbox.get(receipt.inputId) + const transcript = await persistence.stores.messages.loadThread('t') + const everything = JSON.stringify({ events, results, inbox, transcript }) + expect(everything).not.toContain(PASTED_KEY) + expect(everything).toContain('...5678') + await close() + }) +}) diff --git a/testing/e2e/fixtures/harness-first-party/basic.json b/testing/e2e/fixtures/harness-first-party/basic.json index 77195c081f..19361f6e68 100644 --- a/testing/e2e/fixtures/harness-first-party/basic.json +++ b/testing/e2e/fixtures/harness-first-party/basic.json @@ -102,6 +102,14 @@ "response": { "content": "Here is a plan for Rome." } + }, + { + "match": { + "userMessage": "[first-party-keys] hello" + }, + "response": { + "content": "Hello with your key." + } } ] } diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts index 8dc25c4ef9..73db03eef0 100644 --- a/testing/e2e/src/routes/api.harness-first-party.ts +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -4,7 +4,7 @@ import { tmpdir } from 'node:os' import { join } from 'node:path' import { promisify } from 'node:util' import { createFileRoute } from '@tanstack/react-router' -import { defineAgent, toolDefinition } from '@tanstack/ai' +import { defineAgent, keyedAdapter, toolDefinition } from '@tanstack/ai' import { createHarnessHost, defineHarness, @@ -15,13 +15,16 @@ import { agents, permissions, projectInstructions, + providerKeys, title, } from '@tanstack/ai-harness/plugins' +import { createOpenaiChat } from '@tanstack/ai-openai' +import { openaiByok } from '@tanstack/ai-openai/byok' import { memoryPersistence } from '@tanstack/ai-persistence' import { z } from 'zod' import type { AnyTextAdapter } from '@tanstack/ai' import type { HarnessHost } from '@tanstack/ai-harness' -import { createTextAdapter } from '@/lib/providers' +import { LLMOCK_DEFAULT_BASE, createTextAdapter } from '@/lib/providers' const run = promisify(execFile) @@ -194,9 +197,54 @@ async function sessionTitle( return { text: turn.text, title: entry?.title, failures } } +/** The key that the user pastes in the `provider-keys` scenario. */ +const USER_KEY = 'sk-e2e-user-key-7f3a' + +/** + * `providerKeys()`: `/connect:openai` asks for the key as a secret question, + * saves the answer, and the next turn calls the model with that key. The + * spec reads the `authorization` header in aimock's journal. + */ +async function connectedKey(host: HarnessHost, testId: string | undefined) { + const gpt = keyedAdapter(openaiByok, (key) => + createOpenaiChat('gpt-6.1-sol', key, { + baseURL: `${LLMOCK_DEFAULT_BASE}/v1`, + ...(testId ? { defaultHeaders: { 'X-Test-Id': testId } } : {}), + }), + ) + const session = await host.open( + defineHarness({ + name: 'e2e/harness-provider-keys', + adapter: gpt, + plugins: () => [providerKeys({ providers: [openaiByok] })], + }), + { threadId: 'e2e-keys', principal: { id: 'ada' } }, + ) + const connecting = session.command('connect:openai') + const deadline = Date.now() + 5_000 + while ( + session.snapshot().pendingQuestions.length === 0 && + Date.now() < deadline + ) { + await new Promise((resolve) => setTimeout(resolve, 20)) + } + const [question] = session.snapshot().pendingQuestions + // Without a question, `connecting` never settles: fail now. + if (!question) throw new Error('/connect:openai asked no question.') + await session.answer(question.questionId, USER_KEY) + const connected = await connecting + const turn = await session.prompt('[first-party-keys] hello') + return { + secret: question.secret ?? false, + connected, + text: turn.text, + key: USER_KEY, + } +} + /** * First-party harness plugins, one scenario per request: - * `?scenario=agents|plan-subagent|env|title&testId=...`. The models are + * `?scenario=agents|plan-subagent|env|title|provider-keys&testId=...`. The models are * OpenAI adapters against aimock. */ export const Route = createFileRoute('/api/harness-first-party')({ @@ -228,6 +276,9 @@ export const Route = createFileRoute('/api/harness-first-party')({ ).adapter return Response.json(await sessionTitle(host, model, titleModel)) } + if (scenario === 'provider-keys') { + return Response.json(await connectedKey(host, testId)) + } return Response.json( { error: `Unknown scenario ${scenario}` }, { status: 400 }, diff --git a/testing/e2e/tests/harness-first-party.spec.ts b/testing/e2e/tests/harness-first-party.spec.ts index 058cc9419d..de27dbfa2b 100644 --- a/testing/e2e/tests/harness-first-party.spec.ts +++ b/testing/e2e/tests/harness-first-party.spec.ts @@ -3,7 +3,7 @@ import type { APIRequestContext } from '@playwright/test' /** * First-party harness plugins in a session: `agents()`, `permissions()`, - * `projectInstructions()`, and `title()`. Each test runs one scenario of + * `projectInstructions()`, `title()`, and `providerKeys()`. Each test runs one scenario of * `/api/harness-first-party`. The model is the OpenAI adapter against * aimock, and aimock's journal shows what each model call sent. */ @@ -122,4 +122,22 @@ test.describe('harness first-party plugins', () => { failures: [], }) }) + + test('providerKeys(): /connect saves the pasted key, and the next turn uses it', async ({ + request, + testId, + aimockPort, + }) => { + const body = await runScenario(request, 'provider-keys', testId) + expect(body.secret).toBe(true) + expect(body.connected).toMatch(/^Connected to OpenAI \(key \.\.\./) + expect(body.connected).not.toContain(body.key) + expect(body.text).toBe('Hello with your key.') + + const [call] = await modelCalls(request, aimockPort, testId) + const authorization = Object.entries(call?.headers ?? {}).find( + ([key]) => key.toLowerCase() === 'authorization', + )?.[1] + expect(authorization).toBe(`Bearer ${body.key}`) + }) }) From d6d6f8d7fee2ed53198c1a95d1eda9fb1c72bfaa Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 00:28:36 +0200 Subject: [PATCH 04/22] test(ai-e2e): read the key from the adapter factory, aimock redacts the header --- .../e2e/src/routes/api.harness-first-party.ts | 16 ++++++++++------ testing/e2e/tests/harness-first-party.spec.ts | 8 +------- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/testing/e2e/src/routes/api.harness-first-party.ts b/testing/e2e/src/routes/api.harness-first-party.ts index 73db03eef0..b5a54b451d 100644 --- a/testing/e2e/src/routes/api.harness-first-party.ts +++ b/testing/e2e/src/routes/api.harness-first-party.ts @@ -202,16 +202,19 @@ const USER_KEY = 'sk-e2e-user-key-7f3a' /** * `providerKeys()`: `/connect:openai` asks for the key as a secret question, - * saves the answer, and the next turn calls the model with that key. The - * spec reads the `authorization` header in aimock's journal. + * saves the answer, and the next turn builds the model adapter with that key. */ async function connectedKey(host: HarnessHost, testId: string | undefined) { - const gpt = keyedAdapter(openaiByok, (key) => - createOpenaiChat('gpt-6.1-sol', key, { + // The keys the adapter is built with. aimock's journal redacts the + // authorization header, so the spec reads them here. + const keysUsed: Array = [] + const gpt = keyedAdapter(openaiByok, (key) => { + keysUsed.push(key) + return createOpenaiChat('gpt-6.1-sol', key, { baseURL: `${LLMOCK_DEFAULT_BASE}/v1`, ...(testId ? { defaultHeaders: { 'X-Test-Id': testId } } : {}), - }), - ) + }) + }) const session = await host.open( defineHarness({ name: 'e2e/harness-provider-keys', @@ -239,6 +242,7 @@ async function connectedKey(host: HarnessHost, testId: string | undefined) { connected, text: turn.text, key: USER_KEY, + usedUserKey: keysUsed.includes(USER_KEY), } } diff --git a/testing/e2e/tests/harness-first-party.spec.ts b/testing/e2e/tests/harness-first-party.spec.ts index de27dbfa2b..ea841bcdc1 100644 --- a/testing/e2e/tests/harness-first-party.spec.ts +++ b/testing/e2e/tests/harness-first-party.spec.ts @@ -126,18 +126,12 @@ test.describe('harness first-party plugins', () => { test('providerKeys(): /connect saves the pasted key, and the next turn uses it', async ({ request, testId, - aimockPort, }) => { const body = await runScenario(request, 'provider-keys', testId) expect(body.secret).toBe(true) expect(body.connected).toMatch(/^Connected to OpenAI \(key \.\.\./) expect(body.connected).not.toContain(body.key) expect(body.text).toBe('Hello with your key.') - - const [call] = await modelCalls(request, aimockPort, testId) - const authorization = Object.entries(call?.headers ?? {}).find( - ([key]) => key.toLowerCase() === 'authorization', - )?.[1] - expect(authorization).toBe(`Bearer ${body.key}`) + expect(body.usedUserKey).toBe(true) }) }) From 29bb0a487f8b8febebcbd891bead0ab7aa3ba987 Mon Sep 17 00:00:00 2001 From: Patrick L Date: Wed, 14 Oct 2026 12:03:58 +0200 Subject: [PATCH 05/22] fix(ai-persistence): store a tool result once when a run resumes from a client tool (#1744) * fix(ai-persistence): store a tool result once when a run resumes from a client tool * test(e2e): client tool resume stores each tool result once * fix(ai-persistence): keep a later tool result that reuses a stored toolCallId Count stored tool results per toolCallId and let each posted copy use up one, instead of dropping every posted result whose toolCallId the prefix holds. Check before the id check, so a result posted under its stored id also uses up its count. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AGyaNPEJb8zBzY2wY5xWdZ * Revert "fix(ai-persistence): keep a later tool result that reuses a stored toolCallId" Provider tool call ids are unique within a thread, so the reuse case does not occur. Back to the plain toolCallId match. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01AGyaNPEJb8zBzY2wY5xWdZ * fix(ai-persistence): replace a tool result matched by toolCallId in place Skip the toolCallId match when the incoming tool result has no toolCallId, so two id-less results are not treated as one. Replace the matched stored row with the incoming copy instead of dropping it, the same as a match by id. Move the merge-rules comment onto mergeStoredMessages and add direct merge tests. --------- Co-authored-by: Claude Opus 5.5 Co-authored-by: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> --- .changeset/persistence-resume-tool-results.md | 5 + packages/ai-persistence/src/merge-stored.ts | 24 +- .../ai-persistence/tests/merge-stored.test.ts | 42 ++++ .../tests/message-ids-across-calls.test.ts | 81 ++++++ testing/e2e/src/routeTree.gen.ts | 22 ++ .../routes/api.client-tool-resume-store.ts | 230 ++++++++++++++++++ .../tests/client-tool-resume-store.spec.ts | 26 ++ 7 files changed, 425 insertions(+), 5 deletions(-) create mode 100644 .changeset/persistence-resume-tool-results.md create mode 100644 packages/ai-persistence/tests/merge-stored.test.ts create mode 100644 testing/e2e/src/routes/api.client-tool-resume-store.ts create mode 100644 testing/e2e/tests/client-tool-resume-store.spec.ts diff --git a/.changeset/persistence-resume-tool-results.md b/.changeset/persistence-resume-tool-results.md new file mode 100644 index 0000000000..032f553bfd --- /dev/null +++ b/.changeset/persistence-resume-tool-results.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-persistence': patch +--- + +Store a tool result once when a run resumes from a client-tool interrupt. The client posts earlier tool results under interrupt snapshot ids the store never saw, so `withPersistence` appended a second copy and the next provider request failed. diff --git a/packages/ai-persistence/src/merge-stored.ts b/packages/ai-persistence/src/merge-stored.ts index 4c8ba03c81..6b29219b98 100644 --- a/packages/ai-persistence/src/merge-stored.ts +++ b/packages/ai-persistence/src/merge-stored.ts @@ -1,10 +1,5 @@ import type { ModelMessage } from '@tanstack/ai' -// Empty incoming keeps stored. Non-empty: the last incoming id that already -// exists in stored is a cutoff (reload drops the old assistant after that -// user). Same id is replaced in place, but a stored mid-conversation record -// stays when the incoming copy has none. New ids and messages with no id are -// appended. /** How many stored messages the merge keeps. Less than `stored.length` on a reload. */ export function storedCutoff( stored: ReadonlyArray, @@ -19,6 +14,12 @@ export function storedCutoff( return stored.length } +// Empty incoming keeps stored. Non-empty: the last incoming id that already +// exists in stored is a cutoff (reload drops the old assistant after that +// user). Same id, or a tool result with the same toolCallId, is replaced in +// place, but a stored mid-conversation record stays when the incoming copy has +// none. An id-less message that matches the prefix row at its index is +// dropped. Other messages are appended. export function mergeStoredMessages( stored: ReadonlyArray, incoming: ReadonlyArray, @@ -55,6 +56,19 @@ export function mergeStoredMessages( if (!message) continue const id = message.id if (id && storedIds.has(id)) continue + // A resumed run posts stored tool results under the interrupt snapshot's + // ids, which the store never saw. Match tool results by toolCallId. + if (message.role === 'tool' && message.toolCallId !== undefined) { + const storedIndex = prefix.findIndex( + (stored) => + stored.role === 'tool' && stored.toolCallId === message.toolCallId, + ) + const stored = prefix[storedIndex] + if (stored) { + merged[storedIndex] = keepStoredRecord(message, stored) + continue + } + } // Attach/reload can post the stored transcript again with no ids. Keep the // prefix row instead of appending a second copy of the same turn. if (!id) { diff --git a/packages/ai-persistence/tests/merge-stored.test.ts b/packages/ai-persistence/tests/merge-stored.test.ts new file mode 100644 index 0000000000..0e4ca54c0a --- /dev/null +++ b/packages/ai-persistence/tests/merge-stored.test.ts @@ -0,0 +1,42 @@ +import { describe, expect, it } from 'vitest' +import { mergeStoredMessages } from '../src/merge-stored' +import type { ModelMessage } from '@tanstack/ai' + +const user: ModelMessage = { id: 'u1', role: 'user', content: 'hi' } +const assistant: ModelMessage = { + id: 'a1', + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call_1', + type: 'function', + function: { name: 'lookup', arguments: '{}' }, + }, + ], +} +const result = (content: string, id?: string): ModelMessage => ({ + ...(id && { id }), + role: 'tool', + toolCallId: 'call_1', + content, +}) + +describe('mergeStoredMessages: tool results', () => { + it('replaces a stored tool result posted under a new id in place', () => { + const merged = mergeStoredMessages( + [user, assistant, result('old')], + [user, assistant, result('new', 'snapshot_r1_2')], + ) + expect(merged).toEqual([user, assistant, result('new', 'snapshot_r1_2')]) + }) + + it('appends tool results that have no toolCallId', () => { + const stored: ModelMessage = { role: 'tool', content: 'a' } + const incoming: ModelMessage = { id: 'x', role: 'tool', content: 'b' } + expect(mergeStoredMessages([stored], [incoming])).toEqual([ + stored, + incoming, + ]) + }) +}) diff --git a/packages/ai-persistence/tests/message-ids-across-calls.test.ts b/packages/ai-persistence/tests/message-ids-across-calls.test.ts index 259f5e5668..6907a6db25 100644 --- a/packages/ai-persistence/tests/message-ids-across-calls.test.ts +++ b/packages/ai-persistence/tests/message-ids-across-calls.test.ts @@ -119,6 +119,22 @@ function model(withThinking: boolean) { ] } +const toolCall = ( + toolCallId: string, + toolCallName: string, + parentMessageId: string, +): Array => [ + { + type: EventType.TOOL_CALL_START, + toolCallId, + toolCallName, + parentMessageId, + timestamp: t, + }, + { type: EventType.TOOL_CALL_ARGS, toolCallId, delta: '{}', timestamp: t }, + { type: EventType.TOOL_CALL_END, toolCallId, timestamp: t }, +] + const call1Count = (messages: ReadonlyArray) => messages.flatMap((m) => m.toolCalls ?? []).filter((c) => c.id === 'call_1') .length @@ -196,4 +212,69 @@ describe('useChat + withPersistence: message ids across model calls', () => { expect(call1Count(stored)).toBe(1) }) } + + it('stores each tool result once when a run resumes from a client tool', async () => { + const persistence = memoryPersistence() + // Server tool, then a client tool that interrupts the run, then text. + const { adapter, prompts } = mockAdapter([ + run('r1', toolCall('call_1', 'lookup', 'm1'), 'tool_calls'), + run('r1', toolCall('call_2', 'ask', 'm2'), 'tool_calls'), + run('r1', text('m3', 'done'), 'stop'), + ]) + const tools = [ + { name: 'lookup', description: 'Lookup', execute: () => ({ ok: true }) }, + { name: 'ask', description: 'Ask the user' }, + ] + const middleware = [withPersistence(persistence)] + + const processor = new StreamProcessor() + const user1: UIMessage = { + id: 'u1', + role: 'user', + parts: [{ type: 'text', content: 'hi' }], + } + processor.setMessages([user1]) + processor.prepareAssistantMessage() + for await (const chunk of chat({ + adapter, + messages: await serverMessages([user1], 'r1'), + tools, + runId: 'r1', + threadId: 't1', + middleware, + }) as AsyncIterable) { + processor.processChunk(chunk) + } + processor.finalizeStream() + + // The interrupt snapshot gave the stored server tool result a snapshot id. + processor.addToolResult('call_2', { answer: 'yes' }) + for await (const _ of chat({ + adapter, + messages: await serverMessages(processor.getMessages(), 'r1'), + tools, + runId: 'r1', + threadId: 't1', + resume: [ + { + interruptId: 'client_tool_call_2', + status: 'resolved', + payload: { answer: 'yes' }, + }, + ], + middleware, + }) as AsyncIterable) { + // drain + } + + const stored = threadMessages( + await persistence.stores.messages!.loadThread('t1'), + ) + expect( + stored.filter((m) => m.role === 'tool').map((m) => m.toolCallId), + ).toEqual(['call_1', 'call_2']) + expect( + prompts[2]!.filter((m) => m.role === 'tool').map((m) => m.toolCallId), + ).toEqual(['call_1', 'call_2']) + }) }) diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index aa9537f6db..275d0cba66 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -73,6 +73,7 @@ import { Route as ApiByteplusEncryptedReasoningWireRouteImport } from './routes/ import { Route as ApiByteplusSeedance1080pWireRouteImport } from './routes/api.byteplus-seedance-1080p-wire' 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 ApiCompactionUsageWireRouteImport } from './routes/api.compaction-usage-wire' import { Route as ApiCompactionWireRouteImport } from './routes/api.compaction-wire' @@ -514,6 +515,12 @@ const ApiChatCompletionsIncompleteStreamRoute = path: '/api/chat-completions-incomplete-stream', getParentRoute: () => rootRouteImport, } as any) +const ApiClientToolResumeStoreRoute = + ApiClientToolResumeStoreRouteImport.update({ + id: '/api/client-tool-resume-store', + path: '/api/client-tool-resume-store', + getParentRoute: () => rootRouteImport, + } as any) const ApiCloudflareBindingWireRoute = ApiCloudflareBindingWireRouteImport.update({ id: '/api/cloudflare-binding-wire', @@ -1124,6 +1131,7 @@ export interface FileRoutesByFullPath { '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute @@ -1294,6 +1302,7 @@ export interface FileRoutesByTo { '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute @@ -1465,6 +1474,7 @@ export interface FileRoutesById { '/api/byteplus-seedance-1080p-wire': typeof ApiByteplusSeedance1080pWireRoute '/api/chat': typeof ApiChatRoute '/api/chat-completions-incomplete-stream': typeof ApiChatCompletionsIncompleteStreamRoute + '/api/client-tool-resume-store': typeof ApiClientToolResumeStoreRoute '/api/cloudflare-binding-wire': typeof ApiCloudflareBindingWireRoute '/api/compaction-usage-wire': typeof ApiCompactionUsageWireRoute '/api/compaction-wire': typeof ApiCompactionWireRoute @@ -1637,6 +1647,7 @@ export interface FileRouteTypes { | '/api/byteplus-seedance-1080p-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/compaction-usage-wire' | '/api/compaction-wire' @@ -1807,6 +1818,7 @@ export interface FileRouteTypes { | '/api/byteplus-seedance-1080p-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/compaction-usage-wire' | '/api/compaction-wire' @@ -1977,6 +1989,7 @@ export interface FileRouteTypes { | '/api/byteplus-seedance-1080p-wire' | '/api/chat' | '/api/chat-completions-incomplete-stream' + | '/api/client-tool-resume-store' | '/api/cloudflare-binding-wire' | '/api/compaction-usage-wire' | '/api/compaction-wire' @@ -2148,6 +2161,7 @@ export interface RootRouteChildren { ApiByteplusSeedance1080pWireRoute: typeof ApiByteplusSeedance1080pWireRoute ApiChatRoute: typeof ApiChatRoute ApiChatCompletionsIncompleteStreamRoute: typeof ApiChatCompletionsIncompleteStreamRoute + ApiClientToolResumeStoreRoute: typeof ApiClientToolResumeStoreRoute ApiCloudflareBindingWireRoute: typeof ApiCloudflareBindingWireRoute ApiCompactionUsageWireRoute: typeof ApiCompactionUsageWireRoute ApiCompactionWireRoute: typeof ApiCompactionWireRoute @@ -2700,6 +2714,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiChatCompletionsIncompleteStreamRouteImport parentRoute: typeof rootRouteImport } + '/api/client-tool-resume-store': { + id: '/api/client-tool-resume-store' + path: '/api/client-tool-resume-store' + fullPath: '/api/client-tool-resume-store' + preLoaderRoute: typeof ApiClientToolResumeStoreRouteImport + parentRoute: typeof rootRouteImport + } '/api/cloudflare-binding-wire': { id: '/api/cloudflare-binding-wire' path: '/api/cloudflare-binding-wire' @@ -3558,6 +3579,7 @@ const rootRouteChildren: RootRouteChildren = { ApiChatRoute: ApiChatRoute, ApiChatCompletionsIncompleteStreamRoute: ApiChatCompletionsIncompleteStreamRoute, + ApiClientToolResumeStoreRoute: ApiClientToolResumeStoreRoute, ApiCloudflareBindingWireRoute: ApiCloudflareBindingWireRoute, ApiCompactionUsageWireRoute: ApiCompactionUsageWireRoute, ApiCompactionWireRoute: ApiCompactionWireRoute, diff --git a/testing/e2e/src/routes/api.client-tool-resume-store.ts b/testing/e2e/src/routes/api.client-tool-resume-store.ts new file mode 100644 index 0000000000..1c39c0f4bc --- /dev/null +++ b/testing/e2e/src/routes/api.client-tool-resume-store.ts @@ -0,0 +1,230 @@ +import { createFileRoute } from '@tanstack/react-router' +import { + chat, + chatParamsFromRequestBody, + createChatOptions, +} from '@tanstack/ai' +import { StreamProcessor, uiMessagesToWire } from '@tanstack/ai/client' +import { createAnthropicChat } from '@tanstack/ai-anthropic' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' + +const DUMMY_KEY = 'sk-ant-e2e-test-dummy-key' + +type SseEvent = Record & { type: string } + +function messageStart(id: string): SseEvent { + return { + type: 'message_start', + message: { + id, + type: 'message', + role: 'assistant', + content: [], + model: 'claude-sonnet-4-5', + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 5, output_tokens: 0 }, + }, + } +} + +function text(index: number, value: string) { + return [ + { + type: 'content_block_start', + index, + content_block: { type: 'text', text: '' }, + }, + { + type: 'content_block_delta', + index, + delta: { type: 'text_delta', text: value }, + }, + { type: 'content_block_stop', index }, + ] +} + +function toolUse(index: number, id: string, name: string) { + return [ + { + type: 'content_block_start', + index, + content_block: { type: 'tool_use', id, name, input: {} }, + }, + { + type: 'content_block_delta', + index, + delta: { type: 'input_json_delta', partial_json: '{}' }, + }, + { type: 'content_block_stop', index }, + ] +} + +function end(stopReason: string) { + return [ + { + type: 'message_delta', + delta: { stop_reason: stopReason, stop_sequence: null }, + usage: { output_tokens: 20 }, + }, + { type: 'message_stop' }, + ] +} + +function sse(events: Array): Response { + return new Response( + events + .map( + (event) => `event: ${event.type}\ndata: ${JSON.stringify(event)}\n\n`, + ) + .join(''), + { status: 200, headers: { 'Content-Type': 'text/event-stream' } }, + ) +} + +/** Answers each call with the next scripted stream and records its body. */ +function scriptedFetch(streams: Array>) { + const bodies: Array<{ + messages: Array<{ role: string; content: unknown }> + }> = [] + const fetchImpl: typeof fetch = async (input, init) => { + const req = input instanceof Request ? input : new Request(input, init) + bodies.push(JSON.parse(await req.text())) + return sse(streams[bodies.length - 1] ?? []) + } + return { bodies, fetchImpl } +} + +/** Every `tool_result` id the request sends back to Claude. */ +function toolResultIds(body: { messages: Array<{ content: unknown }> }) { + return body.messages.flatMap((message) => + Array.isArray(message.content) + ? message.content + .filter((block: { type: string }) => block.type === 'tool_result') + .map((block: { tool_use_id: string }) => block.tool_use_id) + : [], + ) +} + +/** + * Claude calls a server tool, then a client tool, and the run stops on the + * client-tool interrupt. The browser takes the interrupt `MESSAGES_SNAPSHOT`, + * which gives the stored server tool result a snapshot id, adds the client + * tool's output, and resumes the run. `withPersistence` merges what it posts + * with the stored transcript. + * + * Before the fix, the store had never seen the snapshot id, so the merge + * appended the server tool result a second time and the resumed call sent + * Claude a second `tool_result` for `toolu_lookup`, with no matching `tool_use` + * in the assistant turn before it. + * + * `StreamProcessor` → `uiMessagesToWire` → `chatParamsFromRequestBody` → + * `chat()`, as a browser client and `/api/chat` do, answered by a capturing + * `fetch` (no aimock fixture). + */ +export const Route = createFileRoute('/api/client-tool-resume-store')({ + server: { + handlers: { + POST: async () => { + try { + const claude = scriptedFetch([ + [ + messageStart('msg_call_1'), + ...toolUse(0, 'toolu_lookup', 'lookup'), + ...end('tool_use'), + ], + [ + messageStart('msg_call_2'), + ...toolUse(0, 'toolu_ask', 'ask'), + ...end('tool_use'), + ], + [ + messageStart('msg_call_3'), + ...text(0, 'It is 42.'), + ...end('end_turn'), + ], + ]) + const adapter = createAnthropicChat('claude-sonnet-4-5', DUMMY_KEY, { + fetch: claude.fetchImpl, + }) + const persistence = memoryPersistence() + const tools = [ + { + name: 'lookup', + description: 'Look up the answer.', + execute: () => ({ answer: 42 }), + }, + { name: 'ask', description: 'Ask the user.' }, + ] + const threadId = 'client-tool-resume-store' + const runId = 'client-tool-resume-store-1' + const client = new StreamProcessor({}) + client.addUserMessage('What is the answer?') + + const send = async ( + resume?: Array<{ + interruptId: string + status: 'resolved' + payload: unknown + }>, + ) => { + const params = await chatParamsFromRequestBody({ + threadId, + runId, + // `JSON.parse(JSON.stringify(...))` stands in for the HTTP hop. + messages: JSON.parse( + JSON.stringify(uiMessagesToWire(client.getMessages())), + ), + tools: [], + context: [], + ...(resume && { resume }), + }) + client.prepareAssistantMessage() + for await (const chunk of chat({ + ...createChatOptions({ adapter }), + messages: params.messages, + ...(params.resume && { resume: params.resume }), + tools, + threadId, + runId, + middleware: [withPersistence(persistence)], + stream: true, + })) { + client.processChunk(chunk) + } + client.finalizeStream() + } + + await send() + client.addToolResult('toolu_ask', { confirmed: true }) + await send([ + { + interruptId: 'client_tool_toolu_ask', + status: 'resolved', + payload: { confirmed: true }, + }, + ]) + + const stored = + (await persistence.stores.messages?.loadThread(threadId)) ?? [] + return Response.json({ + ok: true, + resumedToolResultIds: claude.bodies[2] + ? toolResultIds(claude.bodies[2]) + : [], + storedToolResultIds: stored.flatMap((message) => + message.role === 'tool' && message.toolCallId + ? [message.toolCallId] + : [], + ), + }) + } catch (error) { + return Response.json({ + ok: false, + error: error instanceof Error ? error.message : String(error), + }) + } + }, + }, + }, +}) diff --git a/testing/e2e/tests/client-tool-resume-store.spec.ts b/testing/e2e/tests/client-tool-resume-store.spec.ts new file mode 100644 index 0000000000..4fea22c875 --- /dev/null +++ b/testing/e2e/tests/client-tool-resume-store.spec.ts @@ -0,0 +1,26 @@ +import { test, expect } from './fixtures' + +/** + * Claude calls a server tool, then a client tool that interrupts the run. With + * `withPersistence`, resuming the run must send and store each tool result + * once. `/api/client-tool-resume-store` runs the turn and its resume and + * returns what Claude received and what the store kept. + */ +test.describe('client tool resume — tool results in the store', () => { + test('the resumed call and the store hold each tool result once', async ({ + request, + }) => { + const response = await request.post('/api/client-tool-resume-store') + expect(response.ok()).toBe(true) + const result = (await response.json()) as { + ok: boolean + error?: string + resumedToolResultIds: Array + storedToolResultIds: Array + } + if (!result.ok) throw new Error(`Route failed: ${result.error}`) + + expect(result.resumedToolResultIds).toEqual(['toolu_lookup', 'toolu_ask']) + expect(result.storedToolResultIds).toEqual(['toolu_lookup', 'toolu_ask']) + }) +}) From 761a652a3d3f8d31017e11034a4844c8c81af0f6 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:47:10 +1100 Subject: [PATCH 06/22] fix(openai-base): report Responses cache write tokens in usage (#1774) * fix(openai-base): map Responses input_tokens_details.cache_write_tokens to cacheWriteTokens Closes #1734 * fix(openai-base): keep the cache-write fix to one block; add Responses usage e2e --- .changeset/responses-cache-write-tokens.md | 5 ++ packages/openai-base/src/usage.ts | 13 +++- .../openai-base/tests/responses-usage.test.ts | 22 +++++++ testing/e2e/global-setup.ts | 65 +++++++++++++++++++ testing/e2e/src/routeTree.gen.ts | 21 ++++++ .../src/routes/api.openai-responses-usage.ts | 49 ++++++++++++++ .../e2e/tests/openai-responses-usage.spec.ts | 29 +++++++++ 7 files changed, 202 insertions(+), 2 deletions(-) create mode 100644 .changeset/responses-cache-write-tokens.md create mode 100644 packages/openai-base/tests/responses-usage.test.ts create mode 100644 testing/e2e/src/routes/api.openai-responses-usage.ts create mode 100644 testing/e2e/tests/openai-responses-usage.spec.ts diff --git a/.changeset/responses-cache-write-tokens.md b/.changeset/responses-cache-write-tokens.md new file mode 100644 index 0000000000..19c4667afc --- /dev/null +++ b/.changeset/responses-cache-write-tokens.md @@ -0,0 +1,5 @@ +--- +'@tanstack/openai-base': patch +--- + +Responses API usage now reports `promptTokensDetails.cacheWriteTokens` from `input_tokens_details.cache_write_tokens`, the same as the Chat Completions path. diff --git a/packages/openai-base/src/usage.ts b/packages/openai-base/src/usage.ts index 2555494bf0..e0a8387faf 100644 --- a/packages/openai-base/src/usage.ts +++ b/packages/openai-base/src/usage.ts @@ -90,8 +90,8 @@ export function buildChatCompletionsUsage( * `ResponseUsage` object. * * Shared by every provider that routes through - * {@link OpenAIBaseResponsesTextAdapter}. Surfaces cached prompt tokens and - * reasoning detail tokens when present. Returns `undefined` when the provider + * {@link OpenAIBaseResponsesTextAdapter}. Surfaces cache read, cache write, + * and reasoning detail tokens when present. Returns `undefined` when the provider * reported no usage object, so callers omit the field rather than fabricating * zeroed totals. */ @@ -116,6 +116,15 @@ export function buildResponsesUsage( } } + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition + const cacheWriteTokens = usage.input_tokens_details?.cache_write_tokens + if (cacheWriteTokens && cacheWriteTokens > 0) { + result.promptTokensDetails = { + ...result.promptTokensDetails, + cacheWriteTokens, + } + } + // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition const reasoningTokens = usage.output_tokens_details?.reasoning_tokens if (reasoningTokens && reasoningTokens > 0) { diff --git a/packages/openai-base/tests/responses-usage.test.ts b/packages/openai-base/tests/responses-usage.test.ts new file mode 100644 index 0000000000..ad41790bc9 --- /dev/null +++ b/packages/openai-base/tests/responses-usage.test.ts @@ -0,0 +1,22 @@ +import { describe, expect, it } from 'vitest' +import { buildResponsesUsage } from '../src/usage' + +// OpenAI Responses usage for a call that writes the prompt cache (issue #1734). +describe('buildResponsesUsage', () => { + it('maps cache reads and cache writes', () => { + expect( + buildResponsesUsage({ + input_tokens: 1941, + input_tokens_details: { cached_tokens: 3, cache_write_tokens: 1938 }, + output_tokens: 5, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 1946, + }), + ).toEqual({ + promptTokens: 1941, + completionTokens: 5, + totalTokens: 1946, + promptTokensDetails: { cachedTokens: 3, cacheWriteTokens: 1938 }, + }) + }) +}) diff --git a/testing/e2e/global-setup.ts b/testing/e2e/global-setup.ts index f3a2cba472..30b480197b 100644 --- a/testing/e2e/global-setup.ts +++ b/testing/e2e/global-setup.ts @@ -157,6 +157,9 @@ export default async function globalSetup() { // Provider-executed web search responses need native Responses/Gemini wire // items that aimock does not synthesize. These mounts feed the adapter E2E // route deterministic search calls and grounding metadata. + // OpenAI Responses usage for a call that writes the prompt cache (#1734). + // aimock's native Responses helper cannot emit `input_tokens_details`. + mock.mount('/openai-responses-usage', openaiResponsesUsageMount()) mock.mount('/provider-search-openai', openaiProviderSearchMount()) mock.mount('/provider-search-gemini', geminiProviderSearchMount()) @@ -1958,6 +1961,68 @@ function openaiUsageDetailsMount( } } +function openaiResponsesUsageMount(): Mountable { + return { + async handleRequest( + req: http.IncomingMessage, + res: http.ServerResponse, + pathname: string, + ): Promise { + if (req.method !== 'POST' || pathname !== '/v1/responses') { + return false + } + await drainBody(req) + + const message = { + id: 'msg_e2e_usage', + type: 'message', + role: 'assistant', + status: 'completed', + content: [{ type: 'output_text', text: 'Hi', annotations: [] }], + } + const events = [ + { + type: 'response.output_text.delta', + item_id: 'msg_e2e_usage', + output_index: 0, + content_index: 0, + delta: 'Hi', + }, + { + type: 'response.completed', + response: { + id: 'resp_e2e_usage', + object: 'response', + status: 'completed', + model: 'gpt-4o', + output: [message], + usage: { + input_tokens: 1941, + input_tokens_details: { + cached_tokens: 0, + cache_write_tokens: 1938, + }, + output_tokens: 5, + output_tokens_details: { reasoning_tokens: 0 }, + total_tokens: 1946, + }, + }, + }, + ] + + res.statusCode = 200 + res.setHeader('Content-Type', 'text/event-stream') + res.setHeader('Cache-Control', 'no-cache') + for (const event of events) { + res.write(`data: ${JSON.stringify(event)}\n\n`) + } + res.write('data: [DONE]\n\n') + res.end() + return true + }, + } +} + function openaiProviderSearchMount(): Mountable { return { async handleRequest( diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 275d0cba66..6fb7e92c61 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -128,6 +128,7 @@ import { Route as ApiNonStreamingRunErrorRouteImport } from './routes/api.non-st import { Route as ApiOpenaiCompletedResponseTextRouteImport } from './routes/api.openai-completed-response-text' import { Route as ApiOpenaiImage25ModelsRouteImport } from './routes/api.openai-image-2-5-models' import { Route as ApiOpenaiMalformedToolArgumentsRouteImport } from './routes/api.openai-malformed-tool-arguments' +import { Route as ApiOpenaiResponsesUsageRouteImport } from './routes/api.openai-responses-usage' import { Route as ApiOpenaiShellSkillsWireRouteImport } from './routes/api.openai-shell-skills-wire' import { Route as ApiOpenaiStrictToolNullWireRouteImport } from './routes/api.openai-strict-tool-null-wire' import { Route as ApiOpenaiUsageDetailsRouteImport } from './routes/api.openai-usage-details' @@ -800,6 +801,11 @@ const ApiOpenaiMalformedToolArgumentsRoute = path: '/api/openai-malformed-tool-arguments', getParentRoute: () => rootRouteImport, } as any) +const ApiOpenaiResponsesUsageRoute = ApiOpenaiResponsesUsageRouteImport.update({ + id: '/api/openai-responses-usage', + path: '/api/openai-responses-usage', + getParentRoute: () => rootRouteImport, +} as any) const ApiOpenaiShellSkillsWireRoute = ApiOpenaiShellSkillsWireRouteImport.update({ id: '/api/openai-shell-skills-wire', @@ -1186,6 +1192,7 @@ export interface FileRoutesByFullPath { '/api/openai-completed-response-text': typeof ApiOpenaiCompletedResponseTextRoute '/api/openai-image-2-5-models': typeof ApiOpenaiImage25ModelsRoute '/api/openai-malformed-tool-arguments': typeof ApiOpenaiMalformedToolArgumentsRoute + '/api/openai-responses-usage': typeof ApiOpenaiResponsesUsageRoute '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute '/api/openai-strict-tool-null-wire': typeof ApiOpenaiStrictToolNullWireRoute '/api/openai-usage-details': typeof ApiOpenaiUsageDetailsRoute @@ -1357,6 +1364,7 @@ export interface FileRoutesByTo { '/api/openai-completed-response-text': typeof ApiOpenaiCompletedResponseTextRoute '/api/openai-image-2-5-models': typeof ApiOpenaiImage25ModelsRoute '/api/openai-malformed-tool-arguments': typeof ApiOpenaiMalformedToolArgumentsRoute + '/api/openai-responses-usage': typeof ApiOpenaiResponsesUsageRoute '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute '/api/openai-strict-tool-null-wire': typeof ApiOpenaiStrictToolNullWireRoute '/api/openai-usage-details': typeof ApiOpenaiUsageDetailsRoute @@ -1529,6 +1537,7 @@ export interface FileRoutesById { '/api/openai-completed-response-text': typeof ApiOpenaiCompletedResponseTextRoute '/api/openai-image-2-5-models': typeof ApiOpenaiImage25ModelsRoute '/api/openai-malformed-tool-arguments': typeof ApiOpenaiMalformedToolArgumentsRoute + '/api/openai-responses-usage': typeof ApiOpenaiResponsesUsageRoute '/api/openai-shell-skills-wire': typeof ApiOpenaiShellSkillsWireRoute '/api/openai-strict-tool-null-wire': typeof ApiOpenaiStrictToolNullWireRoute '/api/openai-usage-details': typeof ApiOpenaiUsageDetailsRoute @@ -1702,6 +1711,7 @@ export interface FileRouteTypes { | '/api/openai-completed-response-text' | '/api/openai-image-2-5-models' | '/api/openai-malformed-tool-arguments' + | '/api/openai-responses-usage' | '/api/openai-shell-skills-wire' | '/api/openai-strict-tool-null-wire' | '/api/openai-usage-details' @@ -1873,6 +1883,7 @@ export interface FileRouteTypes { | '/api/openai-completed-response-text' | '/api/openai-image-2-5-models' | '/api/openai-malformed-tool-arguments' + | '/api/openai-responses-usage' | '/api/openai-shell-skills-wire' | '/api/openai-strict-tool-null-wire' | '/api/openai-usage-details' @@ -2044,6 +2055,7 @@ export interface FileRouteTypes { | '/api/openai-completed-response-text' | '/api/openai-image-2-5-models' | '/api/openai-malformed-tool-arguments' + | '/api/openai-responses-usage' | '/api/openai-shell-skills-wire' | '/api/openai-strict-tool-null-wire' | '/api/openai-usage-details' @@ -2216,6 +2228,7 @@ export interface RootRouteChildren { ApiOpenaiCompletedResponseTextRoute: typeof ApiOpenaiCompletedResponseTextRoute ApiOpenaiImage25ModelsRoute: typeof ApiOpenaiImage25ModelsRoute ApiOpenaiMalformedToolArgumentsRoute: typeof ApiOpenaiMalformedToolArgumentsRoute + ApiOpenaiResponsesUsageRoute: typeof ApiOpenaiResponsesUsageRoute ApiOpenaiShellSkillsWireRoute: typeof ApiOpenaiShellSkillsWireRoute ApiOpenaiStrictToolNullWireRoute: typeof ApiOpenaiStrictToolNullWireRoute ApiOpenaiUsageDetailsRoute: typeof ApiOpenaiUsageDetailsRoute @@ -3099,6 +3112,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiOpenaiMalformedToolArgumentsRouteImport parentRoute: typeof rootRouteImport } + '/api/openai-responses-usage': { + id: '/api/openai-responses-usage' + path: '/api/openai-responses-usage' + fullPath: '/api/openai-responses-usage' + preLoaderRoute: typeof ApiOpenaiResponsesUsageRouteImport + parentRoute: typeof rootRouteImport + } '/api/openai-shell-skills-wire': { id: '/api/openai-shell-skills-wire' path: '/api/openai-shell-skills-wire' @@ -3634,6 +3654,7 @@ const rootRouteChildren: RootRouteChildren = { ApiOpenaiCompletedResponseTextRoute: ApiOpenaiCompletedResponseTextRoute, ApiOpenaiImage25ModelsRoute: ApiOpenaiImage25ModelsRoute, ApiOpenaiMalformedToolArgumentsRoute: ApiOpenaiMalformedToolArgumentsRoute, + ApiOpenaiResponsesUsageRoute: ApiOpenaiResponsesUsageRoute, ApiOpenaiShellSkillsWireRoute: ApiOpenaiShellSkillsWireRoute, ApiOpenaiStrictToolNullWireRoute: ApiOpenaiStrictToolNullWireRoute, ApiOpenaiUsageDetailsRoute: ApiOpenaiUsageDetailsRoute, diff --git a/testing/e2e/src/routes/api.openai-responses-usage.ts b/testing/e2e/src/routes/api.openai-responses-usage.ts new file mode 100644 index 0000000000..c2a2d6c82b --- /dev/null +++ b/testing/e2e/src/routes/api.openai-responses-usage.ts @@ -0,0 +1,49 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, createChatOptions } from '@tanstack/ai' +import { createOpenaiChat } from '@tanstack/ai-openai' + +const LLMOCK_DEFAULT_BASE = process.env.LLMOCK_URL || 'http://127.0.0.1:4010' +const DUMMY_KEY = 'sk-e2e-test-dummy-key' + +/** + * Drives the OpenAI Responses adapter against a hand-crafted aimock mount + * (`/openai-responses-usage`) whose `response.completed` usage reports + * `input_tokens_details.cache_write_tokens`. The companion spec asserts it + * reaches `RUN_FINISHED.usage` as `promptTokensDetails.cacheWriteTokens`. + */ +export const Route = createFileRoute('/api/openai-responses-usage')({ + server: { + handlers: { + POST: async () => { + const adapter = createOpenaiChat('gpt-4o', DUMMY_KEY, { + baseURL: `${LLMOCK_DEFAULT_BASE}/openai-responses-usage/v1`, + }) + + let usage: Record | undefined + try { + for await (const chunk of chat({ + ...createChatOptions({ adapter }), + messages: [{ role: 'user', content: 'hi' }], + })) { + if (chunk.type === 'RUN_FINISHED') { + usage = chunk.usage as Record | undefined + } + } + } catch (error) { + return new Response( + JSON.stringify({ + ok: false, + error: error instanceof Error ? error.message : String(error), + }), + { status: 200, headers: { 'Content-Type': 'application/json' } }, + ) + } + + return new Response(JSON.stringify({ ok: true, usage }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }) + }, + }, + }, +}) diff --git a/testing/e2e/tests/openai-responses-usage.spec.ts b/testing/e2e/tests/openai-responses-usage.spec.ts new file mode 100644 index 0000000000..0059ef27e7 --- /dev/null +++ b/testing/e2e/tests/openai-responses-usage.spec.ts @@ -0,0 +1,29 @@ +import { test, expect } from './fixtures' + +/** + * Verifies that a Responses API cache write reaches `RUN_FINISHED.usage` + * (#1734). The `/api/openai-responses-usage` route drives the OpenAI Responses + * adapter against an aimock mount whose usage reports + * `input_tokens_details.cache_write_tokens`. + */ +test.describe('openai — Responses usage breakdown', () => { + test('cache write tokens reach RUN_FINISHED.usage', async ({ request }) => { + const res = await request.post('/api/openai-responses-usage') + expect(res.ok()).toBe(true) + + const { ok, usage, error } = (await res.json()) as { + ok: boolean + error?: string + usage?: Record + } + + expect(error ?? null).toBeNull() + expect(ok).toBe(true) + expect(usage).toEqual({ + promptTokens: 1941, + completionTokens: 5, + totalTokens: 1946, + promptTokensDetails: { cacheWriteTokens: 1938 }, + }) + }) +}) From 6a9f0b1e3d4ec3f34d319081d9fbde075847eca3 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:47:49 +1100 Subject: [PATCH 07/22] fix(ai-opencode, ai-sandbox): end the run when a teardown step does not settle (#1773) Bound each OpenCode teardown step to one second and log failures, and close open connections when the host tool bridge closes. Closes #1638 --- .changeset/opencode-teardown-bounded.md | 6 ++ packages/ai-opencode/src/adapters/text.ts | 24 +++++++- packages/ai-opencode/tests/teardown.test.ts | 58 +++++++++++++++++++ packages/ai-sandbox/src/tool-bridge.ts | 11 +++- packages/ai-sandbox/tests/tool-bridge.test.ts | 14 +++++ 5 files changed, 108 insertions(+), 5 deletions(-) create mode 100644 .changeset/opencode-teardown-bounded.md create mode 100644 packages/ai-opencode/tests/teardown.test.ts diff --git a/.changeset/opencode-teardown-bounded.md b/.changeset/opencode-teardown-bounded.md new file mode 100644 index 0000000000..9f4b78fda9 --- /dev/null +++ b/.changeset/opencode-teardown-bounded.md @@ -0,0 +1,6 @@ +--- +'@tanstack/ai-opencode': patch +'@tanstack/ai-sandbox': patch +--- + +`RUN_FINISHED` now arrives when a teardown step does not settle. The OpenCode adapter gives each teardown step (session dispose, server dispose, tool bridge close) one second and logs a step that fails. The host tool bridge's `close()` now closes open connections, so an agent that holds a request open on the bridge no longer blocks the end of a Claude Code, Codex, or OpenCode run. diff --git a/packages/ai-opencode/src/adapters/text.ts b/packages/ai-opencode/src/adapters/text.ts index 0d3636fdcc..bb0d39cb12 100644 --- a/packages/ai-opencode/src/adapters/text.ts +++ b/packages/ai-opencode/src/adapters/text.ts @@ -54,6 +54,7 @@ import type { OpencodeTextProviderOptions } from '../provider-options' const DEFAULT_WORKDIR = '/workspace' const DEFAULT_PORT = 4096 +const TEARDOWN_STEP_MS = 1000 export interface OpencodeTextConfig { /** Working directory inside the sandbox. Defaults to `/workspace`. */ @@ -544,9 +545,26 @@ export class OpencodeTextAdapter< externalSignal.removeEventListener('abort', onAbort) } channel.close() - await handle?.dispose() - await server?.dispose() - await bridge?.close() + // Each step gets a time limit: a step that never settles (a sandbox + // `kill()` that never returns, a bridge request held open) must not + // hold back RUN_FINISHED. + const bounded = (step: () => Promise | undefined) => + Promise.race([ + Promise.resolve() + .then(step) + .catch((error: unknown) => + logger.errors('opencode teardown failed', { + error, + source: 'opencode.chatStream', + }), + ), + new Promise((resolve) => + setTimeout(resolve, TEARDOWN_STEP_MS).unref?.(), + ), + ]) + await bounded(() => handle?.dispose()) + await bounded(() => server?.dispose()) + await bounded(() => bridge?.close()) } } diff --git a/packages/ai-opencode/tests/teardown.test.ts b/packages/ai-opencode/tests/teardown.test.ts new file mode 100644 index 0000000000..546f47db67 --- /dev/null +++ b/packages/ai-opencode/tests/teardown.test.ts @@ -0,0 +1,58 @@ +/** + * A teardown step that never settles must not hold back the end of the run + * (issue #1638). The server's `dispose()` and the bridge's `close()` both + * hang here; the stream must still end. + */ +import { describe, expect, it, vi } from 'vitest' +import { + provideSandbox, + provideToolBridgeProvisioner, +} from '@tanstack/ai-sandbox' +import { resolveDebugOption } from '@tanstack/ai/adapter-internals' +import { opencodeText } from '../src/index' +import type { CapabilityContext } from '@tanstack/ai' +import type { SandboxHandle } from '@tanstack/ai-sandbox' + +const never = () => new Promise(() => {}) + +vi.mock('../src/process/sandbox-server', () => ({ + startOpencodeServerInSandbox: () => + Promise.resolve({ baseUrl: 'http://127.0.0.1:1', dispose: never }), +})) +vi.mock('../src/process/server', () => ({ + startOpencodeSession: () => Promise.reject(new Error('no session')), +})) + +const MODEL = 'anthropic/claude-sonnet-4-5' + +describe('opencode teardown', () => { + it('ends the stream when dispose() and close() never settle', async () => { + const capabilities = { + capabilities: { markProvided: () => {}, has: () => true }, + } as unknown as CapabilityContext + provideSandbox(capabilities, { + id: 'sbx', + provider: 'mock', + } as unknown as SandboxHandle) + provideToolBridgeProvisioner(capabilities, { + provision: () => + Promise.resolve({ + name: 'tanstack', + url: 'http://127.0.0.1:1/mcp', + token: 'token', + close: never, + }), + }) + const types: Array = [] + for await (const chunk of opencodeText(MODEL).chatStream({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + logger: resolveDebugOption(false), + capabilities, + tools: [{ name: 'getWeather', description: 'weather' }], + })) { + types.push(chunk.type) + } + expect(types).toContain('RUN_ERROR') + }) +}) diff --git a/packages/ai-sandbox/src/tool-bridge.ts b/packages/ai-sandbox/src/tool-bridge.ts index 5d7a117a31..6d370716ca 100644 --- a/packages/ai-sandbox/src/tool-bridge.ts +++ b/packages/ai-sandbox/src/tool-bridge.ts @@ -363,8 +363,15 @@ export async function startHostToolBridge( name: BRIDGED_MCP_SERVER_NAME, url, token, - close: () => - new Promise((resolve) => httpServer.close(() => resolve())), + close: () => { + const closed = new Promise((resolve) => + httpServer.close(() => resolve()), + ) + // `close()` alone waits for every open connection. An agent that keeps + // a request or a keep-alive socket open would hold the run's teardown. + httpServer.closeAllConnections() + return closed + }, } } diff --git a/packages/ai-sandbox/tests/tool-bridge.test.ts b/packages/ai-sandbox/tests/tool-bridge.test.ts index da8d9dab65..827691812d 100644 --- a/packages/ai-sandbox/tests/tool-bridge.test.ts +++ b/packages/ai-sandbox/tests/tool-bridge.test.ts @@ -3,6 +3,7 @@ * client (no `claude` needed): list tools + call a tool, asserting the host * `execute()` runs and the result comes back, plus bearer-token enforcement. */ +import { request } from 'node:http' import { afterEach, describe, expect, it } from 'vitest' import { Client } from '@modelcontextprotocol/sdk/client/index.js' import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' @@ -96,6 +97,19 @@ describe('startHostToolBridge', () => { await expect(client.connect(transport)).rejects.toThrow() await client.close().catch(() => {}) }) + + it('closes while a request is still open', async () => { + const open = await startHostToolBridge([], { hostForSandbox: '127.0.0.1' }) + // Send headers and part of a body, then never finish the request. + const req = request(open.url, { + method: 'POST', + headers: { authorization: `Bearer ${open.token}` }, + }) + req.on('error', () => {}) + req.write('{') + await new Promise((resolve) => setTimeout(resolve, 50)) + await open.close() + }) }) describe('hostForSandbox', () => { From 3170517d1c3faa0c1ee6fb004b8a551978fa2a7b Mon Sep 17 00:00:00 2001 From: Alem Tuzlak Date: Wed, 14 Oct 2026 13:49:48 +0200 Subject: [PATCH 08/22] feat(ai): check tool input after approval edits and middleware, just before it runs (#1750) * feat(ai): check tool input at the last moment, after approval edits and middleware * test(ai): fix the boundary client tool tests * test(ai): the subagent tool input is checked after middleware * test(ai): subagent schema errors also show the received arguments * fix(ai): keep returned client results and run tool hooks once at boundaries - A client result that comes back is kept as is. Checking the raw model input again dropped the real result when middleware had repaired it. - An afterTools boundary reuses the tool phase result, so onBeforeToolCall and the tool results do not run twice. - The boundary drains executeToolCalls through drainToolCallGenerator and passes the resume state; the hooks and resume state come from shared helpers. - The validation error shows the model's own arguments in the tool result only, so middleware-injected values do not reach the model or tracing. - Changeset, middleware docs, and args JSDoc corrected. * feat(ai): check JSON Schema and MCP tool input with @cfworker/json-schema (#1751) * feat(ai): validate JSON Schema tool input with @cfworker/json-schema * chore: keep one @google/genai copy for ai-vertex and ai-gemini * test(ai): the JSON Schema of a subagent tool checks background * fix(ai): keep the page URL out of a JSON Schema tool error A failed $ref included the validator base URI. In a browser that URI is the page, query included. * fix(ai): JSON Schema tool check accepts OpenAPI 3.0 and keeps data intact - Strip `format` only from schema positions, so a `format` key inside `const` or `enum` data still matches. - Convert boolean `exclusiveMinimum`/`exclusiveMaximum` per schema, so they work under any `$schema` and numeric bounds elsewhere still apply. - Support OpenAPI 3.0 `nullable: true`: null joins `type`/`enum`, and a nullable `$ref`/`allOf`/`anyOf`/`oneOf`/`not` schema is wrapped as `if: null, else: schema`. Local `$ref`s into a wrapped schema move with it. - Check an input with no JSON form as null, list at most 10 errors, and drop more wrapper errors (`if`, `$recursiveRef`, `additionalItems`, `unevaluatedItems`). - Read the first error line with `split('\n', 1)`, which type-checks under `noUncheckedIndexedAccess`. * feat(ai): repair common model mistakes in tool input, with chat({ repairToolInput }) (#1752) * feat(ai): repair common model mistakes in tool input before the final check * fix(ai): type the repair order list * test(ai): use an input the repair cannot fix in the subagent schema test * feat(ai): chat({ repairToolInput }) to turn off or replace the input repair false turns the built-in tool input repair off. A function replaces it: it gets { tool, toolCall, input, error } for an input that failed the check, and returns the input to check again, or undefined to send the validation error to the model. * feat(ai, adapters): keep the raw tool argument string for replay (#1753) * feat(ai, adapters): keep the raw tool argument string for replay * refactor(ai): one tool-argument reconcile helper, args sent once finishToolCallArguments replaces the three reconcile copies in ToolCallManager, TextEngine and StreamProcessor. TOOL_CALL_END now sends its arguments once, as metadata.tanstack.args reconciled to the input; in-process input is parsed back from it. Drops jsonInput and isToolInputJsonLossless. A null input no longer replaces the raw args. * fix(ai-bedrock, ai-gemini): replay malformed tool arguments as {} Malformed arguments now stay in the history next to their tool error. The Bedrock and Gemini Interactions history converters send them as {} instead of failing the request. * fix(ai-gemini): send the args snapshot only when no deltas stream The initial step.arguments snapshot went out as a TOOL_CALL_ARGS delta before the streamed deltas, so the deltas did not add up to the final arguments on the client. * refactor(adapters): always send args on TOOL_CALL_END Drop the redundant rawArgs guards and hasArguments flags. Core treats empty args the same as missing args. * test(ai): keep the JSON Schema check tests off the repair path The OpenAPI 3.0 and $ref tests from feat/json-schema-tool-input send inputs the repair now fixes (a null optional field, "3" for an integer). They check the schema alone, with repairToolInput: false. Adds a test that the repair walks a nullable schema, and restores the subagent 'yes' input. --------- Co-authored-by: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> * fix(ai): give held calls a result and keep child interrupts at boundaries - When the only call that waits for approval fails its input check, the calls that the batch gating held back now run, so each call has a result. - An approved client tool is sent to the client while another call waits for approval, as a plain client tool already was. - The validation error tells the model when an approval edit or middleware changed its arguments, so it does not retry the same valid call. - An afterTools boundary interrupt keeps the subagent interrupts of the tool phase. - Changeset and server tools docs: a bad approval edit is rejected when the run resumes. --------- Co-authored-by: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> --- .changeset/json-schema-tool-input.md | 17 + .changeset/raw-tool-arguments.md | 19 + .changeset/tool-input-coercion.md | 12 + .changeset/tool-input-final-check.md | 22 + docs/advanced/middleware.md | 21 +- docs/chat/stream-events.md | 2 +- docs/config.json | 4 +- docs/tools/server-tools.md | 105 +- packages/ai-anthropic/src/adapters/text.ts | 27 +- .../tests/anthropic-adapter.test.ts | 93 + .../src/converse/message-converter.ts | 17 +- .../src/converse/stream-processor.ts | 5 + .../tests/converse/message-converter.test.ts | 35 +- .../tests/converse/stream-processor.test.ts | 62 + packages/ai-gemini/src/adapters/text.ts | 16 +- .../experimental/text-interactions/adapter.ts | 92 +- .../ai-gemini/tests/gemini-adapter.test.ts | 77 + .../ai-gemini/tests/replay-parity.test.ts | 39 + .../tests/text-interactions-adapter.test.ts | 103 +- packages/ai-mistral/src/adapters/text.ts | 18 +- .../ai-mistral/tests/mistral-adapter.test.ts | 103 ++ packages/ai-ollama/src/adapters/text.ts | 8 +- packages/ai-ollama/tests/text-adapter.test.ts | 47 + .../src/adapters/responses-text.ts | 15 +- packages/ai-openrouter/src/adapters/text.ts | 10 +- .../tests/openrouter-adapter.test.ts | 48 + .../openrouter-responses-adapter.test.ts | 102 ++ packages/ai/package.json | 1 + .../ai/src/activities/chat/agents/spawn.ts | 4 +- packages/ai/src/activities/chat/index.ts | 223 +-- .../src/activities/chat/middleware/types.ts | 9 +- .../src/activities/chat/stream/processor.ts | 84 +- .../activities/chat/tools/input-validation.ts | 557 ++++++ .../src/activities/chat/tools/tool-calls.ts | 360 ++-- packages/ai/src/middlewares/otel.ts | 3 + packages/ai/src/types.ts | 26 + .../utilities/normalize-stream-chunk.test.ts | 31 + .../src/utilities/normalize-stream-chunk.ts | 16 +- .../utilities/restore-inbound-chunk.test.ts | 15 + .../ai/src/utilities/restore-inbound-chunk.ts | 8 +- .../ai/src/utilities/tool-call-arguments.ts | 244 +++ .../ai/tests/chat-client-tool-wait.test.ts | 6 +- packages/ai/tests/subagent-interrupts.test.ts | 46 + packages/ai/tests/subagent-tool-input.test.ts | 44 +- packages/ai/tests/tool-call-arguments.test.ts | 223 +++ .../ai/tests/tool-input-validation.test.ts | 1527 +++++++++++++++++ .../src/adapters/chat-completions-text.ts | 7 +- .../src/adapters/responses-text.ts | 15 +- .../tests/chat-completions-text.test.ts | 61 +- .../openai-base/tests/responses-text.test.ts | 102 ++ pnpm-lock.yaml | 64 +- .../middleware-test/raw-tool-args.json | 27 + testing/e2e/src/routeTree.gen.ts | 21 + .../src/routes/api.json-schema-tool-input.ts | 87 + testing/e2e/src/routes/api.middleware-test.ts | 14 +- testing/e2e/src/routes/api.tools-test.ts | 5 +- testing/e2e/src/routes/middleware-test.tsx | 2 + .../e2e/tests/json-schema-tool-input.spec.ts | 31 + testing/e2e/tests/middleware.spec.ts | 92 + 59 files changed, 4554 insertions(+), 520 deletions(-) create mode 100644 .changeset/json-schema-tool-input.md create mode 100644 .changeset/raw-tool-arguments.md create mode 100644 .changeset/tool-input-coercion.md create mode 100644 .changeset/tool-input-final-check.md create mode 100644 packages/ai/src/activities/chat/tools/input-validation.ts create mode 100644 packages/ai/src/utilities/tool-call-arguments.ts create mode 100644 packages/ai/tests/tool-call-arguments.test.ts create mode 100644 packages/ai/tests/tool-input-validation.test.ts create mode 100644 testing/e2e/fixtures/middleware-test/raw-tool-args.json create mode 100644 testing/e2e/src/routes/api.json-schema-tool-input.ts create mode 100644 testing/e2e/tests/json-schema-tool-input.spec.ts diff --git a/.changeset/json-schema-tool-input.md b/.changeset/json-schema-tool-input.md new file mode 100644 index 0000000000..7b930dbbca --- /dev/null +++ b/.changeset/json-schema-tool-input.md @@ -0,0 +1,17 @@ +--- +'@tanstack/ai': minor +--- + +Check a raw JSON Schema tool input before the tool runs. This covers plain JSON Schema tools and every MCP tool. + +Behavior change: + +- A tool with a JSON Schema `inputSchema` now gets a tool error (state `output-error`) when the input does not fit the schema. The tool does not run. Before, any input ran. +- The check runs at the same place as the Standard Schema check: after approval edits and middleware `onBeforeToolCall`. +- The draft comes from `$schema` (draft-04, draft-07, 2019-09, 2020-12). With no `$schema`, the check uses 2020-12. `format` is not checked. +- OpenAPI 3.0 schemas work: `nullable: true` lets the field be `null`, and a boolean `exclusiveMinimum` or `exclusiveMaximum` works in every draft. +- The error names each failed location (up to 10) and shows the received arguments: `Input validation failed for tool : #/count: ...`. +- A schema that the check cannot use, such as one with a remote `$ref`, gives a tool error on each call. +- The check uses `@cfworker/json-schema`. It does not compile code, so it runs on Cloudflare Workers. + +Migration: if a JSON Schema tool must accept input that its schema rejects, widen the schema. diff --git a/.changeset/raw-tool-arguments.md b/.changeset/raw-tool-arguments.md new file mode 100644 index 0000000000..9f97815a79 --- /dev/null +++ b/.changeset/raw-tool-arguments.md @@ -0,0 +1,19 @@ +--- +'@tanstack/ai': minor +'@tanstack/openai-base': patch +'@tanstack/ai-openrouter': patch +'@tanstack/ai-bedrock': patch +'@tanstack/ai-mistral': patch +'@tanstack/ai-gemini': patch +'@tanstack/ai-ollama': patch +'@tanstack/ai-anthropic': patch +--- + +Keep the provider's raw tool argument string for replay. + +- Adapters now send the raw argument string as `args` on `TOOL_CALL_END`. +- The tool call in the history keeps that exact string when it is the same JSON value as the parsed input, so the next request sends the provider's own bytes back (for example `1.50`, or an id past `Number.MAX_SAFE_INTEGER`). When the parsed input differs, only the changed values are written again. +- The `TOOL_CALL_END` that reaches the client carries the same string in `metadata.tanstack.args`, so the client history matches the server history. It no longer also carries `metadata.tanstack.input`: in-process `chat()` and `StreamProcessor` parse `input` from `args`. Run the client and the server on the same release: an older client ignores `args`, so it builds the tool input from the streamed deltas only, and a tool call with no deltas (for example an Anthropic server tool) gets no input. +- History converters for Bedrock and the Gemini Interactions API send malformed tool arguments as `{}` instead of failing the request. +- Gemini Interactions: the initial `step.arguments` snapshot goes out as a `TOOL_CALL_ARGS` delta only when no argument deltas follow, so the deltas always add up to the final arguments. +- Malformed tool arguments are no longer turned into `{}` by an adapter. The tool gets a "Failed to parse tool arguments as JSON" error that the model can repair. Mistral still ends the run with an error on malformed arguments, as before. diff --git a/.changeset/tool-input-coercion.md b/.changeset/tool-input-coercion.md new file mode 100644 index 0000000000..3e85c9f7c5 --- /dev/null +++ b/.changeset/tool-input-coercion.md @@ -0,0 +1,12 @@ +--- +'@tanstack/ai': minor +--- + +Repair common model mistakes in a tool input before the final input check. + +- When a JSON Schema tool's input (and every MCP tool's input) fails the check, it is repaired and checked again: `"5"` as `5` (decimal text only), `"true"` as `true`, `5` as `"5"` for a string, and `null` as `""`, `0`, or `false` for a required field. A `null` for an optional field that does not allow `null` is removed. The repair follows unions, `$ref`, `if`/`then`/`else`, `prefixItems`, `patternProperties`, and `additionalProperties`. +- A Standard Schema runs first as written. When it fails and the schema can export JSON Schema (Zod, ArkType, Valibot), the input is repaired and the schema runs again. A failure reports the issues of the first run. +- The received arguments do not change. The tool gets the repaired copy. +- New `chat({ repairToolInput })` option. `false` turns the repair off. A function replaces the built-in repair: it gets `{ tool, toolCall, input, error }` and returns the input to check again, or `undefined` to send the validation error to the model. + +Behavior change: a call that failed the check only because of these mistakes now runs with the repaired input. diff --git a/.changeset/tool-input-final-check.md b/.changeset/tool-input-final-check.md new file mode 100644 index 0000000000..39fb39220b --- /dev/null +++ b/.changeset/tool-input-final-check.md @@ -0,0 +1,22 @@ +--- +'@tanstack/ai': minor +--- + +Check a tool input at the last moment, just before the tool runs or the client gets the call. + +Behavior change: + +- The tool's Standard Schema `inputSchema` now checks the input after approval `editedArgs` and after middleware `onBeforeToolCall`. Before, a `transformArgs` value reached `execute` or the client without a check. An input that does not fit the schema now becomes a tool error (state `output-error`), and the tool does not run. An approval edit is also still checked when the run resumes. +- `onBeforeToolCall` gets the parsed input before the schema check. It has no schema defaults or transforms yet, and it can be invalid. +- Async Standard Schema input validators are awaited. Before, every call to such a tool failed with an input validation error. +- `onBeforeToolCall` now also runs for a client tool before the server sends the call to the client. This includes a client request at an interrupt boundary. It does not run while a call waits for approval or for a denied call. When the client result comes back, the result is kept, with no second hook run and no second check. `otelMiddleware` records no `execute_tool` span for a client tool, because the server does not run it. +- A validation error after the hook runs `onAfterToolCall` with `ok: false`. +- The tool error that the model sees now also shows the arguments that the model sent. When an approval edit or middleware changed them, the error says so. +- While a call waits for approval, an approved client tool in the same batch is now sent to the client. Before, it waited for the next round trip. +- When the only call that waits for approval fails its input check, the other calls of the batch now run. Before, they got no result. + +Migration: + +- If a middleware `transformArgs` or an approval edit sends a value that the schema rejects, make the value fit the schema, or widen the schema. +- If `onBeforeToolCall` reads `hookCtx.args`, check its shape first. Do not expect schema defaults there. +- If a middleware counts or blocks calls in `onBeforeToolCall` and must see only server tools, return early when `hookCtx.tool?.execute` is missing. diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index e0800eb5e5..c52d5e90ea 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -668,7 +668,26 @@ system prompt, is in [Apply Answers](../interrupts/apply-answers). ### onBeforeToolCall -Called before each tool executes. The first middleware that returns a non-void decision short-circuits — remaining middleware are skipped for that tool call. +Called before a server tool runs, and before a client tool goes to the client. The first middleware that returns a non-void decision skips the remaining middleware for that call. + +Each tool call goes through these steps: + +1. The engine parses the JSON arguments from the model. +2. `onBeforeToolCall` gets the parsed input. +3. The tool's `inputSchema` checks the final input. An async validator is awaited. A schema transform applies once, here. +4. The server tool runs, or the client gets the call, with the checked input. + +The schema checks the input after your hook: + +- `hookCtx.args` has no schema defaults or transforms yet, and it can be invalid. Check its shape before you read it. +- A `transformArgs` value or an in-place edit must fit the schema. If it does not fit, the call gets a tool error, `onAfterToolCall` runs with `ok: false`, and the tool does not run. + +When the hook runs: + +- A client tool runs the hook before the server sends the call to the client. The hook does not run when the client result comes back. +- A client tool gets `onAfterToolCall` only when the hook skips the call or the schema check fails. The result from the client does not run it. If your middleware pairs the two hooks, return early when `hookCtx.tool?.execute` is missing. +- A call that waits for approval does not run the hook. A denied call does not run it either. +- After approval, the hook gets the approved input, including any `editedArgs`. See [Tool Approval](../interrupts/tool-approval). ```typescript import { type ChatMiddleware } from "@tanstack/ai"; diff --git a/docs/chat/stream-events.md b/docs/chat/stream-events.md index a88110295d..0c6fd3ff54 100644 --- a/docs/chat/stream-events.md +++ b/docs/chat/stream-events.md @@ -238,7 +238,7 @@ The media generation hooks take a `threadId` too. There it names a slot, not a c ## Tool input and output -SSE and HTTP `TOOL_CALL_END` does not carry parsed `input`. In-process `chat()` still has `input`. Tool input and output also live on `UIMessage` parts. +SSE and HTTP `TOOL_CALL_END` does not carry parsed `input`. `metadata.tanstack.args` carries the tool argument string as the provider sent it. In-process `chat()` still has `input`, parsed from that string. Tool input and output also live on `UIMessage` parts. On the server, feed chunks into `StreamProcessor`. On the client, read `useChat` `messages`. diff --git a/docs/config.json b/docs/config.json index 98cc0f1dbd..25cd364bcd 100644 --- a/docs/config.json +++ b/docs/config.json @@ -294,7 +294,7 @@ "label": "Server Tools", "to": "tools/server-tools", "addedAt": "2026-04-15", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-14" }, { "label": "Client Tools", @@ -353,7 +353,7 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-10-13" + "updatedAt": "2026-10-14" }, { "label": "Built-in Middleware", diff --git a/docs/tools/server-tools.md b/docs/tools/server-tools.md index 12aaf118a7..9106cb0e00 100644 --- a/docs/tools/server-tools.md +++ b/docs/tools/server-tools.md @@ -355,7 +355,96 @@ const getUserData = getUserDataDef.server(async ({ userId }) => { **Throwing vs. returning an error:** if your `.server()` function throws, the SDK catches it and surfaces it as a tool-result *error* (the model sees the failure but you lose control over the message). Returning a structured `{ error }` shape keeps the model in control of how to recover and is usually preferable. Either way, when an `outputSchema` is defined the returned value is validated against it (Zod) before being added to the conversation — so include the `error` field in your `outputSchema` if you return it. -Malformed JSON arguments and Standard Schema input validation failures are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. +Malformed JSON arguments and input validation failures (Standard Schema or JSON Schema) are also surfaced as tool-result errors. The tool implementation is not called, and the agent loop can return the error to the model so it can repair the tool call. This also applies to Chat Completions tools with no required input fields. + +The `inputSchema` checks the input just before the tool runs: + +- The check runs after approval edits and after middleware `onBeforeToolCall`. An approval edit that does not fit the schema stops the resume with an `invalid-edited-args` error. A middleware change that does not fit becomes a tool error. +- An async validator is awaited. +- The tool gets the schema's output, so a transform such as `z.string().transform(Number)` applies once. + +Models often send a value with the wrong type, such as `"5"` for a number. The check repairs these against the schema before it fails the call: + +| The model sends | The schema wants | The tool gets | +| --- | --- | --- | +| `"5"` | `number` or `integer` | `5` | +| `"true"` | `boolean` | `true` | +| `5` | `string` | `"5"` | +| `null` for an optional field | a type with no `null` | no field | +| `null` for a required field | `string`, `number`, or `boolean` | `""`, `0`, or `false` | + +- A JSON Schema tool, and an MCP tool, gets the repair when its input fails the check. A valid input runs as it is. +- A Standard Schema that can export JSON Schema (Zod, ArkType, Valibot) runs first as written. If it fails, the input is repaired and the schema runs again. If it still fails, the error shows the issues of the first run. +- An input that cannot be repaired, such as `"many"` for a number, becomes a tool error. + +### Turn off or replace the repair + +Sometimes the built-in repair is the wrong move. Maybe your tool moves money, and a `null` must never turn into `0`. Or maybe your model makes the same odd mistake every time, and you know the fix. Pass `repairToolInput` to `chat()` and take control. + +Turn the repair off, so a bad input always goes back to the model as an error: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { transferFunds } from "./tools"; + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-6.1-sol"), + messages, + tools: [transferFunds], + repairToolInput: false, + }); + + return toServerSentEventsResponse(stream); +} +``` + +Or write your own fix. Say the model keeps sending `"limit": "all"` to `search_products`. Map it to a number: + +```typescript +import { chat, toServerSentEventsResponse } from "@tanstack/ai"; +import { openaiText } from "@tanstack/ai-openai"; +import { searchProducts } from "./tools"; + +export async function POST(request: Request) { + const { messages } = await request.json(); + + const stream = chat({ + adapter: openaiText("gpt-6.1-sol"), + messages, + tools: [searchProducts], + repairToolInput: ({ tool, input }) => { + if ( + tool.name === "search_products" && + typeof input === "object" && + input !== null && + "limit" in input && + input.limit === "all" + ) { + return { ...input, limit: 100 }; + } + return undefined; + }, + }); + + return toServerSentEventsResponse(stream); +} +``` + +How your function works: + +- It runs only when an input fails the check. A valid input never gets to it. +- It replaces the built-in repair. A `"5"` stays a string unless your function changes it. +- It gets `tool`, `toolCall`, `input`, and `error` (the validation error). +- Return the new input, and the check runs again. If that check fails, the model gets the original error. +- Return `undefined` to send the original error to the model. If your function throws, the model gets the thrown message. + +The option works the same for Standard Schema tools, JSON Schema tools, and MCP tools. + +A client tool, or a tool that waits for approval, is checked twice: once before it goes out, and once when it comes back. So your function can run twice for one call. ## Using JSON Schema @@ -403,14 +492,24 @@ const getUserData = getUserDataDef.server(async (args) => { }); ``` -> **Note:** JSON Schema tools skip runtime validation. Zod schemas are recommended for full type safety and validation. +The JSON Schema checks each call before the tool runs. MCP tools get the same check. If the input does not fit, the call gets a tool error, and the error names each failed field: + +```text +Input validation failed for tool get_user_data: #: Instance does not have required property "userId". +``` + +- The check uses the draft that `$schema` names: draft-04, draft-07, 2019-09, or 2020-12. With no `$schema`, it uses 2020-12. +- Got an OpenAPI 3.0 schema? `nullable: true` works, so that field can be `null`. A boolean `exclusiveMinimum` or `exclusiveMaximum` works too. +- The check does not test `format` (such as `uri` or `email`). A `format` only describes the value. +- The check does not compile code, so it also runs on Cloudflare Workers. +- `args` is still typed as `unknown`. A Zod schema gives you the type too. > **Tip:** When you pass `.client()` tools to `useChat`, a check on `part.name` narrows `part.input` and `part.output`. See [Type-safe tool call events](../chat/stream-events#type-safe-tool-call-events). ## Best Practices 1. **Keep tools focused** - Each tool should do one thing well -2. **Validate inputs** - Use Zod schemas to ensure type safety (JSON Schema skips validation) +2. **Validate inputs** - The `inputSchema` checks every call. A Zod schema also gives you types 3. **Handle errors** - Return meaningful error messages 4. **Use descriptions** - Clear descriptions help the model use tools correctly 5. **Secure sensitive operations** - Never expose API keys or secrets to the client diff --git a/packages/ai-anthropic/src/adapters/text.ts b/packages/ai-anthropic/src/adapters/text.ts index d850bbe1f3..d02313231b 100644 --- a/packages/ai-anthropic/src/adapters/text.ts +++ b/packages/ai-anthropic/src/adapters/text.ts @@ -1512,7 +1512,13 @@ export class AnthropicTextAdapter< let accumulatedSignature = '' const toolCallsMap = new Map< number, - { id: string; name: string; input: string; started: boolean } + { + id: string + name: string + input: string + started: boolean + hasInputDelta: boolean + } >() let currentToolIndex = -1 // Server-side tools share the `input_json_delta` wire format with client @@ -1577,8 +1583,11 @@ export class AnthropicTextAdapter< toolCallsMap.set(currentToolIndex, { id: event.content_block.id, name: event.content_block.name, - input: '', + // The initial input snapshot. The first `input_json_delta` + // replaces it. + input: JSON.stringify(event.content_block.input) ?? '', started: false, + hasInputDelta: false, }) } else if (event.content_block.type === 'server_tool_use') { currentServerTool = { @@ -1851,6 +1860,10 @@ export class AnthropicTextAdapter< } } + if (!existing.hasInputDelta) { + existing.input = '' + existing.hasInputDelta = true + } existing.input += event.delta.partial_json yield { @@ -1932,13 +1945,14 @@ export class AnthropicTextAdapter< } } - // Emit TOOL_CALL_END - let parsedInput: unknown = {} + // Emit TOOL_CALL_END. Malformed JSON sends no `input`, so core + // reports the error from the raw `args`. + let parsedInput: unknown try { const parsed = existing.input ? JSON.parse(existing.input) : {} parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = {} + // Keep `parsedInput` undefined for the tool error. } yield { @@ -1948,7 +1962,8 @@ export class AnthropicTextAdapter< toolName: existing.name, model, timestamp: Date.now(), - input: parsedInput, + args: existing.input, + ...(parsedInput !== undefined && { input: parsedInput }), } // Reset so a new TEXT_MESSAGE_START is emitted if text follows tool calls diff --git a/packages/ai-anthropic/tests/anthropic-adapter.test.ts b/packages/ai-anthropic/tests/anthropic-adapter.test.ts index 86e950aa5a..4c51a2db89 100644 --- a/packages/ai-anthropic/tests/anthropic-adapter.test.ts +++ b/packages/ai-anthropic/tests/anthropic-adapter.test.ts @@ -12,6 +12,7 @@ import { AnthropicTextAdapter } from '../src/adapters/text' import type { AnthropicTextProviderOptions } from '../src/adapters/text' import type { AnthropicDocumentMetadata } from '../src/message-types' import { ANTHROPIC_MAX_NONSTREAMING_TOKENS } from '../src/model-meta' +import { createSilentLogger } from './utils/logger' import { z } from 'zod' const mocks = vi.hoisted(() => { @@ -2520,6 +2521,98 @@ describe('Anthropic stream processing', () => { 'text B', ]) }) + + /** The TOOL_CALL_END of one streamed `tool_use` block. */ + async function toolCallEnd( + initialInput: Record | undefined, + deltas: Array, + ) { + mocks.betaMessagesCreate.mockResolvedValueOnce( + (async function* () { + yield { + type: 'content_block_start', + index: 0, + content_block: { + type: 'tool_use', + id: 'toolu_raw', + name: 'inspect', + ...(initialInput !== undefined && { input: initialInput }), + }, + } + for (const partial_json of deltas) { + yield { + type: 'content_block_delta', + index: 0, + delta: { type: 'input_json_delta', partial_json }, + } + } + yield { type: 'content_block_stop', index: 0 } + yield { + type: 'message_delta', + delta: { stop_reason: 'tool_use' }, + usage: { output_tokens: 1 }, + } + yield { type: 'message_stop' } + })(), + ) + + const chunks: Array = [] + for await (const chunk of createAdapter('claude-opus-4-1').chatStream({ + model: 'claude-opus-4-1', + messages: [{ role: 'user', content: 'Go' }], + logger: createSilentLogger(), + })) { + chunks.push(chunk) + } + return chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + { + name: 'the streamed input', + initial: {}, + deltas: ['{"count":2}'], + args: '{"count":2}', + input: { count: 2 }, + }, + { + name: 'the streamed bytes as sent', + initial: {}, + deltas: ['{"n": ', '1.50}'], + args: '{"n": 1.50}', + input: { n: 1.5 }, + }, + { + name: 'the initial input snapshot when no delta arrives', + initial: { value: 'start' }, + deltas: [], + args: '{"value":"start"}', + input: { value: 'start' }, + }, + { + name: 'the streamed input over the initial snapshot', + initial: { value: 'start' }, + deltas: ['{"count":2}'], + args: '{"count":2}', + input: { count: 2 }, + }, + ])( + 'sends $name as the raw TOOL_CALL_END args', + async ({ initial, deltas, args, input }) => { + expect(await toolCallEnd(initial, deltas)).toMatchObject({ args, input }) + }, + ) + + it('sends malformed tool input raw on TOOL_CALL_END with no input', async () => { + const end = await toolCallEnd({}, ['{"a":']) + expect(end).toMatchObject({ args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + + it('sends empty input as {} with empty raw args', async () => { + const end = await toolCallEnd(undefined, []) + expect(end).toMatchObject({ input: {}, args: '' }) + }) }) describe('Anthropic adapter error handling', () => { diff --git a/packages/ai-bedrock/src/converse/message-converter.ts b/packages/ai-bedrock/src/converse/message-converter.ts index 602a0adf15..7900075848 100644 --- a/packages/ai-bedrock/src/converse/message-converter.ts +++ b/packages/ai-bedrock/src/converse/message-converter.ts @@ -276,11 +276,10 @@ function messageToBlocks( } // null → no text blocks - // Append toolUse blocks for assistant tool calls. Malformed or non-object - // arguments come from a prior assistant turn the engine already accepted, so - // they signal a real upstream problem — throw rather than silently coercing to - // `{}` and forwarding a corrupted tool call (the adapter's catch surfaces it - // as a RUN_ERROR). + // Append toolUse blocks for assistant tool calls. Malformed JSON stays in + // the history next to its tool error, so it goes out as `{}`. Non-object + // arguments signal a real upstream problem and throw (the adapter's catch + // surfaces it as a RUN_ERROR). if (msg.role === 'assistant' && msg.toolCalls) { for (const call of msg.toolCalls) blocks.push(toolUseBlock(call)) } @@ -292,13 +291,11 @@ function toolUseBlock( call: NonNullable[number], ): ContentBlock { const rawArguments = call.function.arguments || '{}' - let parsed: unknown + let parsed: unknown = {} try { parsed = JSON.parse(rawArguments) - } catch (error: unknown) { - throw new Error( - `Bedrock Converse: tool call "${call.function.name}" has malformed JSON arguments (${String(error)}). Raw: ${rawArguments}`, - ) + } catch { + // The tool result already told the model the arguments were malformed. } if (parsed === null || typeof parsed !== 'object' || Array.isArray(parsed)) { throw new Error( diff --git a/packages/ai-bedrock/src/converse/stream-processor.ts b/packages/ai-bedrock/src/converse/stream-processor.ts index 911d90d7e3..2be2caa5cb 100644 --- a/packages/ai-bedrock/src/converse/stream-processor.ts +++ b/packages/ai-bedrock/src/converse/stream-processor.ts @@ -72,6 +72,7 @@ export async function* processConverseStream( id: string name: string started: boolean + arguments: string } >() @@ -167,6 +168,7 @@ export async function* processConverseStream( id, name, started: true, + arguments: '', }) yield { type: EventType.TOOL_CALL_START, @@ -189,6 +191,7 @@ export async function* processConverseStream( if (delta && 'toolUse' in delta && delta.toolUse?.input !== undefined) { const toolCall = toolCallsByIndex.get(index) if (toolCall?.started) { + toolCall.arguments += delta.toolUse.input yield { type: EventType.TOOL_CALL_ARGS, toolCallId: toolCall.id, @@ -295,6 +298,7 @@ export async function* processConverseStream( yield { type: EventType.TOOL_CALL_END, toolCallId: toolCall.id, + args: toolCall.arguments, toolCallName: toolCall.name, toolName: toolCall.name, timestamp: Date.now(), @@ -337,6 +341,7 @@ export async function* processConverseStream( yield { type: EventType.TOOL_CALL_END, toolCallId: toolCall.id, + args: toolCall.arguments, toolCallName: toolCall.name, toolName: toolCall.name, timestamp: Date.now(), diff --git a/packages/ai-bedrock/tests/converse/message-converter.test.ts b/packages/ai-bedrock/tests/converse/message-converter.test.ts index 0001e8b107..95409cfa18 100644 --- a/packages/ai-bedrock/tests/converse/message-converter.test.ts +++ b/packages/ai-bedrock/tests/converse/message-converter.test.ts @@ -117,22 +117,25 @@ describe('toConverseMessages', () => { }) }) - it('throws on malformed JSON tool-call arguments instead of forwarding {}', () => { - expect(() => - toConverseMessages([ - { - role: 'assistant', - content: '', - toolCalls: [ - { - id: 't1', - type: 'function', - function: { name: 'getX', arguments: '{"a":' }, - }, - ], - }, - ]), - ).toThrow(/malformed JSON/) + it('sends malformed JSON tool-call arguments as {}', () => { + const { + messages: [message], + } = toConverseMessages([ + { + role: 'assistant', + content: '', + toolCalls: [ + { + id: 't1', + type: 'function', + function: { name: 'getX', arguments: '{"a":' }, + }, + ], + }, + ]) + expect(message?.content).toEqual([ + { toolUse: { toolUseId: 't1', name: 'getX', input: {} } }, + ]) }) it('throws when tool-call arguments parse to a non-object', () => { diff --git a/packages/ai-bedrock/tests/converse/stream-processor.test.ts b/packages/ai-bedrock/tests/converse/stream-processor.test.ts index 620b304d45..9d546c403c 100644 --- a/packages/ai-bedrock/tests/converse/stream-processor.test.ts +++ b/packages/ai-bedrock/tests/converse/stream-processor.test.ts @@ -377,4 +377,66 @@ describe('processConverseStream', () => { events.filter((e) => e.type === EventType.TOOL_CALL_END), ).toHaveLength(2) }) + + it('puts the raw streamed argument string on TOOL_CALL_END', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '{"n": ' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '1.50}' } }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + { messageStop: { stopReason: 'tool_use' } }, + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '{"n": 1.50}' }) + }) + + it('puts malformed raw arguments on TOOL_CALL_END without an input', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { + contentBlockDelta: { + delta: { toolUse: { input: '{"a":' } }, + contentBlockIndex: 0, + }, + }, + // Truncated: the drained END still carries the raw string. + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + + it('sends empty args on TOOL_CALL_END when no argument fragment arrived', async () => { + const events = await collect( + { + contentBlockStart: { + start: { toolUse: { toolUseId: 't1', name: 'getX' } }, + contentBlockIndex: 0, + }, + }, + { contentBlockStop: { contentBlockIndex: 0 } }, + ) + const end = events.find((e) => e.type === EventType.TOOL_CALL_END) + expect(end).toMatchObject({ toolCallId: 't1', args: '' }) + }) }) diff --git a/packages/ai-gemini/src/adapters/text.ts b/packages/ai-gemini/src/adapters/text.ts index 8286eb1624..16499b1845 100644 --- a/packages/ai-gemini/src/adapters/text.ts +++ b/packages/ai-gemini/src/adapters/text.ts @@ -305,7 +305,12 @@ function buildInteractionsInput( const signature = metadata?.thoughtSignature if (typeof signature === 'string' && signature) steps.push({ type: 'thought', signature, summary: [] }) - const args: unknown = JSON.parse(toolCall.function.arguments || '{}') + // Malformed JSON stays in the history next to its tool error, so it + // goes out as `{}`. + let args: unknown = {} + try { + args = JSON.parse(toolCall.function.arguments || '{}') + } catch {} if (typeof args !== 'object' || args === null || Array.isArray(args)) throw new Error( `Gemini Interactions function ${toolCall.function.name} requires object arguments`, @@ -1158,12 +1163,14 @@ export class GeminiTextAdapter< // this chunk (including UNEXPECTED_TOOL_CALL finishes) were already // registered and started by the per-part loop above. for (const [toolCallId, toolCallData] of toolCallMap.entries()) { - let parsedInput: unknown = {} + // Malformed JSON sends no `input`, so core reports the error from + // the raw `args`. + let parsedInput: unknown try { const parsed = JSON.parse(toolCallData.args) parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = {} + // Keep `parsedInput` undefined for the tool error. } yield { @@ -1173,7 +1180,8 @@ export class GeminiTextAdapter< toolName: toolCallData.name, model, timestamp: Date.now(), - input: parsedInput, + args: toolCallData.args, + ...(parsedInput !== undefined && { input: parsedInput }), } } diff --git a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts index 40dc9a26ce..4da75e0593 100644 --- a/packages/ai-gemini/src/experimental/text-interactions/adapter.ts +++ b/packages/ai-gemini/src/experimental/text-interactions/adapter.ts @@ -920,6 +920,22 @@ function parsePartialToolArguments( } } +// The raw `args` and parsed `input` of a TOOL_CALL_END. The streamed argument +// text wins over the initial `step.arguments` snapshot. Malformed JSON sends +// no `input`, so core reports the error from the raw `args`. +function toolCallEndArguments( + streamed: string | undefined, + state: ToolCallState, +): { args: string; input?: unknown } { + const args = streamed?.trim() ? streamed : JSON.stringify(state.args) + try { + const parsed: unknown = JSON.parse(args) + return { args, input: parsed && typeof parsed === 'object' ? parsed : {} } + } catch { + return { args } + } +} + // `satisfies` pins these arrays to the SDK's narrow mime-type unions: if // Google removes a format the build breaks, and if they add one ours keeps // working (we just won't accept the new one until added here). @@ -1235,6 +1251,37 @@ async function* translateInteractionEvents( } } + // Ends a tool call. With no streamed argument text, the initial + // `step.arguments` snapshot goes out as the only TOOL_CALL_ARGS delta, so + // the deltas always add up to the final arguments. + const endToolCall = function* ( + toolCallId: string, + state: ToolCallState, + ): Generator { + const end = toolCallEndArguments( + argStringByToolCallId.get(toolCallId), + state, + ) + if (!argStringByToolCallId.get(toolCallId)?.trim() && end.args !== '{}') { + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId, + model, + timestamp, + delta: end.args, + args: end.args, + } + } + yield { + type: EventType.TOOL_CALL_END, + toolCallId, + toolName: state.name, + model, + timestamp, + ...end, + } + } + // Seals any in-flight messages and tool calls. Called both on the // normal terminal path (`interaction.complete`) and on the error path // (`error` SSE event + premature EOF) so the StreamProcessor never @@ -1245,14 +1292,7 @@ async function* translateInteractionEvents( for (const [toolCallId, state] of toolCalls) { if (state.ended) continue state.ended = true - yield { - type: EventType.TOOL_CALL_END, - toolCallId, - toolName: state.name, - model, - timestamp, - input: state.args, - } + yield* endToolCall(toolCallId, state) } if (hasEmittedTextMessageStart) { hasEmittedTextMessageStart = false @@ -1307,8 +1347,9 @@ async function* translateInteractionEvents( indexToToolCallId.set(index, toolCallId) // `step.arguments` is required on FunctionCallStep but may // be an empty `{}` placeholder when streaming, where the - // real args arrive as `arguments_delta` events. Treat both - // uniformly: stash whatever we got, stringify once. + // real args arrive as `arguments_delta` events. Keep this + // initial snapshot apart from the streamed text: it is sent + // only when no argument delta arrives (see `endToolCall`). const initialArgs = step.arguments const state: ToolCallState = { name: step.name, @@ -1318,12 +1359,6 @@ async function* translateInteractionEvents( ended: false, } toolCalls.set(toolCallId, state) - argStringByToolCallId.set( - toolCallId, - Object.keys(initialArgs).length > 0 - ? JSON.stringify(initialArgs) - : '', - ) yield { type: EventType.TOOL_CALL_START, toolCallId, @@ -1337,17 +1372,6 @@ async function* translateInteractionEvents( timestamp, index: state.index, } - if (Object.keys(initialArgs).length > 0) { - const argsJson = JSON.stringify(initialArgs) - yield { - type: EventType.TOOL_CALL_ARGS, - toolCallId, - model, - timestamp, - delta: argsJson, - args: argsJson, - } - } break } case 'thought': { @@ -1579,8 +1603,9 @@ async function* translateInteractionEvents( break } const state = toolCalls.get(toolCallId) - if (!state) break - const fragment = delta.arguments ?? '' + // A delta with no arguments keeps the initial snapshot. + if (!state || typeof delta.arguments !== 'string') break + const fragment = delta.arguments const buffer = (argStringByToolCallId.get(toolCallId) ?? '') + fragment argStringByToolCallId.set(toolCallId, buffer) @@ -1693,14 +1718,7 @@ async function* translateInteractionEvents( const state = toolCalls.get(toolCallId) if (state && !state.ended) { state.ended = true - yield { - type: EventType.TOOL_CALL_END, - toolCallId, - toolName: state.name, - model, - timestamp, - input: state.args, - } + yield* endToolCall(toolCallId, state) } indexToToolCallId.delete(event.index) } diff --git a/packages/ai-gemini/tests/gemini-adapter.test.ts b/packages/ai-gemini/tests/gemini-adapter.test.ts index 86fee86930..f3c4dc586a 100644 --- a/packages/ai-gemini/tests/gemini-adapter.test.ts +++ b/packages/ai-gemini/tests/gemini-adapter.test.ts @@ -570,6 +570,83 @@ describe('GeminiAdapter through AI', () => { expect(endIds.sort()).toEqual(startIds.sort()) }) + it.each([ + { args: { count: 2 }, raw: '{"count":2}', input: { count: 2 } }, + { args: '{"n": 1.50}', raw: '{"n": 1.50}', input: { n: 1.5 } }, + ])( + 'sends the raw function arguments $raw on TOOL_CALL_END', + async ({ args, raw, input }) => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [ + { + content: { + parts: [ + { functionCall: { id: 'call_raw', name: 'inspect', args } }, + ], + }, + finishReason: 'STOP', + }, + ], + }, + ]), + ) + + const chunks: Array = [] + for await (const chunk of createTextAdapter().chatStream({ + model: 'gemini-2.5-pro', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + })) { + chunks.push(chunk) + } + + expect(chunks.find((c) => c.type === 'TOOL_CALL_END')).toMatchObject({ + args: raw, + input, + }) + }, + ) + + it('sends malformed function arguments raw on TOOL_CALL_END with no input', async () => { + mocks.generateContentStreamSpy.mockResolvedValue( + createStream([ + { + candidates: [ + { + content: { + parts: [ + { + functionCall: { + id: 'call_bad', + name: 'inspect', + args: '{"a":', + }, + }, + ], + }, + finishReason: 'STOP', + }, + ], + }, + ]), + ) + + const chunks: Array = [] + for await (const chunk of createTextAdapter().chatStream({ + model: 'gemini-2.5-pro', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + })) { + chunks.push(chunk) + } + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('emits grounded web search sources and omits the provider call on the next turn', async () => { const groundingMetadata = { webSearchQueries: ['latest release'], diff --git a/packages/ai-gemini/tests/replay-parity.test.ts b/packages/ai-gemini/tests/replay-parity.test.ts index f4ba64fe63..201e2ade36 100644 --- a/packages/ai-gemini/tests/replay-parity.test.ts +++ b/packages/ai-gemini/tests/replay-parity.test.ts @@ -723,6 +723,45 @@ describe('Gemini replay parity', () => { ).toBe('interaction-real') }) + it('replays malformed tool arguments to the agentic API as {}', async () => { + mocks.interactions.mockResolvedValue({ id: 'i', output_text: 'ok' }) + await collect( + chat({ + adapter: new GeminiTextAdapter({ apiKey: 'key' }, 'gemini-3.7-flash'), + messages: [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call-1', + type: 'function', + function: { name: 'check', arguments: '{"a":' }, + }, + ], + }, + { role: 'tool', toolCallId: 'call-1', content: 'bad arguments' }, + { + role: 'user', + content: [ + { + type: 'video', + source: { type: 'url', value: 'u', mimeType: 'video/mp4' }, + metadata: { processing: 'agentic' }, + }, + ], + }, + ], + }), + ) + expect(mocks.interactions.mock.calls[0]?.[0].input).toContainEqual({ + type: 'function_call', + id: 'call-1', + name: 'check', + arguments: {}, + }) + }) + it('keeps API source independent across concurrent calls on one adapter', async () => { mocks.interactions.mockResolvedValue({ id: 'interaction-concurrent', diff --git a/packages/ai-gemini/tests/text-interactions-adapter.test.ts b/packages/ai-gemini/tests/text-interactions-adapter.test.ts index 30dffa518c..0bff21caed 100644 --- a/packages/ai-gemini/tests/text-interactions-adapter.test.ts +++ b/packages/ai-gemini/tests/text-interactions-adapter.test.ts @@ -563,11 +563,9 @@ describe('GeminiTextInteractionsAdapter', () => { expect(endEvent.input).toEqual({ location: 'Berlin', unit: 'celsius' }) }) - it('preserves the last good args when the arguments stream truncates mid-fragment', async () => { - // If the stream ends after an incomplete fragment, the tool input must - // retain whatever was parsed so far rather than being reset to `{}` — a - // strict JSON.parse would throw on the partial buffer and lose the - // already-merged keys. + it('keeps truncated arguments raw on TOOL_CALL_END with no input', async () => { + // A partial value can appear while the arguments stream. It is not the + // final input: core reports the malformed raw arguments as a tool error. mocks.interactionsCreateSpy.mockResolvedValue( mkStream([ { @@ -602,21 +600,104 @@ describe('GeminiTextInteractionsAdapter', () => { ]), ) - const adapter = createAdapter() const chunks = await collectChunks( - chat({ - adapter, + createAdapter().chatStream({ + model: 'gemini-2.5-flash', + logger: testLogger, messages: [{ role: 'user', content: 'Weather in London?' }], tools: [{ name: 'lookup_weather', description: 'Return the weather' }], }), ) - // No parse-failure RUN_ERROR, and the completed key survives truncation. expect(chunks.find((c) => c.type === 'RUN_ERROR')).toBeUndefined() - const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') as any - expect(endEvent.input).toEqual({ city: 'London' }) + const endEvent = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(endEvent).toMatchObject({ args: '{"city":"London","unit":' }) + expect(endEvent).not.toHaveProperty('input') }) + const rawArgumentCases: Array<{ + name: string + deltas: Array + args: string + input: Record + }> = [ + { + name: 'the initial snapshot when no argument delta arrives', + deltas: [], + args: '{"initial":1}', + input: { initial: 1 }, + }, + { + name: 'the streamed text over the initial snapshot', + deltas: ['{"n": ', '1.50}'], + args: '{"n": 1.50}', + input: { n: 1.5 }, + }, + { + name: 'the initial snapshot when a delta has no arguments', + deltas: [undefined], + args: '{"initial":1}', + input: { initial: 1 }, + }, + ] + + it.each(rawArgumentCases)( + 'sends $name as the raw TOOL_CALL_END args', + async ({ deltas, args, input }) => { + mocks.interactionsCreateSpy.mockResolvedValue( + mkStream([ + { + event_type: 'interaction.created', + interaction: { id: 'int_raw', status: 'in_progress' }, + }, + { + event_type: 'step.start', + index: 0, + step: { + type: 'function_call', + id: 'call_raw', + name: 'inspect', + arguments: { initial: 1 }, + }, + }, + ...deltas.map((fragment) => ({ + event_type: 'step.delta', + index: 0, + delta: { + type: 'arguments_delta', + ...(fragment !== undefined && { arguments: fragment }), + }, + })), + { event_type: 'step.stop', index: 0 }, + { + event_type: 'interaction.completed', + interaction: { id: 'int_raw', status: 'completed' }, + }, + ]), + ) + + const chunks = await collectChunks( + createAdapter().chatStream({ + model: 'gemini-2.5-flash', + logger: testLogger, + messages: [{ role: 'user', content: 'Go' }], + tools: [{ name: 'inspect', description: 'Inspect' }], + }), + ) + + expect(chunks.find((c) => c.type === 'TOOL_CALL_END')).toMatchObject({ + args, + input, + }) + // A client that only reads the deltas gets the same arguments. + const streamed = chunks + .filter((c) => c.type === 'TOOL_CALL_ARGS') + .map((c) => ('delta' in c ? c.delta : '')) + .join('') + expect(streamed).toBe(args) + }, + ) + it('translates thought_summary deltas into REASONING events', async () => { mocks.interactionsCreateSpy.mockResolvedValue( mkStream([ diff --git a/packages/ai-mistral/src/adapters/text.ts b/packages/ai-mistral/src/adapters/text.ts index 40845ff0b5..b141e7f72e 100644 --- a/packages/ai-mistral/src/adapters/text.ts +++ b/packages/ai-mistral/src/adapters/text.ts @@ -750,6 +750,7 @@ export class MistralTextAdapter< toolName: toolCall.name, model: chunkModel, timestamp, + args: toolCall.arguments, input: parsedInput, }) } @@ -834,6 +835,9 @@ export class MistralTextAdapter< } for (const [, toolCall] of toolCallsInProgress) { if (toolCall.started && !toolCall.ended) { + // Parse before marking ended, so a parse error still lets the + // catch below send TOOL_CALL_END with the raw args. + const parsedInput = parseToolCallInput(toolCall, normalizeToolInput) toolCall.ended = true hasEmittedToolCall = true yield asChunk({ @@ -843,7 +847,8 @@ export class MistralTextAdapter< toolName: toolCall.name, model: lastChunkModel, timestamp, - input: parseToolCallInput(toolCall, normalizeToolInput), + args: toolCall.arguments, + input: parsedInput, }) } } @@ -900,9 +905,9 @@ export class MistralTextAdapter< for (const [, toolCall] of toolCallsInProgress) { if (toolCall.started && !toolCall.ended) { toolCall.ended = true - // Best-effort parse for the partial args; if invalid, surface - // empty input rather than throwing inside the cleanup path. - let partialInput: unknown = {} + // Best-effort parse for the partial args. If invalid, leave input + // out and send the raw args, so core reports the parse error. + let partialInput: unknown try { partialInput = toolCall.arguments ? normalizeToolInput( @@ -911,7 +916,7 @@ export class MistralTextAdapter< ) : {} } catch { - partialInput = {} + partialInput = undefined } yield asChunk({ type: 'TOOL_CALL_END', @@ -920,7 +925,8 @@ export class MistralTextAdapter< toolName: toolCall.name, model: lastChunkModel, timestamp, - input: partialInput, + args: toolCall.arguments, + ...(partialInput !== undefined && { input: partialInput }), }) } } diff --git a/packages/ai-mistral/tests/mistral-adapter.test.ts b/packages/ai-mistral/tests/mistral-adapter.test.ts index 5212814b9a..0d0fea0483 100644 --- a/packages/ai-mistral/tests/mistral-adapter.test.ts +++ b/packages/ai-mistral/tests/mistral-adapter.test.ts @@ -1087,6 +1087,109 @@ describe('Mistral AG-UI event emission', () => { expect(caught?.message).toMatch(/lookup_weather/) }) + it('puts the raw argument string on TOOL_CALL_END', async () => { + setupMockStream([ + { + id: 'cmpl-raw', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [ + { + index: 0, + id: 'tc_raw', + function: { + name: 'lookup_weather', + arguments: '{"location": "Berlin", ', + }, + }, + ], + }, + finishReason: null, + }, + ], + }, + { + id: 'cmpl-raw', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [{ index: 0, function: { arguments: '"n": 1.50}' } }], + }, + finishReason: 'tool_calls', + }, + ], + }, + ]) + + const adapter = createMistralText('mistral-large-latest', 'test-api-key') + const chunks: Array = [] + for await (const chunk of adapter.chatStream( + chatOpts({ + model: 'mistral-large-latest', + messages: [{ role: 'user', content: 'Weather in Berlin?' }], + tools: [weatherTool], + }), + )) { + chunks.push(chunk) + } + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ + toolCallId: 'tc_raw', + args: '{"location": "Berlin", "n": 1.50}', + input: { location: 'Berlin' }, + }) + }) + + it('sends malformed raw arguments on TOOL_CALL_END without an input', async () => { + setupMockStream([ + { + id: 'cmpl-bad', + model: 'mistral-large-latest', + choices: [ + { + index: 0, + delta: { + toolCalls: [ + { + index: 0, + id: 'tc_bad', + function: { name: 'lookup_weather', arguments: '{"a":' }, + }, + ], + }, + finishReason: 'tool_calls', + }, + ], + }, + ]) + + const adapter = createMistralText('mistral-large-latest', 'test-api-key') + const chunks: Array = [] + await expect( + (async () => { + for await (const chunk of adapter.chatStream( + chatOpts({ + model: 'mistral-large-latest', + messages: [{ role: 'user', content: 'x' }], + tools: [weatherTool], + }), + )) { + chunks.push(chunk) + } + })(), + ).rejects.toThrow(/Failed to parse tool call arguments/) + + const end = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(end).toMatchObject({ toolCallId: 'tc_bad', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('sends stream_options.include_usage so Mistral returns usage on streaming', async () => { let capturedBody: unknown vi.stubGlobal( diff --git a/packages/ai-ollama/src/adapters/text.ts b/packages/ai-ollama/src/adapters/text.ts index bc87d6cfee..ca1b59df65 100644 --- a/packages/ai-ollama/src/adapters/text.ts +++ b/packages/ai-ollama/src/adapters/text.ts @@ -355,7 +355,7 @@ export class OllamaTextAdapter extends BaseTextAdapter< } // Serialize arguments to a string for the TOOL_CALL_ARGS event - let parsedInput: unknown = {} + let parsedInput: unknown const argsStr = typeof actualToolCall.function.arguments === 'string' ? actualToolCall.function.arguments @@ -364,7 +364,8 @@ export class OllamaTextAdapter extends BaseTextAdapter< const parsed = JSON.parse(argsStr) parsedInput = parsed && typeof parsed === 'object' ? parsed : {} } catch { - parsedInput = actualToolCall.function.arguments + // Leave input out. Core reports the malformed raw args. + parsedInput = undefined } // Emit TOOL_CALL_ARGS with full args (Ollama doesn't stream args incrementally) @@ -385,7 +386,8 @@ export class OllamaTextAdapter extends BaseTextAdapter< toolName: actualToolCall.function.name || '', model: chunk.model, timestamp: Date.now(), - input: parsedInput, + args: argsStr, + ...(parsedInput !== undefined && { input: parsedInput }), }) return events diff --git a/packages/ai-ollama/tests/text-adapter.test.ts b/packages/ai-ollama/tests/text-adapter.test.ts index dd759c34b7..14dde3a044 100644 --- a/packages/ai-ollama/tests/text-adapter.test.ts +++ b/packages/ai-ollama/tests/text-adapter.test.ts @@ -276,6 +276,53 @@ describe('OllamaTextAdapter.chatStream (tool calls)', () => { expect(startChunk!.toolCallId).toMatch(/^search_\d+/) }) + async function toolCallEnd(args: unknown) { + chatMock.mockResolvedValueOnce( + asyncIterable([ + { + message: { + role: 'assistant', + content: '', + tool_calls: [ + { id: 'tc-raw', function: { name: 'search', arguments: args } }, + ], + }, + done: true, + done_reason: 'stop', + }, + ]), + ) + const chunks = await collectStream( + createOllamaChat('llama3.2').chatStream({ + logger: testLogger, + model: 'llama3.2', + messages: [{ role: 'user', content: 'find cats' }], + tools: [searchTool], + }), + ) + return chunks.find((c) => c.type === 'TOOL_CALL_END') + } + + it('puts the raw argument string on TOOL_CALL_END', async () => { + const end = await toolCallEnd('{"q": "cats", "n": 1.50}') + expect(end).toMatchObject({ + toolCallId: 'tc-raw', + args: '{"q": "cats", "n": 1.50}', + input: { q: 'cats', n: 1.5 }, + }) + }) + + it('puts the serialized object arguments on TOOL_CALL_END', async () => { + const end = await toolCallEnd({ q: 'cats' }) + expect(end).toMatchObject({ args: '{"q":"cats"}', input: { q: 'cats' } }) + }) + + it('sends malformed raw arguments on TOOL_CALL_END without an input', async () => { + const end = await toolCallEnd('{"a":') + expect(end).toMatchObject({ toolCallId: 'tc-raw', args: '{"a":' }) + expect(end).not.toHaveProperty('input') + }) + it('forwards tools to the ollama client in provider format', async () => { chatMock.mockResolvedValueOnce( asyncIterable([ diff --git a/packages/ai-openrouter/src/adapters/responses-text.ts b/packages/ai-openrouter/src/adapters/responses-text.ts index 89a9356964..2c6eb555ab 100644 --- a/packages/ai-openrouter/src/adapters/responses-text.ts +++ b/packages/ai-openrouter/src/adapters/responses-text.ts @@ -1471,7 +1471,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: chunk.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -1482,7 +1482,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: chunk.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } } @@ -1552,7 +1553,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1562,7 +1563,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -1679,7 +1681,7 @@ export class OpenRouterResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1689,7 +1691,8 @@ export class OpenRouterResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined diff --git a/packages/ai-openrouter/src/adapters/text.ts b/packages/ai-openrouter/src/adapters/text.ts index 677d8dc636..884272434f 100644 --- a/packages/ai-openrouter/src/adapters/text.ts +++ b/packages/ai-openrouter/src/adapters/text.ts @@ -1167,7 +1167,7 @@ export class OpenRouterTextAdapter< rawArguments: toolCall.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -1179,7 +1179,8 @@ export class OpenRouterTextAdapter< toolName: toolCall.name, model: chunk.model || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1230,7 +1231,7 @@ export class OpenRouterTextAdapter< rawArguments: toolCall.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -1240,7 +1241,8 @@ export class OpenRouterTextAdapter< toolName: toolCall.name, model: lastModel || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } diff --git a/packages/ai-openrouter/tests/openrouter-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-adapter.test.ts index 00b6235e78..ad9bf2bc3f 100644 --- a/packages/ai-openrouter/tests/openrouter-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-adapter.test.ts @@ -522,6 +522,54 @@ describe('OpenRouter adapter option mapping', () => { } }) + // finishReason=null closes the call in the post-stream drain. + it.each([ + { raw: '{"n": 1.50}', finishReason: 'tool_calls' }, + { raw: '{"n": 1.50}', finishReason: null }, + { raw: '{"a":', finishReason: 'tool_calls' }, + { raw: '{"a":', finishReason: null }, + ])( + 'sends the raw argument string $raw as args on TOOL_CALL_END (finishReason=$finishReason)', + async ({ raw, finishReason }) => { + setupMockSdkClient([ + { + id: 'chatcmpl-raw', + model: 'openai/gpt-4o-mini', + choices: [ + { + delta: { + toolCalls: [ + { + index: 0, + id: 'call_raw', + type: 'function', + function: { name: 'lookup_weather', arguments: raw }, + }, + ], + }, + finishReason, + }, + ], + }, + ]) + + const chunks: Array = [] + for await (const chunk of createAdapter().chatStream({ + model: 'openai/gpt-4o-mini', + messages: [{ role: 'user', content: 'Weather?' }], + tools: [weatherTool], + logger: testLogger, + })) { + chunks.push(chunk) + } + + const toolEnd = chunks.find((c) => c.type === 'TOOL_CALL_END') + expect(toolEnd).toMatchObject({ toolCallId: 'call_raw', args: raw }) + if (raw === '{"a":') expect(toolEnd).not.toHaveProperty('input') + else expect(toolEnd).toMatchObject({ input: { n: 1.5 } }) + }, + ) + it('emits parentMessageId on tool-first tool calls matching the assistant message id', async () => { // Tool call streams before any text content. parentMessageId must bind the // tool call to the same assistant message id the eventual TEXT_MESSAGE_START diff --git a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts index e02554a996..caeaf5bc9d 100644 --- a/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts +++ b/packages/ai-openrouter/tests/openrouter-responses-adapter.test.ts @@ -1353,6 +1353,108 @@ describe('OpenRouter responses adapter — stream event bridge', () => { }) }) + describe('raw tool arguments', () => { + type EndSite = 'arguments.done' | 'output_item.done' | 'response.completed' + + // Each site emits TOOL_CALL_END from a different event. Without a name on + // output_item.added, args.done is deferred to a later backfill. + function rawArgsStream( + site: EndSite, + raw: string, + ): Array> { + const named = { + type: 'function_call', + id: 'item_raw', + callId: 'call_raw', + name: 'lookup_weather', + } + return [ + { + type: 'response.output_item.added', + sequenceNumber: 0, + outputIndex: 0, + item: + site === 'arguments.done' + ? { ...named, arguments: '' } + : { + type: 'function_call', + id: 'item_raw', + callId: 'call_raw', + arguments: '', + }, + }, + { + type: 'response.function_call_arguments.done', + sequenceNumber: 1, + itemId: 'item_raw', + outputIndex: 0, + arguments: raw, + }, + ...(site === 'output_item.done' + ? [ + { + type: 'response.output_item.done', + sequenceNumber: 2, + outputIndex: 0, + item: { ...named, arguments: raw }, + }, + ] + : []), + { + type: 'response.completed', + sequenceNumber: 3, + response: { + model: 'm', + output: [{ ...named, arguments: raw }], + usage: { inputTokens: 0, outputTokens: 0, totalTokens: 0 }, + }, + }, + ] + } + + async function toolEnds(site: EndSite, raw: string) { + setupMockSdkClient(rawArgsStream(site, raw)) + const chunks: Array = [] + for await (const chunk of createAdapter().chatStream({ + model: 'openai/gpt-4o-mini', + messages: [{ role: 'user', content: 'hi' }], + tools: [weatherTool], + logger: testLogger, + })) { + chunks.push(chunk) + } + return chunks.filter((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])( + 'sends the raw argument string as args on TOOL_CALL_END (%s)', + async (site) => { + const ends = await toolEnds(site, '{"n": 1.50}') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ + toolCallId: 'call_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])('sends malformed arguments as args with no input (%s)', async (site) => { + const ends = await toolEnds(site, '{"a":') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ toolCallId: 'call_raw', args: '{"a":' }) + expect(ends[0]).not.toHaveProperty('input') + }) + }) + it('round-trips distinct Responses item and call IDs through a server tool', async () => { const firstTurn = [ { diff --git a/packages/ai/package.json b/packages/ai/package.json index 9f2eca7f77..81fd734828 100644 --- a/packages/ai/package.json +++ b/packages/ai/package.json @@ -102,6 +102,7 @@ ], "dependencies": { "@ag-ui/core": "1.0.0", + "@cfworker/json-schema": "^4.1.1", "@standard-schema/spec": "^1.1.0", "@tanstack/ai-event-client": "workspace:^", "@tanstack/ai-utils": "workspace:^", diff --git a/packages/ai/src/activities/chat/agents/spawn.ts b/packages/ai/src/activities/chat/agents/spawn.ts index f8d5bd3753..7637cd4de5 100644 --- a/packages/ai/src/activities/chat/agents/spawn.ts +++ b/packages/ai/src/activities/chat/agents/spawn.ts @@ -1073,8 +1073,8 @@ function agentToolSchema( } /** - * The input schema of an agent tool. The tool loop checks a Standard Schema - * before middleware runs, as for any tool. With reserved fields, a wrapper + * The input schema of an agent tool. The tool loop checks it after + * middleware, just before the tool runs, as for any tool. With reserved fields, a wrapper * checks the agent's fields and passes the reserved ones through. */ function agentToolInput( diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 31b483eb63..5fead436a2 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -85,6 +85,7 @@ import { SUBAGENT_PLAN_KEY, readSubagentTurn } from './agents/turn' import { MiddlewareAbortError, ToolCallManager, + approvalResolution, executeToolCalls, } from './tools/tool-calls' import { @@ -136,8 +137,11 @@ import type { import type { ApprovalRequest, ClientToolRequest, + ExecuteToolCallsResult, McpInputRequest, + ToolExecutionMiddlewareHooks, ToolResult, + ToolResumeExecutionState, } from './tools/tool-calls' import type { ApprovalSchemaConfig, @@ -615,6 +619,11 @@ export interface TextActivityOptions< agentLoopStrategy?: TextOptions['agentLoopStrategy'] /** How the server tools of one model turn run. Default `'parallel'`. */ toolExecution?: TextOptions['toolExecution'] + /** + * How a tool input that fails the final check is repaired. Default: the + * built-in repair. `false`: no repair. A function replaces it. + */ + repairToolInput?: TextOptions['repairToolInput'] /** * Retry a rate-limited model call after the wait the provider asks for. * @@ -2352,8 +2361,19 @@ class TextEngine< } private handleToolCallEndEvent(chunk: ToolCallEndEvent): void { - this.toolCallManager.completeToolCall(chunk) + const final = this.toolCallManager.completeToolCall(chunk) const end = chunk as AdapterYieldChunk + if (final) { + // The client stores the same argument string as the server history. + // The wire carries only `args`; `input` is parsed back from it. + end.args = final.args + end.input = final.input + const metadata = tanstackMetadata(end) + if (metadata) { + delete metadata.args + delete metadata.input + } + } const state = end.state ?? tanstackMetadata(end)?.state if (state !== 'output-error' || end.result === undefined) return this.handleToolCallResultEvent({ @@ -2607,44 +2627,12 @@ class TextEngine< clientToolResults, (eventName, data, options) => this.createCustomEventChunk(eventName, data, options), - { - onBeforeToolCall: async (toolCall, tool, args) => { - this.logger.tools(`phase=before name=${toolCall.function.name}`, { - name: toolCall.function.name, - args, - }) - const hookCtx = { - toolCall, - tool, - args, - toolName: toolCall.function.name, - toolCallId: toolCall.id, - } - return this.middlewareRunner.runOnBeforeToolCall( - this.middlewareCtx, - hookCtx, - ) - }, - onAfterToolCall: async (info) => { - this.logger.tools(`phase=after name=${info.toolName}`, { - name: info.toolName, - result: info.result, - }) - return this.middlewareRunner.runOnAfterToolCall( - this.middlewareCtx, - info, - ) - }, - }, + this.toolMiddlewareHooks(), this.middlewareCtx.context, this.toolAbortSignal, - { - clientToolErrors: this.resumeClientToolErrors, - deniedToolResults: this.resumeDeniedToolResults, - cancelledToolCallIds: this.resumeCancelledToolCallIds, - inputResponses: this.resumeInputResponses, - }, + this.toolResumeState(), this.params.toolExecution, + this.params.repairToolInput, ) // Consume the async generator, yielding custom events and collecting the return value @@ -2873,44 +2861,12 @@ class TextEngine< clientToolResults, (eventName, data, options) => this.createCustomEventChunk(eventName, data, options), - { - onBeforeToolCall: async (toolCall, tool, args) => { - this.logger.tools(`phase=before name=${toolCall.function.name}`, { - name: toolCall.function.name, - args, - }) - const hookCtx = { - toolCall, - tool, - args, - toolName: toolCall.function.name, - toolCallId: toolCall.id, - } - return this.middlewareRunner.runOnBeforeToolCall( - this.middlewareCtx, - hookCtx, - ) - }, - onAfterToolCall: async (info) => { - this.logger.tools(`phase=after name=${info.toolName}`, { - name: info.toolName, - result: info.result, - }) - return this.middlewareRunner.runOnAfterToolCall( - this.middlewareCtx, - info, - ) - }, - }, + this.toolMiddlewareHooks(), this.middlewareCtx.context, this.toolAbortSignal, - { - clientToolErrors: this.resumeClientToolErrors, - deniedToolResults: this.resumeDeniedToolResults, - cancelledToolCallIds: this.resumeCancelledToolCallIds, - inputResponses: this.resumeInputResponses, - }, + this.toolResumeState(), this.params.toolExecution, + this.params.repairToolInput, ) // Consume the async generator, yielding custom events and collecting the return value @@ -2957,7 +2913,8 @@ class TextEngine< finishEvent, toolCalls, afterToolRequests, - executionResult.inputRequired, + // Its results already went out with the chunks above. + { ...executionResult, results: [] }, ) this.setToolPhase('wait') return @@ -3907,7 +3864,8 @@ class TextEngine< requests?: ReadonlyArray< GenericInterruptRequest> >, - inputRequired: ReadonlyArray = [], + // The tool phase result, so an afterTools boundary runs no hook twice. + executed?: ExecuteToolCallsResult, ): AsyncGenerator { this.middlewareCtx.phase = phase const boundaryRequests = @@ -3937,62 +3895,107 @@ class TextEngine< this.addTerminalAssistantMessages() } } - const actionable = this.getBoundaryActionableToolRequests(toolCalls) + const actionable = + executed ?? + (yield* this.drainToolCallGenerator(this.boundaryToolRequests(toolCalls))) + for (const chunk of this.buildToolResultChunks( + actionable.results, + finishEvent, + )) { + yield* this.pipeThroughMiddleware(chunk) + } yield* this.emitActionableInterruptBoundary( finishEvent, - actionable.approvals, - actionable.clientRequests, + actionable.needsApproval, + actionable.needsClientExecution, boundaryRequests, - inputRequired, + actionable.inputRequired, + actionable.subagentInterrupts, ) return true } - private getBoundaryActionableToolRequests( + /** + * The approval previews and client requests a boundary pauses with. They get + * the same middleware hook and final input check as the tool phase, and a + * skip or an invalid input becomes a tool result. No server tool runs here: + * a server tool is in the list only while it waits for approval. + */ + private boundaryToolRequests( toolCalls: ReadonlyArray, - ): { - approvals: Array - clientRequests: Array - } { + ): ReturnType { const { approvals, clientToolResults } = this.collectClientState() - const approvalRequests: Array = [] - const clientRequests: Array = [] - for (const toolCall of toolCalls) { + const tools = this.resolveExecutableTools(toolCalls) + const pendingCalls = toolCalls.filter((toolCall) => { // Provider-executed calls are complete; never surface them as client work. - if (isProviderExecutedToolCall(toolCall)) continue - const tool = this.resolveExecutableTools([toolCall]).find( + if (isProviderExecutedToolCall(toolCall)) return false + const tool = tools.find( (candidate) => candidate.name === toolCall.function.name, - ) as RuntimeToolWithApproval | undefined - if (!tool) continue - let input: unknown = {} - try { - const parsed = JSON.parse(toolCall.function.arguments.trim() || '{}') - input = parsed && typeof parsed === 'object' ? parsed : {} - } catch { - input = {} + ) + if (!tool) return false + if ( + tool.needsApproval && + approvalResolution(approvals, toolCall.id) === undefined + ) { + return true } - const approvalId = `approval_${toolCall.id}` - if (tool.needsApproval && !approvals.has(approvalId)) { - approvalRequests.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - approvalId, - }) - } else if ( + return ( !tool.execute && !clientToolResults.has(toolCall.id) && !this.resumeClientToolErrors.has(toolCall.id) && !this.resumeCancelledToolCallIds.has(toolCall.id) - ) { - clientRequests.push({ - toolCallId: toolCall.id, + ) + }) + return executeToolCalls( + pendingCalls, + tools, + approvals, + clientToolResults, + undefined, + this.toolMiddlewareHooks(), + this.middlewareCtx.context, + this.toolAbortSignal, + this.toolResumeState(), + undefined, + this.params.repairToolInput, + ) + } + + private toolMiddlewareHooks(): ToolExecutionMiddlewareHooks { + return { + onBeforeToolCall: async (toolCall, tool, args) => { + this.logger.tools(`phase=before name=${toolCall.function.name}`, { + name: toolCall.function.name, + args, + }) + return this.middlewareRunner.runOnBeforeToolCall(this.middlewareCtx, { + toolCall, + tool, + args, toolName: toolCall.function.name, - input, + toolCallId: toolCall.id, }) - } + }, + onAfterToolCall: async (info) => { + this.logger.tools(`phase=after name=${info.toolName}`, { + name: info.toolName, + result: info.result, + }) + return this.middlewareRunner.runOnAfterToolCall( + this.middlewareCtx, + info, + ) + }, + } + } + + private toolResumeState(): ToolResumeExecutionState { + return { + clientToolErrors: this.resumeClientToolErrors, + deniedToolResults: this.resumeDeniedToolResults, + cancelledToolCallIds: this.resumeCancelledToolCallIds, + inputResponses: this.resumeInputResponses, } - return { approvals: approvalRequests, clientRequests } } private completeEphemeralInterruptBindings(chunk: StreamChunk): StreamChunk { diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 2acb594465..101b920a6f 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -439,7 +439,10 @@ export interface ToolCallHookContext { toolCall: ToolCall /** The resolved tool definition, if found */ tool: Tool | undefined - /** Parsed arguments for the tool call */ + /** + * Parsed arguments for the tool call, before the `inputSchema` check: no + * schema defaults or transforms yet, and possibly invalid. + */ args: unknown /** Name of the tool */ toolName: string @@ -845,8 +848,10 @@ export interface ChatMiddleware< | Promise | null> /** - * Called before a tool is executed. + * Called before a server tool runs and before a client tool is sent to the + * client. Not called while the call waits for approval. * Can observe, transform args, skip execution, or abort the run. + * The tool's `inputSchema` checks the args after this hook. */ onBeforeToolCall?: ( ctx: ChatMiddlewareContext, diff --git a/packages/ai/src/activities/chat/stream/processor.ts b/packages/ai/src/activities/chat/stream/processor.ts index 12c960686d..cafc8a8e6f 100644 --- a/packages/ai/src/activities/chat/stream/processor.ts +++ b/packages/ai/src/activities/chat/stream/processor.ts @@ -35,6 +35,10 @@ import { withTanstackMetadata, } from '../../../utilities/merge-metadata' import { getChunkRunId } from '../../../utilities/chunk-ids' +import { + finishToolCallArguments, + parseToolArguments, +} from '../../../utilities/tool-call-arguments' import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk' import { isContentPartArray, @@ -214,38 +218,6 @@ function interruptBatchHasGeneric(interrupts: Array): boolean { }) } -/** - * The canonical arguments string for a `TOOL_CALL_END.input`, or `undefined` - * when JSON cannot carry it: `JSON.stringify` returns `undefined` (despite - * its declared type) for a top-level function or symbol and throws on BigInt - * and circular references. `null` is not tool arguments either, so it also - * keeps the streamed value instead of writing `arguments = "null"`. - */ -function serializeToolInput(input: unknown): string | undefined { - if (input === undefined || input === null) return undefined - try { - const serialized: unknown = JSON.stringify(input) - return typeof serialized === 'string' ? serialized : undefined - } catch { - return undefined - } -} - -/** - * The input of a tool call from its raw arguments: `{}` for an empty string, - * `undefined` when the arguments are not complete JSON (for example cut off) - * or are `null`. - */ -function parseToolArguments(args: string): unknown { - if (args.trim() === '') return {} - try { - const parsed: unknown = JSON.parse(args) - return parsed ?? undefined - } catch { - return undefined - } -} - /** * StreamProcessor - State machine for processing AI response streams * @@ -2480,9 +2452,9 @@ export class StreamProcessor { * Handle TOOL_CALL_END event — arguments are finalized (input-complete). * Tool output arrives on TOOL_CALL_RESULT, not on this event. * - * If TOOL_CALL_END carries parsed `input`, it is the canonical arguments: - * write it into the accumulated string and override the rendered part's - * `input` with it. That covers adapters that deliver the whole input on END + * The END's raw `args` and parsed `input` replace the streamed arguments + * (see `finishToolCallArguments`), and the part's `input` is parsed from + * the result. That covers adapters that deliver the whole input on END * (e.g. Anthropic server_tool_use / web_search — issue #839) and adapters * that stream the wire arguments and then normalize them (OpenAI strict-mode * null widening, undone by the adapter after #939), so `arguments` and @@ -2499,47 +2471,19 @@ export class StreamProcessor { const msgState = this.getMessageState(messageId) if (!msgState) return - // The parsed input can ride on the spec `input` field or, for adapters - // that only stamp it into TanStack metadata, on `metadata.tanstack.input`. - const input = - chunk.input !== undefined - ? chunk.input - : (tanstackMetadata(chunk)?.input as unknown) - // Transition the tool call to input-complete (the authoritative completion signal) const existingToolCall = msgState.toolCalls.get(chunk.toolCallId) if (existingToolCall && existingToolCall.state !== 'input-complete') { - // The parsed input replaces the accumulated arguments string, so - // completeToolCall's strict parse surfaces the canonical value on the - // ToolCallPart even when the streamed deltas carried a provider-side - // reshaping of it. An input JSON cannot carry is not canonical: both - // `arguments` and `input` then stay with the streamed value, so the - // two never disagree. - const serializedInput = serializeToolInput(input) - if (serializedInput !== undefined) { - existingToolCall.arguments = serializedInput - } + // The END's arguments replace the streamed ones, so completeToolCall's + // strict parse surfaces the canonical value on the ToolCallPart even + // when the streamed deltas carried a provider-side reshaping of it. + existingToolCall.arguments = finishToolCallArguments( + existingToolCall.arguments, + chunk, + ).args const index = msgState.toolCallOrder.indexOf(chunk.toolCallId) this.completeToolCall(messageId, index, existingToolCall) - - // Canonicalize on the parsed input: overrides the accumulated-args parse - // that completeToolCall wrote (adapters may coerce values differently - // between streamed args and the final structured input). - if (serializedInput !== undefined) { - existingToolCall.parsedArguments = input - this.messages = updateToolCallPart(this.messages, messageId, { - id: existingToolCall.id, - name: existingToolCall.name, - arguments: existingToolCall.arguments, - state: 'input-complete', - input, - ...(existingToolCall.metadata !== undefined && { - metadata: existingToolCall.metadata, - }), - }) - this.emitMessagesChange() - } } } diff --git a/packages/ai/src/activities/chat/tools/input-validation.ts b/packages/ai/src/activities/chat/tools/input-validation.ts new file mode 100644 index 0000000000..b84d083b55 --- /dev/null +++ b/packages/ai/src/activities/chat/tools/input-validation.ts @@ -0,0 +1,557 @@ +import { dereference, initialBaseURI, validate } from '@cfworker/json-schema' +import { + StandardSchemaValidationError, + convertSchemaToJsonSchema, + isStandardJSONSchema, + isStandardSchema, + validateWithStandardSchema, +} from './schema-converter' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { Schema, SchemaDraft } from '@cfworker/json-schema' +import type { JSONSchema, SchemaInput } from '../../../types' + +/** A JSON Schema ready for `validate`: the copy, its draft, its `$ref` lookup. */ +interface CompiledSchema { + schema: Schema | boolean + draft: SchemaDraft + lookup: Record +} + +const compiled = new WeakMap() + +/** Keywords whose value is data, not a schema. */ +const DATA_KEYWORDS = new Set([ + 'const', + 'enum', + 'default', + 'examples', + 'example', +]) + +/** Keywords whose value maps names to schemas. */ +const SCHEMA_MAP_KEYWORDS = new Set([ + 'properties', + 'patternProperties', + '$defs', + 'definitions', + 'dependentSchemas', + 'dependencies', +]) + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null && !Array.isArray(value) +} + +/** Draft-04 and OpenAPI 3.0 write these bounds as booleans. */ +const BOOLEAN_BOUNDS = [ + ['exclusiveMinimum', 'minimum'], + ['exclusiveMaximum', 'maximum'], +] as const + +/** Keywords that can reject `null` even when `type` allows it. */ +const NULL_REJECTING_KEYWORDS = ['$ref', 'allOf', 'anyOf', 'oneOf', 'not'] + +/** Keywords that name a schema. They stay on a wrapped schema. */ +const NAMING_KEYWORDS = new Set(['$id', 'id', '$schema', '$anchor']) + +/** A key as a JSON Pointer token, URI-encoded as cfworker stores paths. */ +function pointerToken(key: string) { + return encodeURI(key.replaceAll('~', '~0').replaceAll('/', '~1')) +} + +/** Calls `visit` on each schema in `node`, inner schemas first. */ +function eachSchema( + node: unknown, + path: string, + visit: (schema: Record, path: string) => void, +): void { + if (Array.isArray(node)) { + node.forEach((item, index) => eachSchema(item, `${path}/${index}`, visit)) + return + } + if (!isRecord(node)) return + for (const [key, value] of Object.entries(node)) { + if (DATA_KEYWORDS.has(key)) continue + const keyPath = `${path}/${pointerToken(key)}` + if (SCHEMA_MAP_KEYWORDS.has(key) && isRecord(value)) { + for (const [name, schema] of Object.entries(value)) { + eachSchema(schema, `${keyPath}/${pointerToken(name)}`, visit) + } + } else { + eachSchema(value, keyPath, visit) + } + } + visit(node, path) +} + +/** + * Make one schema of the JSON copy ready for the validator, in place: + * - `format` goes away. It is an annotation, as in 2019-09 and later: a tool + * that declares `format: 'uri'` still gets a plain path. + * - A boolean bound (`exclusiveMinimum: true`, draft-04 and OpenAPI 3.0) + * becomes the number form, so it works in every draft. + * - OpenAPI 3.0 `nullable: true` lets the value be `null`: `null` joins + * `type` and `enum`. A schema with a `$ref`, `allOf`, `anyOf`, `oneOf`, or + * `not` becomes `{ if: { type: 'null' }, else: schema }`, and its path goes + * in `wrapped`. A failed `else` is an `if` error, which the error text + * leaves out, so the model never reads "Expected null". + */ +function prepareSchema( + schema: Record, + path: string, + wrapped: Array, +) { + if (typeof schema.format === 'string') delete schema.format + for (const [bound, limit] of BOOLEAN_BOUNDS) { + if (typeof schema[bound] !== 'boolean') continue + if (schema[bound] && typeof schema[limit] === 'number') { + schema[bound] = schema[limit] + delete schema[limit] + } else { + delete schema[bound] + } + } + if (schema.nullable !== true) return + if (schema.type !== undefined) { + schema.type = [...new Set([schema.type, 'null'].flat())] + } + if (Array.isArray(schema.enum) && !schema.enum.includes(null)) { + schema.enum.push(null) + } + if (!NULL_REJECTING_KEYWORDS.some((key) => key in schema)) return + const inner: Record = {} + for (const [key, value] of Object.entries(schema)) { + if (NAMING_KEYWORDS.has(key)) continue + inner[key] = value + delete schema[key] + } + schema.if = { type: 'null' } + schema.else = inner + wrapped.push(path) +} + +/** Any base works. A ref resolves as a URL, as cfworker reads it. */ +const REF_BASE = 'https://ref.invalid/' + +/** A `$ref` into a wrapped schema, moved to where that schema went. */ +function movedRef(ref: string, wrappedPath: string) { + if (!ref.startsWith('#')) return ref + const target = new URL(ref, REF_BASE).href + const from = new URL(wrappedPath, REF_BASE).href + return target.startsWith(`${from}/`) + ? `${from}/else${target.slice(from.length)}`.slice(REF_BASE.length) + : ref +} + +/** The draft named by `$schema`, else 2020-12. */ +function draftOf(schema: JSONSchema): SchemaDraft { + const uri = typeof schema.$schema === 'string' ? schema.$schema : '' + // Boolean bounds are numbers by now, and cfworker reads numeric bounds as + // booleans in its draft-04 mode. Draft 7 checks the rest the same way. + if (/draft-0[467]/.test(uri)) return '7' + if (uri.includes('2019-09')) return '2019-09' + return '2020-12' +} + +function compile(schema: JSONSchema | boolean): CompiledSchema { + if (typeof schema === 'boolean') { + return { schema, draft: '2020-12', lookup: dereference(schema) } + } + const cached = compiled.get(schema) + if (cached) return cached + // `dereference` writes lookup data onto the schema it gets. A JSON copy + // keeps the tool's own schema, which also goes to the provider, unchanged. + const copy: Schema = JSON.parse(JSON.stringify(schema)) + const wrapped: Array = [] + eachSchema(copy, '#', (node, path) => prepareSchema(node, path, wrapped)) + // Inner schemas were wrapped first, so their moves apply first. + // ponytail: refs are read as paths from the root. An absolute ref, and a + // ref inside a nullable schema with its own `$id`, are not moved. + if (wrapped.length > 0) { + eachSchema(copy, '#', (node) => { + if (typeof node.$ref === 'string') { + node.$ref = wrapped.reduce(movedRef, node.$ref) + } + }) + } + const result = { + schema: copy, + draft: draftOf(schema), + lookup: dereference(copy), + } + compiled.set(schema, result) + return result +} + +const DECIMAL = /^-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?$/ + +function convertPrimitive(value: unknown, type: string) { + switch (type) { + case 'number': + case 'integer': { + if (value === null) return 0 + if (typeof value === 'boolean') return value ? 1 : 0 + // Decimal text only: no hex, and no integer past the safe range. + if (typeof value !== 'string' || !DECIMAL.test(value)) return value + const number = Number(value) + const valid = + type === 'integer' + ? Number.isSafeInteger(number) + : Number.isFinite(number) + return valid ? number : value + } + case 'boolean': + if (value === null || value === 'false' || value === 0) return false + if (value === 'true' || value === 1) return true + return value + case 'string': + if (value === null) return '' + return typeof value === 'number' || typeof value === 'boolean' + ? String(value) + : value + case 'null': + return value === '' || value === 0 || value === false ? null : value + default: + return value + } +} + +function matchesType(value: unknown, type: string) { + switch (type) { + case 'null': + return value === null + case 'array': + return Array.isArray(value) + case 'object': + return ( + typeof value === 'object' && value !== null && !Array.isArray(value) + ) + case 'integer': + return typeof value === 'number' && Number.isInteger(value) + default: + return typeof value === type + } +} + +/** + * Repair common model mistakes against the schema before the check: `"5"` + * for 5, `"true"` for true, `null` for an optional field. It walks unions, + * `$ref`, `if`/`then`/`else`, tuples, and pattern and additional properties. + * A value that cannot be repaired stays as it is, so the check reports it. + */ +function coerceArguments(value: unknown, root: CompiledSchema) { + const { draft, lookup } = root + + function check(schema: Schema | boolean, candidate: unknown) { + if (typeof schema === 'boolean') return schema + return validate(candidate, schema, draft, lookup).valid + } + + function refTarget(schema: Schema) { + const uri = schema.__absolute_ref__ + const target = uri === undefined ? undefined : lookup[uri] + return typeof target === 'object' ? target : undefined + } + + function coerceUnion( + candidate: unknown, + schemas: Array, + seen: Set, + required: Set, + ) { + if (schemas.some((schema) => check(schema, candidate))) return candidate + for (const schema of schemas) { + const converted = coerce( + structuredClone(candidate), + schema, + seen, + required, + ) + if (check(schema, converted)) return converted + } + return candidate + } + + function requiredKeys( + schema: Schema | boolean, + candidate: unknown, + seen = new Set(), + ): Set { + if (typeof schema === 'boolean' || seen.has(schema)) return new Set() + const nextSeen = new Set(seen).add(schema) + const required = new Set( + Array.isArray(schema.required) ? schema.required : [], + ) + const children: Array = [...(schema.allOf ?? [])] + const target = refTarget(schema) + if (target) children.push(target) + if (schema.if !== undefined) { + const branch = check(schema.if, candidate) ? schema.then : schema.else + if (branch) children.push(branch) + } + for (const child of children) { + for (const key of requiredKeys(child, candidate, nextSeen)) { + required.add(key) + } + } + return required + } + + function coerce( + candidate: unknown, + schema: Schema | boolean, + seen = new Set(), + inheritedRequired = new Set(), + ): unknown { + if (typeof schema === 'boolean') return candidate + // A reference cycle at the same value cannot add another conversion. + if (seen.has(schema)) return candidate + const nextSeen = new Set(seen).add(schema) + const required = requiredKeys(schema, candidate) + for (const key of inheritedRequired) required.add(key) + let next = candidate + const target = refTarget(schema) + if (target) next = coerce(next, target, nextSeen, required) + for (const nested of schema.allOf ?? []) { + next = coerce(next, nested, nextSeen, required) + } + if (schema.anyOf) next = coerceUnion(next, schema.anyOf, nextSeen, required) + if (schema.oneOf) next = coerceUnion(next, schema.oneOf, nextSeen, required) + + const types = + typeof schema.type === 'string' ? [schema.type] : (schema.type ?? []) + if (!types.some((type) => matchesType(next, type))) { + // `null` last: `0` for `['null', 'string']` is `"0"`, not `null`. + const order: Array = types.filter((type) => type !== 'null') + if (types.includes('null')) order.push('null') + for (const type of order) { + const converted = convertPrimitive(next, type) + if (converted !== next) { + next = converted + break + } + } + } + + if (typeof next === 'object' && next !== null && !Array.isArray(next)) { + // Own keys only, `__proto__` included. + const object: Record = Object.fromEntries( + Object.entries(next), + ) + const properties = Object.entries(schema.properties ?? {}) + const knownKeys = new Set(properties.map(([key]) => key)) + for (const [key, property] of properties) { + if (!Object.hasOwn(object, key)) continue + if ( + object[key] === null && + !required.has(key) && + !check(property, null) + ) { + // A strict-mode provider sends `null` for an optional field. + delete object[key] + } else { + object[key] = coerce(object[key], property) + } + } + for (const [pattern, property] of Object.entries( + schema.patternProperties ?? {}, + )) { + const expression = new RegExp(pattern, 'u') + for (const key of Object.keys(object)) { + if (!expression.test(key)) continue + knownKeys.add(key) + object[key] = coerce(object[key], property) + } + } + const additional = schema.additionalProperties + if (typeof additional === 'object') { + for (const [key, entry] of Object.entries(object)) { + if (!knownKeys.has(key)) object[key] = coerce(entry, additional) + } + } + next = object + } + if (Array.isArray(next)) { + const tuple = + schema.prefixItems ?? + (Array.isArray(schema.items) ? schema.items : undefined) + const rest = Array.isArray(schema.items) ? undefined : schema.items + next = next.map((item: unknown, index) => { + const itemSchema = tuple?.[index] ?? rest + return itemSchema === undefined ? item : coerce(item, itemSchema) + }) + } + if (schema.if !== undefined) { + const branch = check(schema.if, next) ? schema.then : schema.else + if (branch) next = coerce(next, branch, nextSeen, required) + } + return next + } + + return coerce(value, root.schema) +} + +/** Errors that only say that a nested check failed. The nested error says why. */ +const WRAPPER_KEYWORDS = new Set([ + 'properties', + 'patternProperties', + 'items', + 'prefixItems', + 'additionalItems', + 'unevaluatedItems', + '$ref', + '$recursiveRef', + 'allOf', + // "Instance does not match "then" schema." (or "else") + 'if', +]) + +/** More errors than this are counted, not listed. */ +const MAX_ERRORS = 10 + +function validateJsonInput( + schema: JSONSchema | boolean, + received: unknown, + toolName: string, + repair: boolean, +) { + let errors: string + try { + // Check the JSON form: a key set to `undefined` is the same as no key, + // and an input with no JSON form (`undefined`) is `null`. + const json = JSON.stringify(received) ?? 'null' + const value: unknown = JSON.parse(json) + const root = compile(schema) + const check = (candidate: unknown) => + validate(candidate, root.schema, root.draft, root.lookup, false) + const result = check(value) + if (result.valid) return received + // Repair only an input that fails. The error reports the input as sent. + if (repair) { + let converted = value + try { + converted = coerceArguments(value, root) + } catch { + // A schema that cannot be used: the first check reports it. + } + if (JSON.stringify(converted) !== json && check(converted).valid) { + return converted + } + } + const leaves = result.errors.filter( + (error) => !WRAPPER_KEYWORDS.has(error.keyword), + ) + const shown = leaves.length > 0 ? leaves : result.errors + errors = shown + .slice(0, MAX_ERRORS) + .map((error) => `${error.instanceLocation}: ${error.error}`) + .join(', ') + if (shown.length > MAX_ERRORS) { + errors += `, and ${shown.length - MAX_ERRORS} more` + } + } catch (error) { + // A schema the validator cannot use, such as an unresolved `$ref`, or a + // value with no JSON form. The first line is enough for the model. + // Drop the validator base (in a browser, the page URL and its query). + const message = error instanceof Error ? error.message : String(error) + errors = message + .split('\n', 1) + .join('') + .replaceAll(initialBaseURI.href, '') + .replace(/ +Absolute URI "[^"]*"\.?/, '') + .trim() + } + throw new Error(`Input validation failed for tool ${toolName}: ${errors}`) +} + +function standardInputError( + issues: ReadonlyArray, + toolName: string, +) { + const error = new StandardSchemaValidationError(issues) + error.message = `Input validation failed for tool ${toolName}: ${error.message}` + return error +} + +/** + * How a failed tool input is repaired. `undefined`: the built-in repair. + * `false`: no repair. A function: called instead of the built-in repair. It + * returns the input to check again, or `undefined` to keep the error. + */ +export type ToolInputRepair = + | false + | ((input: unknown, error: Error) => unknown | Promise) + +/** + * The final check of a tool input, just before the tool runs or the client + * gets the call. A Standard Schema runs here (an async validator too), and + * its output is the input, so transforms apply once. A JSON Schema (an MCP + * tool, a plain JSON Schema tool) is checked with `@cfworker/json-schema`. + * A failed input is repaired as `repair` says, then checked again. The error + * of the first check goes back to the model as the tool result. + */ +export async function validateToolInput( + schema: SchemaInput | boolean | undefined, + received: unknown, + toolName: string, + repair?: ToolInputRepair, +): Promise { + try { + return await checkToolInput( + schema, + received, + toolName, + repair === undefined, + ) + } catch (error) { + if (typeof repair !== 'function') throw error + const repaired = await repair( + received, + error instanceof Error ? error : new Error(String(error)), + ) + if (repaired === undefined) throw error + try { + return await checkToolInput(schema, repaired, toolName, false) + } catch { + throw error + } + } +} + +async function checkToolInput( + schema: SchemaInput | boolean | undefined, + received: unknown, + toolName: string, + repair: boolean, +): Promise { + if (schema === undefined) return received + if (typeof schema === 'boolean') { + return validateJsonInput(schema, received, toolName, repair) + } + if (!isStandardSchema(schema)) { + const jsonSchema = convertSchemaToJsonSchema(schema, { io: 'input' }) + return jsonSchema + ? validateJsonInput(jsonSchema, received, toolName, repair) + : received + } + const first = await validateWithStandardSchema(schema, received) + if (first.success) return first.data + // The authored schema runs first. When it fails and it can export JSON + // Schema, repair the input and run the authored schema again. A failure + // reports the issues of the first run. + const original = standardInputError(first.issues, toolName) + if (!repair || !isStandardJSONSchema(schema)) throw original + let converted: unknown + try { + const jsonSchema = convertSchemaToJsonSchema(schema, { io: 'input' }) + if (!jsonSchema) throw original + converted = validateJsonInput(jsonSchema, received, toolName, true) + } catch { + throw original + } + // Nothing was repaired: a second run gives the same issues. + if (converted === received) throw original + const retried = await validateWithStandardSchema(schema, converted) + if (!retried.success) throw original + return retried.data +} diff --git a/packages/ai/src/activities/chat/tools/tool-calls.ts b/packages/ai/src/activities/chat/tools/tool-calls.ts index 2b889b2934..bd697e8edf 100644 --- a/packages/ai/src/activities/chat/tools/tool-calls.ts +++ b/packages/ai/src/activities/chat/tools/tool-calls.ts @@ -1,7 +1,9 @@ +import { finishToolCallArguments } from '../../../utilities/tool-call-arguments' import { normalizeToolResult } from '../../../utilities/tool-result' import { tanstackMetadata } from '../../../utilities/merge-metadata' import { isProviderExecutedToolCall } from '../../../utilities/provider-executed' import { mergeStreams } from '../../../utilities/merge-streams' +import { validateToolInput } from './input-validation' import type { AdapterYieldChunk } from '../../../utilities/adapter-yield-chunk' import { StandardSchemaValidationError, @@ -351,15 +353,16 @@ export class ToolCallManager< * Complete a tool call with its final input * Called when TOOL_CALL_END is received */ - completeToolCall(event: ToolCallEndEvent): void { + completeToolCall( + event: ToolCallEndEvent, + ): { args: string; input: unknown } | undefined { for (const toolCall of this.toolCallsMap.values()) { if (toolCall.id !== event.toolCallId) continue - if (event.input === undefined) return - const normalized = - event.input && typeof event.input === 'object' ? event.input : {} - toolCall.function.arguments = JSON.stringify(normalized) - return + const final = finishToolCallArguments(toolCall.function.arguments, event) + toolCall.function.arguments = final.args + return final } + return undefined } /** @@ -415,20 +418,7 @@ export class ToolCallManager< ) } - // Validate input against inputSchema (for Standard Schema compliant schemas) - if (tool.inputSchema && isStandardSchema(tool.inputSchema)) { - try { - args = parseWithStandardSchema(tool.inputSchema, args) - } catch (validationError: unknown) { - const message = - validationError instanceof Error - ? validationError.message - : 'Validation failed' - throw new Error( - `Input validation failed for tool ${tool.name}: ${message}`, - ) - } - } + args = await validateToolInput(tool.inputSchema, args, tool.name) // Execute the tool const executionContext = { @@ -600,7 +590,7 @@ export interface ToolResumeExecutionState { inputResponses?: ReadonlyMap } -function approvalResolution( +export function approvalResolution( approvals: ReadonlyMap, toolCallId: string, ): ToolApprovalResolution | undefined { @@ -625,7 +615,7 @@ function deniedApprovalResult(resolution: ToolApprovalResolution): unknown { : { error: 'User declined tool execution' } } -interface ExecuteToolCallsResult { +export interface ExecuteToolCallsResult { /** Tool results ready to send to LLM */ results: Array /** Tools that need user approval before execution */ @@ -989,6 +979,8 @@ async function buildClientToolResult( * @param toolExecution - `'parallel'` (default) prepares every call in call * order, then starts the server tools together. `'sequential'` runs one * call at a time. Results come back in call order either way. + * @param repairToolInput - How an input that fails the final check is + * repaired: the built-in repair (default), none (`false`), or a function. */ export async function* executeToolCalls( toolCalls: Array, @@ -1005,6 +997,7 @@ export async function* executeToolCalls( abortSignal?: AbortSignal, resumeState?: ToolResumeExecutionState, toolExecution: NonNullable = 'parallel', + repairToolInput?: TextOptions['repairToolInput'], ): AsyncGenerator { const results: Array = [] const needsApproval: Array = [] @@ -1091,6 +1084,8 @@ export async function* executeToolCalls( !resumeState?.cancelledToolCallIds?.has(tc.id) ) }) + // The calls that the gating holds back. + const held: Array = [] for (const toolCall of toolCalls) { // Provider-executed tools (Anthropic web_search / web_fetch) already ran @@ -1112,13 +1107,17 @@ export async function* executeToolCalls( continue } - // Skip non-pending tools while approvals are outstanding + // Skip non-pending tools while approvals are outstanding. A client tool + // with no approval, or an approved one, is still sent to the client. if (hasPendingApprovals) { - const isPendingApproval = - tool.needsApproval && - approvalResolution(approvals, toolCall.id) === undefined - const isPlainClientRequest = !tool.needsApproval && !tool.execute - if (!isPendingApproval && !isPlainClientRequest) { + const resolution = approvalResolution(approvals, toolCall.id) + const isPendingApproval = tool.needsApproval && resolution === undefined + const isClientRequest = + !tool.execute && + (!tool.needsApproval || + (resolution !== undefined && isApproved(resolution))) + if (!isPendingApproval && !isClientRequest) { + held.push(toolCall) continue } } @@ -1155,28 +1154,135 @@ export async function* executeToolCalls( continue } } + // The input that the model sent, to tell later edits apart from it. + const sentInput = JSON.stringify(input) + + const resolution = tool.needsApproval + ? approvalResolution(approvals, toolCall.id) + : undefined + // A call that waits for approval shows its checked input. Middleware + // runs only once the call is approved. + const needsPreview = tool.needsApproval && resolution === undefined + if (resolution !== undefined && !isApproved(resolution)) { + // User declined + results.push({ + toolCallId: toolCall.id, + toolName, + result: + resumeState?.deniedToolResults?.get(toolCall.id) ?? + deniedApprovalResult(resolution), + input, + state: 'output-error', + outcome: 'denied', + }) + continue + } + if (resolution !== undefined) { + input = editedApprovalArgs(resolution) ?? input + } + // A client tool whose result came back already ran with the input that + // the hook and the check gave it before dispatch. Checking the raw model + // input again would drop a real result. + const clientError = resumeState?.clientToolErrors?.get(toolCall.id) + if ( + !tool.execute && + (clientResults.has(toolCall.id) || clientError !== undefined) + ) { + results.push( + await buildClientToolResult( + toolCall.id, + toolName, + tool, + clientResults.get(toolCall.id), + input, + clientError, + ), + ) + continue + } - // Validate input against inputSchema (for Standard Schema compliant schemas) - if (tool.inputSchema && isStandardSchema(tool.inputSchema)) { + if (!needsPreview && middlewareHooks) { + let decision: Awaited> try { - input = parseWithStandardSchema(tool.inputSchema, input) - } catch (validationError: unknown) { - const message = - validationError instanceof Error - ? validationError.message - : 'Validation failed' - results.push({ + decision = await applyBeforeToolCallDecision( + toolCall, + tool, + input, + toolName, + middlewareHooks, + results, + ) + } catch (error) { + // The calls prepared so far still run, as they would one at a time. + failures.push(error) + break + } + if (!decision.proceed) continue + input = decision.input + } + + // The final check, after approval edits and middleware: invalid input + // becomes a tool error, and the checked value is what runs. + try { + input = await validateToolInput( + tool.inputSchema, + input, + toolName, + typeof repairToolInput === 'function' + ? (failed, error) => + repairToolInput({ tool, toolCall, input: failed, error }) + : repairToolInput, + ) + } catch (error) { + // The model sees its own arguments, never values that middleware put + // in. They stay out of the error, so tracing does not record them. + // When an edit or middleware changed the input, the model is told so, + // so it does not retry the same valid call. + let changed: boolean + try { + changed = JSON.stringify(input) !== sentInput + } catch { + changed = true + } + const message = `${ + error instanceof Error ? error.message : 'Validation failed' + }\nReceived arguments: ${argsStr}${ + changed + ? '\nAn approval edit or middleware changed these arguments before the check. The changed input failed the check.' + : '' + }` + await pushResultAndRunAfterHook( + results, + { toolCallId: toolCall.id, toolName, - result: { - error: `Input validation failed for tool ${tool.name}: ${message}`, - }, - // raw parse may have failed validation — still attach best-effort input + result: { error: message }, input, state: 'output-error', - }) - continue - } + duration: 0, + }, + { + toolCall, + tool, + toolName, + toolCallId: toolCall.id, + ok: false, + duration: 0, + error, + }, + // A preview call runs no hooks. + needsPreview ? undefined : middlewareHooks, + ) + continue + } + if (needsPreview) { + needsApproval.push({ + toolCallId: toolCall.id, + toolName, + input, + approvalId: `approval_${toolCall.id}`, + }) + continue } // Create a ToolExecutionContext for this tool call with event emission @@ -1207,142 +1313,13 @@ export async function* executeToolCalls( }, } as ToolExecutionContext - // CASE 1: Client-side tool (no execute function) + // Client-side tool (no execute function), approved or with no approval if (!tool.execute) { - // Check if tool needs approval - if (tool.needsApproval) { - const approvalId = `approval_${toolCall.id}` - const resolution = approvalResolution(approvals, toolCall.id) - - // Check if approval decision exists - if (resolution !== undefined) { - const approved = isApproved(resolution) - - if (approved) { - input = editedApprovalArgs(resolution) ?? input - // Approved - check if client has executed - const clientError = resumeState?.clientToolErrors?.get(toolCall.id) - if (clientResults.has(toolCall.id) || clientError !== undefined) { - results.push( - await buildClientToolResult( - toolCall.id, - toolName, - tool, - clientResults.get(toolCall.id), - input, - clientError, - ), - ) - } else { - // Approved but not executed yet - request client execution - needsClientExecution.push({ - toolCallId: toolCall.id, - toolName, - input, - }) - } - } else { - // User declined - results.push({ - toolCallId: toolCall.id, - toolName, - result: - resumeState?.deniedToolResults?.get(toolCall.id) ?? - deniedApprovalResult(resolution), - input, - state: 'output-error', - outcome: 'denied', - }) - } - } else { - // Need approval first - needsApproval.push({ - toolCallId: toolCall.id, - toolName: toolCall.function.name, - input, - approvalId, - }) - } - } else { - // No approval needed - check if client has executed - const clientError = resumeState?.clientToolErrors?.get(toolCall.id) - if (clientResults.has(toolCall.id) || clientError !== undefined) { - results.push( - await buildClientToolResult( - toolCall.id, - toolName, - tool, - clientResults.get(toolCall.id), - input, - clientError, - ), - ) - } else { - // Request client execution - needsClientExecution.push({ - toolCallId: toolCall.id, - toolName, - input, - }) - } - } + needsClientExecution.push({ toolCallId: toolCall.id, toolName, input }) continue } - // CASE 2: Server tool with approval required - if (tool.needsApproval) { - const approvalId = `approval_${toolCall.id}` - const resolution = approvalResolution(approvals, toolCall.id) - - if (resolution === undefined) { - // Need approval - needsApproval.push({ - toolCallId: toolCall.id, - toolName, - input, - approvalId, - }) - continue - } - if (!isApproved(resolution)) { - // User declined - results.push({ - toolCallId: toolCall.id, - toolName, - result: - resumeState?.deniedToolResults?.get(toolCall.id) ?? - deniedApprovalResult(resolution), - input, - state: 'output-error', - outcome: 'denied', - }) - continue - } - // Approved: run it like any other server tool below. - input = editedApprovalArgs(resolution) ?? input - } - - // CASE 3: Server tool, approved or with no approval - if (middlewareHooks) { - let decision: Awaited> - try { - decision = await applyBeforeToolCallDecision( - toolCall, - tool, - input, - toolName, - middlewareHooks, - results, - ) - } catch (error) { - // The calls prepared so far still run, as they would one at a time. - failures.push(error) - break - } - if (!decision.proceed) continue - input = decision.input - } - + // Server tool, approved or with no approval const run = runServerTool( toolCall, tool, @@ -1362,12 +1339,36 @@ export async function* executeToolCalls( yield* mergeStreams(runs) if (failures.length > 0) throw failures[0] + // No approval waits after all: each held-back call's preview failed its + // check or its arguments. The held calls run now, so each call has a result. + let heldInterrupts: Array = [] + if (held.length > 0 && needsApproval.length === 0) { + const rest = yield* executeToolCalls( + held, + tools, + approvals, + clientResults, + createCustomEventChunk, + middlewareHooks, + userContext, + abortSignal, + resumeState, + toolExecution, + repairToolInput, + ) + results.push(...rest.results) + needsClientExecution.push(...rest.needsClientExecution) + inputRequired.push(...rest.inputRequired) + heldInterrupts = rest.subagentInterrupts + } + // Parallel tools finish in any order. The model gets the results, input // requests, and interrupts in the order it made the calls. const callOrder = new Map(toolCalls.map((tc, index) => [tc.id, index])) const byCallOrder = (a: { toolCallId: string }, b: { toolCallId: string }) => (callOrder.get(a.toolCallId) ?? 0) - (callOrder.get(b.toolCallId) ?? 0) results.sort(byCallOrder) + needsClientExecution.sort(byCallOrder) inputRequired.sort(byCallOrder) return { @@ -1375,8 +1376,9 @@ export async function* executeToolCalls( needsApproval, needsClientExecution, inputRequired, - subagentInterrupts: toolCalls.flatMap( - (tc) => interruptsByCall.get(tc.id) ?? [], - ), + subagentInterrupts: [ + ...toolCalls.flatMap((tc) => interruptsByCall.get(tc.id) ?? []), + ...heldInterrupts, + ], } } diff --git a/packages/ai/src/middlewares/otel.ts b/packages/ai/src/middlewares/otel.ts index cef80e2e8d..00994530f7 100644 --- a/packages/ai/src/middlewares/otel.ts +++ b/packages/ai/src/middlewares/otel.ts @@ -899,6 +899,9 @@ export function otelMiddleware( onBeforeToolCall(ctx, hookCtx) { safeCall('otel.onBeforeToolCall', () => { + // A client tool runs in the client, not here, so it gets no + // execute_tool span. Middleware sees its input before dispatch. + if (hookCtx.tool && !hookCtx.tool.execute) return const state = stateByCtx.get(ctx) if (!state) return const parent = state.currentIterationSpan ?? state.rootSpan diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 36a4345902..336fcbe2a1 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -773,6 +773,8 @@ export interface TanStackRunMetadata { state?: ToolOutputState /** Parsed `TOOL_CALL_END` input. Spec `TOOL_CALL_END` has no top-level `input`. */ input?: unknown + /** The raw tool argument string on `TOOL_CALL_END`, as the provider sent it. */ + args?: string } /** @@ -1277,6 +1279,20 @@ export interface ResolvedPromptCache { */ export type FetchWrapper = (next: typeof fetch) => typeof fetch +/** + * What a `repairToolInput` function gets: a tool input that failed the check. + */ +export interface ToolInputRepairContext { + /** The tool the model called. */ + tool: AnyTool + /** The call, with its id and the raw arguments text. */ + toolCall: ToolCall + /** The input that failed, after approval edits and middleware. */ + input: unknown + /** The validation error. A Standard Schema error also has `issues`. */ + error: Error +} + /** * Options passed into the SDK and further piped to the AI provider. */ @@ -1314,6 +1330,16 @@ export interface TextOptions< * starts them together, and `'sequential'` runs them one at a time. */ toolExecution?: 'parallel' | 'sequential' + /** + * How a tool input that fails the final check is repaired. By default, + * common model mistakes such as `"5"` for a number are repaired against the + * schema. `false` turns the repair off. A function replaces the built-in + * repair: it returns the input to check again, or `undefined` to send the + * validation error to the model. + */ + repairToolInput?: + | false + | ((context: ToolInputRepairContext) => unknown | Promise) /** * Retry a model call when the provider error has a wait time (usually a * 429). `chat()` waits for the `retryAfterMs` time of the `RUN_ERROR`, then diff --git a/packages/ai/src/utilities/normalize-stream-chunk.test.ts b/packages/ai/src/utilities/normalize-stream-chunk.test.ts index 7954dc147a..27e42095cd 100644 --- a/packages/ai/src/utilities/normalize-stream-chunk.test.ts +++ b/packages/ai/src/utilities/normalize-stream-chunk.test.ts @@ -343,6 +343,37 @@ describe('normalizeStreamChunk', () => { expect(out).not.toHaveProperty('toolCallName') }) + it('sends TOOL_CALL_END args without a second copy as input', () => { + const out = normalizeOne({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + args: '{"q": 1}', + input: { q: 1 }, + } as AdapterYieldChunk) + + expect(out).toEqual({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + metadata: { tanstack: { args: '{"q": 1}' } }, + }) + }) + + it('rewrites TOOL_CALL_END args to match an input that differs', () => { + const out = normalizeOne({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + args: '{"a": 1, "b": null}', + input: { a: 1 }, + metadata: { tanstack: { input: { a: 1 } } }, + } as AdapterYieldChunk) + + expect(out).toEqual({ + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + metadata: { tanstack: { args: '{"a":1}' } }, + }) + }) + it('stringifies an array TOOL_CALL_END result onto RESULT content', () => { const parts = [{ type: 'text', content: 'hello' }] const out = normalizeAll({ diff --git a/packages/ai/src/utilities/normalize-stream-chunk.ts b/packages/ai/src/utilities/normalize-stream-chunk.ts index a706252269..6d4cbcbe78 100644 --- a/packages/ai/src/utilities/normalize-stream-chunk.ts +++ b/packages/ai/src/utilities/normalize-stream-chunk.ts @@ -3,7 +3,8 @@ import type { StreamChunk, ToolCallResultEvent } from '../types' import type { AdapterYieldChunk } from './adapter-yield-chunk' import { isTanstackUsage, toSpecTokenUsage } from './ag-ui-usage' import type { MetadataRecord } from './merge-metadata' -import { withTanstackMetadata } from './merge-metadata' +import { tanstackMetadata, withTanstackMetadata } from './merge-metadata' +import { finishToolCallArguments } from './tool-call-arguments' import { reasoningEncryptedValue } from './reasoning-encrypted-value' import { specKeysFor } from './spec-event-keys' @@ -136,6 +137,17 @@ export function normalizeStreamChunk( if (chunk.type === EventType.TOOL_CALL_START) { skipLeftover.add('toolName') } + // A TOOL_CALL_END sends its arguments once, as `args` made to match + // `input`. `restorePublicUsage` parses `input` back from it. + let sendsArgs = false + if ( + chunk.type === EventType.TOOL_CALL_END && + typeof (chunk.args ?? tanstackMetadata(chunk)?.args) === 'string' + ) { + sendsArgs = true + tanstack.args = finishToolCallArguments('', chunk).args + skipLeftover.add('input') + } for (const key of Object.keys(chunk)) { if (specKeys.has(key) || skipLeftover.has(key)) continue if (tanstack[key] !== undefined) continue @@ -164,6 +176,8 @@ export function normalizeStreamChunk( ? specChunk : withTanstackMetadata(specChunk, tanstack) + if (sendsArgs) delete tanstackMetadata(normalized)?.input + const extras = encryptedValueExtras(chunk) const main: Array = [normalized as StreamChunk] diff --git a/packages/ai/src/utilities/restore-inbound-chunk.test.ts b/packages/ai/src/utilities/restore-inbound-chunk.test.ts index 26d8a451ba..01f632f4e8 100644 --- a/packages/ai/src/utilities/restore-inbound-chunk.test.ts +++ b/packages/ai/src/utilities/restore-inbound-chunk.test.ts @@ -80,6 +80,21 @@ describe('restorePublicUsage', () => { } expect(chunk.input).toEqual({ q: 'sf' }) }) + + it('restores TOOL_CALL_END input from metadata.tanstack.args', () => { + const chunk: StreamChunk = { + type: EventType.TOOL_CALL_END, + toolCallId: 'tc1', + metadata: { tanstack: { args: '{"q": "sf"}' } }, + } + + restorePublicUsage(chunk) + + if (chunk.type !== EventType.TOOL_CALL_END) { + throw new Error('expected TOOL_CALL_END') + } + expect(chunk.input).toEqual({ q: 'sf' }) + }) }) describe('restoreInboundChunk', () => { diff --git a/packages/ai/src/utilities/restore-inbound-chunk.ts b/packages/ai/src/utilities/restore-inbound-chunk.ts index 460ee8137a..ccf24f3849 100644 --- a/packages/ai/src/utilities/restore-inbound-chunk.ts +++ b/packages/ai/src/utilities/restore-inbound-chunk.ts @@ -3,6 +3,7 @@ import type { StreamChunk } from '../types' import type { AdapterYieldChunk } from './adapter-yield-chunk' import { isTanstackUsage, rebuildTokenUsage } from './ag-ui-usage' import { tanstackMetadata } from './merge-metadata' +import { parseToolArguments } from './tool-call-arguments' /** * Rebuild TanStack `TokenUsage` (`promptTokens`) from spec `usage[]` plus @@ -34,7 +35,12 @@ export function restorePublicUsage(chunk: StreamChunk): StreamChunk { } if (chunk.type === EventType.TOOL_CALL_END && chunk.input === undefined) { - const input = tanstackMetadata(chunk)?.input + const tanstack = tanstackMetadata(chunk) + const input = + tanstack?.input ?? + (typeof tanstack?.args === 'string' + ? parseToolArguments(tanstack.args) + : undefined) if (input !== undefined) { chunk.input = input } diff --git a/packages/ai/src/utilities/tool-call-arguments.ts b/packages/ai/src/utilities/tool-call-arguments.ts new file mode 100644 index 0000000000..d90f8d8bb9 --- /dev/null +++ b/packages/ai/src/utilities/tool-call-arguments.ts @@ -0,0 +1,244 @@ +import { tanstackMetadata } from './merge-metadata' +import type { ToolCallEndEvent } from '../types' +import type { AdapterYieldChunk } from './adapter-yield-chunk' + +type JsonNode = { + start: number + end: number + value: unknown + keyToken?: string + properties?: Map + items?: Array +} + +function sameJsonValue(left: unknown, right: unknown): boolean { + if (Object.is(left, right)) return true + if (Array.isArray(left) && Array.isArray(right)) + return ( + left.length === right.length && + left.every((value, index) => sameJsonValue(value, right[index])) + ) + if ( + left === null || + right === null || + typeof left !== 'object' || + typeof right !== 'object' || + Array.isArray(left) || + Array.isArray(right) + ) + return false + const leftKeys = Object.keys(left) + const rightKeys = Object.keys(right) + return ( + leftKeys.length === rightKeys.length && + leftKeys.every( + (key) => + Object.prototype.hasOwnProperty.call(right, key) && + sameJsonValue(Reflect.get(left, key), Reflect.get(right, key)), + ) + ) +} + +function scanJson(raw: string, parsed: unknown): JsonNode { + let cursor = 0 + const whitespace = () => { + while ( + raw[cursor] === ' ' || + raw[cursor] === '\t' || + raw[cursor] === '\r' || + raw[cursor] === '\n' + ) + cursor++ + } + const stringToken = () => { + const start = cursor++ + while (cursor < raw.length) { + if (raw[cursor] === '\\') cursor += 2 + else if (raw[cursor++] === '"') break + } + return raw.slice(start, cursor) + } + const value = (parsedValue: unknown): JsonNode => { + whitespace() + const start = cursor + if (raw[cursor] === '{') { + cursor++ + whitespace() + const properties = new Map() + while (raw[cursor] !== '}') { + const keyToken = stringToken() + const key: unknown = JSON.parse(keyToken) + whitespace() + cursor++ + const child = + parsedValue !== null && + typeof parsedValue === 'object' && + typeof key === 'string' + ? Reflect.get(parsedValue, key) + : undefined + const childNode = value(child) + if (typeof key === 'string') + properties.set(key, { ...childNode, keyToken }) + whitespace() + if (raw[cursor] !== ',') break + cursor++ + whitespace() + } + cursor++ + return { start, end: cursor, value: parsedValue, properties } + } + if (raw[cursor] === '[') { + cursor++ + whitespace() + const items: Array = [] + while (raw[cursor] !== ']') { + items.push( + value( + Array.isArray(parsedValue) ? parsedValue[items.length] : undefined, + ), + ) + whitespace() + if (raw[cursor] !== ',') break + cursor++ + } + cursor++ + return { start, end: cursor, value: parsedValue, items } + } + if (raw[cursor] === '"') stringToken() + else + while (cursor < raw.length && !/[\s,}\]]/.test(raw[cursor] ?? '')) + cursor++ + return { start, end: cursor, value: parsedValue } + } + return value(parsed) +} + +function serializeValue( + value: unknown, + raw: string, + node: JsonNode | undefined, + seen: Set, +): string | undefined { + if (node && sameJsonValue(node.value, value)) + return raw.slice(node.start, node.end) + if (typeof value === 'number') { + if (Object.is(value, -0)) return '-0' + if (value === Infinity) return '1e999' + if (value === -Infinity) return '-1e999' + return JSON.stringify(value) + } + if (value === null || typeof value !== 'object') return JSON.stringify(value) + if (seen.has(value)) + throw new TypeError('Cannot serialize circular tool input') + seen.add(value) + try { + if (Array.isArray(value)) + return ( + '[' + + Array.from( + value, + (item, index) => + serializeValue(item, raw, node?.items?.[index], seen) ?? 'null', + ).join(',') + + ']' + ) + const properties: Array = [] + for (const key of Object.keys(value)) { + const serialized = serializeValue( + Reflect.get(value, key), + raw, + node?.properties?.get(key), + seen, + ) + if (serialized !== undefined) + properties.push( + (node?.properties?.get(key)?.keyToken ?? JSON.stringify(key)) + + ':' + + serialized, + ) + } + return '{' + properties.join(',') + '}' + } finally { + seen.delete(value) + } +} + +/** + * Keep `raw` when it is the same JSON value as `input`. When they differ, + * write `input` again and keep the raw tokens of the values that did not + * change. Empty or missing raw arguments use the input. + */ +export function reconcileToolCallArguments( + raw: string | undefined, + input: unknown, +): string { + if (input === undefined) return raw ?? '' + if (raw === undefined) + return serializeValue(input, '', undefined, new Set()) ?? '' + let parsed: unknown + try { + parsed = JSON.parse(raw) + } catch { + // Empty raw arguments are no arguments: the input is the value to keep. + if (raw.trim() === '') + return serializeValue(input, '', undefined, new Set()) ?? raw + return raw + } + if (sameJsonValue(parsed, input)) return raw + return serializeValue(input, raw, scanJson(raw, parsed), new Set()) ?? raw +} + +/** + * The input of a tool call from its raw arguments: `{}` for an empty string, + * `undefined` when the arguments are not complete JSON (for example cut off) + * or are `null`. + */ +export function parseToolArguments(args: string): unknown { + if (args.trim() === '') return {} + try { + const parsed: unknown = JSON.parse(args) + return parsed ?? undefined + } catch { + return undefined + } +} + +/** + * The final arguments and input of a tool call at TOOL_CALL_END. `streamed` + * is the text of its TOOL_CALL_ARGS deltas. The END's raw `args` (the + * provider's own string) win over it, and stay as they are when they are the + * same JSON value as the END's `input`. Malformed arguments also stay, so the + * tool reports the error and the model can repair them. `input` is always + * the parse of `args`. + */ +export function finishToolCallArguments( + streamed: string, + end: ToolCallEndEvent, +): { args: string; input: unknown } { + const metadata = tanstackMetadata(end) + // Wire data: only a string is raw arguments. + const endArgs: unknown = (end as AdapterYieldChunk).args ?? metadata?.args + const rawArgs = typeof endArgs === 'string' ? endArgs : undefined + const endInput = end.input ?? metadata?.input + // Only an object or array is tool input. Anything else (`null`, a + // function) keeps the raw or streamed arguments. + const input = + endInput !== null && typeof endInput === 'object' ? endInput : undefined + let raw: string | undefined = rawArgs ?? streamed + // Streamed text that is not JSON gives way to the input. + if (rawArgs === undefined && input !== undefined) { + try { + JSON.parse(raw) + } catch { + raw = undefined + } + } + let args: string + try { + args = reconcileToolCallArguments(raw, input) + } catch { + // An input with no JSON form (a cycle, a bigint): keep the raw arguments. + args = raw ?? streamed + } + return { args, input: parseToolArguments(args) } +} diff --git a/packages/ai/tests/chat-client-tool-wait.test.ts b/packages/ai/tests/chat-client-tool-wait.test.ts index a49217ae20..24ee54158d 100644 --- a/packages/ai/tests/chat-client-tool-wait.test.ts +++ b/packages/ai/tests/chat-client-tool-wait.test.ts @@ -214,7 +214,11 @@ describe('client-tool wait lifecycle', () => { }), ) - expect(fakeTracer.spans).toHaveLength(2) + // The client runs the tool, so the server records no execute_tool span. + expect(fakeTracer.spans.map((span) => span.name)).toEqual([ + 'chat test-model', + 'chat test-model #0', + ]) expect(fakeTracer.spans.every((span) => span.ended)).toBe(true) expect( fakeMeter.records.filter( diff --git a/packages/ai/tests/subagent-interrupts.test.ts b/packages/ai/tests/subagent-interrupts.test.ts index e7c1aabee6..28aa6a5a59 100644 --- a/packages/ai/tests/subagent-interrupts.test.ts +++ b/packages/ai/tests/subagent-interrupts.test.ts @@ -2,6 +2,8 @@ import { describe, expect, it, vi } from 'vitest' import { z } from 'zod' import { chat } from '../src/activities/chat' import { defineAgent } from '../src/activities/chat/agents/define-agent' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { defineInterrupt } from '../src/interrupt-definition' import { StreamProcessor } from '../src/activities/chat/stream/processor' import { uiMessagesToWire } from '../src/utilities/ag-ui-wire' import { fromSpecTokenUsage } from '../src/utilities/ag-ui-usage' @@ -493,6 +495,50 @@ describe('subagent interrupts', () => { expect(second.at(-1)).not.toHaveProperty('outcome.type', 'interrupt') }) + it('keeps the child interrupt next to an afterTools boundary interrupt', async () => { + const { agent } = cleanerAgent() + const review = defineInterrupt({ + id: 'after-review', + responseSchema: z.boolean(), + }) + const chunks = await collectChunks( + chat({ + adapter: toolParent('{}'), + threadId: 't', + runId: 'run-1', + messages: [user], + subagents: { agents: [agent] }, + interrupts: [review], + middleware: [ + defineChatMiddleware({ + name: 'after-review', + onInterruptBoundary(ctx) { + if (ctx.phase !== 'afterTools') return undefined + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review', + }), + ], + } + }, + }), + ], + }) as AsyncIterable, + ) + expect(chunks.at(-1)).toMatchObject({ + type: EventType.RUN_FINISHED, + outcome: { + type: 'interrupt', + interrupts: expect.arrayContaining([ + expect.objectContaining({ id: 'approval_call_c' }), + ]), + }, + }) + }) + it('gives a resumed child the same input the model wrote', async () => { const { run, execute } = cleanerChild() const inputs: Array<{ folder: string }> = [] diff --git a/packages/ai/tests/subagent-tool-input.test.ts b/packages/ai/tests/subagent-tool-input.test.ts index ac2f30bacb..c60c0899b1 100644 --- a/packages/ai/tests/subagent-tool-input.test.ts +++ b/packages/ai/tests/subagent-tool-input.test.ts @@ -220,7 +220,7 @@ describe('agent tool input', () => { }, }) - it('checks the agent fields before middleware, with sessionId in the args', async () => { + it('checks the agent fields after middleware, with sessionId in the args', async () => { const deep = defineAgent({ name: 'deep', description: 'Takes a depth', @@ -250,11 +250,34 @@ describe('agent tool input', () => { }) as AsyncIterable, ) - // The schema default is in place when middleware runs. - expect(seen).toEqual([{ topic: 'tides', depth: 2 }]) + // Middleware sees the model's input. The final check adds the default. + expect(seen).toEqual([{ topic: 'tides' }]) expect(runs[0]?.input).toEqual({ topic: 'tides', depth: 2 }) }) + it('checks a middleware rewrite of the agent fields before the child starts', async () => { + const rewrite: ChatMiddleware = { + name: 'rewrite', + onBeforeToolCall: () => ({ + type: 'transformArgs', + args: { topic: { text: 'tides' } }, + }), + } + const { chunks, result } = await callSubagent( + 'researcher', + { topic: 'tides' }, + { middleware: [rewrite] }, + ) + + expect(result.error).toContain( + 'Input validation failed for tool researcher', + ) + expect(runs).toEqual([]) + expect( + chunks.some((chunk) => chunk.type === EventType.SUBAGENT_STARTED), + ).toBe(false) + }) + it('throws when an inputSchema uses a reserved field that the run uses', () => { const { adapter } = parentCalling('taken', {}) @@ -463,9 +486,10 @@ describe('agent tool input', () => { [ 'a field that fails the agent schema', 'researcher', - { topic: 3 }, + // A number would be repaired into a string. A list cannot be. + { topic: [] }, {}, - 'Input validation failed for tool researcher: Validation failed: Invalid input: expected string, received number', + 'Input validation failed for tool researcher: Validation failed: Invalid input: expected string, received array', ], [ 'an unknown sessionId', @@ -484,9 +508,11 @@ describe('agent tool input', () => { [ 'a background that is not a boolean', 'writer', - { background: 'true' }, + // 'true' would be repaired into true. 'yes' cannot be. + { background: 'yes' }, { binding: host }, - 'prompt and sessionId must be strings, and background a boolean.', + // The JSON Schema of the tool checks the reserved fields too. + 'Input validation failed for tool writer: #/background: Instance type "string" is invalid. Expected "boolean".', ], [ 'background with a sessionId', @@ -501,7 +527,9 @@ describe('agent tool input', () => { async (_label, agent, args, options, error) => { const { chunks, result } = await callSubagent(agent, args, options) - expect(result).toEqual({ error }) + // A schema error also shows the received arguments after the message. + expect(Object.keys(result)).toEqual(['error']) + expect(result.error.startsWith(error)).toBe(true) expect(runs).toEqual([]) // The parent run goes on after the tool error. expect(chunks.some((chunk) => chunk.type === EventType.RUN_ERROR)).toBe( diff --git a/packages/ai/tests/tool-call-arguments.test.ts b/packages/ai/tests/tool-call-arguments.test.ts new file mode 100644 index 0000000000..f0ed922b14 --- /dev/null +++ b/packages/ai/tests/tool-call-arguments.test.ts @@ -0,0 +1,223 @@ +import { describe, expect, it } from 'vitest' +import { chat } from '../src/activities/chat/index' +import { EventType } from '../src/types' +import type { ToolCallEndEvent } from '../src/types' +import { tanstackMetadata } from '../src/utilities/merge-metadata' +import { + finishToolCallArguments, + reconcileToolCallArguments, +} from '../src/utilities/tool-call-arguments' +import { collectChunks, createMockAdapter, ev } from './test-utils' + +describe('lexical tool argument reconciliation', () => { + it.each([ + ' 1e999 ', + '-0', + '9007199254740993', + '1.20e+3', + '"\\u0061"', + '{"a":1e999,"b":[-0,9007199254740993,"\\u0061"]}', + ])('keeps unchanged raw JSON exactly: %s', (raw) => { + expect(reconcileToolCallArguments(raw, JSON.parse(raw))).toBe(raw) + }) + it('reuses unchanged nested values after removing an optional null', () => { + const raw = + '{"removed":null,"nested":{"number":1e999,"text":"\\u0061"},"array":[-0,9007199254740993]}' + expect( + reconcileToolCallArguments(raw, { + nested: { number: Infinity, text: 'a' }, + array: [-0, 9007199254740992], + }), + ).toBe( + '{"nested":{"number":1e999,"text":"\\u0061"},"array":[-0,9007199254740993]}', + ) + }) + it('retains unchanged array positions while correcting another position', () => { + expect(reconcileToolCallArguments('[1e999,2,-0]', [Infinity, 3, -0])).toBe( + '[1e999,3,-0]', + ) + }) + it('uses valid JSON spelling for corrected numbers', () => { + expect( + reconcileToolCallArguments('{}', { + positive: Infinity, + negative: -Infinity, + zero: -0, + }), + ).toBe('{"positive":1e999,"negative":-1e999,"zero":-0}') + }) + it('does not repair malformed raw JSON with a parsed projection', () => { + expect(reconcileToolCallArguments('{"broken":', {})).toBe('{"broken":') + }) + it('serializes END-only input when no raw arguments exist', () => { + expect(reconcileToolCallArguments(undefined, null)).toBe('null') + expect(reconcileToolCallArguments(undefined, [1, true])).toBe('[1,true]') + }) + it('leaves raw arguments alone when input is undefined', () => { + expect(reconcileToolCallArguments('1e999', undefined)).toBe('1e999') + }) +}) + +describe('raw JSON key spelling', () => { + it('keeps an escaped key when removing an unrelated null', () => { + expect( + reconcileToolCallArguments('{"\\u0061":9007199254740993,"drop":null}', { + a: 9007199254740992, + }), + ).toBe('{"\\u0061":9007199254740993}') + }) + it('uses the last duplicate key value and its original spelling', () => { + expect( + reconcileToolCallArguments('{"a":1,"\\u0061":1e999,"drop":null}', { + a: Infinity, + }), + ).toBe('{"\\u0061":1e999}') + }) + it('serializes sparse custom array positions as null slots', () => { + const values: Array = new Array(2) + values[1] = -0 + expect(reconcileToolCallArguments(undefined, values)).toBe('[null,-0]') + }) +}) + +describe('finishToolCallArguments', () => { + const end = ( + fields: Partial & { args?: string }, + ): ToolCallEndEvent => ({ + type: EventType.TOOL_CALL_END, + toolCallId: 'call-1', + timestamp: 0, + ...fields, + }) + + it('keeps raw args when the input is null', () => { + expect( + finishToolCallArguments('', end({ args: '{"q":"x"}', input: null })), + ).toEqual({ args: '{"q":"x"}', input: { q: 'x' } }) + }) + + it('keeps malformed raw args with no input', () => { + expect(finishToolCallArguments('', end({ args: '{"q":' }))).toEqual({ + args: '{"q":', + input: undefined, + }) + }) + + it('keeps the streamed bytes when they match the input', () => { + expect( + finishToolCallArguments('{"n": 1.50}', end({ input: { n: 1.5 } })), + ).toEqual({ args: '{"n": 1.50}', input: { n: 1.5 } }) + }) + + it('ignores args that are not a string', () => { + expect( + finishToolCallArguments( + '{"a":1}', + end({ metadata: { tanstack: { args: { a: 1 } } } } as never), + ), + ).toEqual({ args: '{"a":1}', input: { a: 1 } }) + }) + + it('reads args and input from metadata.tanstack', () => { + expect( + finishToolCallArguments( + '', + end({ + metadata: { + tanstack: { args: '{"a": 1, "b": null}', input: { a: 1 } }, + }, + }), + ), + ).toEqual({ args: '{"a":1}', input: { a: 1 } }) + }) +}) + +describe('raw tool arguments through chat()', () => { + async function run(raw: string, input?: unknown) { + const executed: Array = [] + const { adapter, calls } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call-1', 'check'), + ev.toolArgs('call-1', raw), + { + type: EventType.TOOL_CALL_END, + toolCallId: 'call-1', + timestamp: Date.now(), + args: raw, + ...(input !== undefined ? { input } : {}), + }, + ev.runFinished('tool_calls'), + ], + [ev.runStarted('run-2'), ev.runFinished('stop', 'run-2')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'check' }], + tools: [ + { + name: 'check', + description: 'Check', + execute: (value: unknown) => { + executed.push(value) + return 'ok' + }, + }, + ], + }), + ) + return { chunks, calls, executed } + } + + function replayedArguments( + calls: ReturnType['calls'], + ) { + return calls[1]?.messages + .flatMap((message) => message.toolCalls ?? []) + .map((call) => call.function.arguments) + } + + it('replays the exact provider bytes', async () => { + const raw = '{"count": 1.50, "id": 9007199254740993}' + const { calls, executed } = await run(raw, JSON.parse(raw)) + expect(replayedArguments(calls)).toEqual([raw]) + expect(executed).toEqual([{ count: 1.5, id: 9007199254740992 }]) + }) + + it('sends the raw arguments to the client in TOOL_CALL_END metadata', async () => { + const raw = '{"count": 1.50}' + const { chunks } = await run(raw, JSON.parse(raw)) + const end = chunks.find((chunk) => chunk.type === EventType.TOOL_CALL_END) + const metadata = end ? tanstackMetadata(end) : undefined + expect(metadata && 'args' in metadata ? metadata.args : undefined).toBe(raw) + // `input` is not sent twice; in process it is parsed back from `args`. + expect(metadata && 'input' in metadata).toBe(false) + expect(end && 'input' in end ? end.input : undefined).toEqual({ + count: 1.5, + }) + }) + + it('keeps an in-process input that JSON spells as -0 or 1e999', async () => { + const raw = '{"zero": -0, "big": 1e999}' + const { chunks } = await run(raw, JSON.parse(raw)) + const end = chunks.find((chunk) => chunk.type === EventType.TOOL_CALL_END) + expect(end && 'input' in end ? end.input : undefined).toEqual({ + zero: -0, + big: Infinity, + }) + }) + + it('reports malformed raw arguments instead of running with {}', async () => { + const { calls, executed } = await run('{"count":') + expect(executed).toEqual([]) + const toolMessage = calls[1]?.messages.find( + (message) => message.role === 'tool', + ) + expect(JSON.stringify(toolMessage)).toContain( + 'Failed to parse tool arguments as JSON', + ) + }) +}) diff --git a/packages/ai/tests/tool-input-validation.test.ts b/packages/ai/tests/tool-input-validation.test.ts new file mode 100644 index 0000000000..14b5dec60a --- /dev/null +++ b/packages/ai/tests/tool-input-validation.test.ts @@ -0,0 +1,1527 @@ +import { describe, expect, it, vi } from 'vitest' +import { z } from 'zod' +import { chat } from '../src/activities/chat/index' +import { defineChatMiddleware } from '../src/activities/chat/middleware/define' +import { defineInterrupt } from '../src/interrupt-definition' +import { executeToolCalls } from '../src/activities/chat/tools/tool-calls' +import { validateToolInput } from '../src/activities/chat/tools/input-validation' +import { otelMiddleware } from '../src/middlewares/otel' +import { EventType } from '../src/types' +import { collectChunks, createMockAdapter, ev } from './test-utils' +import { createFakeMeter, createFakeTracer } from './middlewares/fake-otel' +import type { StandardSchemaV1 } from '@standard-schema/spec' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' +import type { ToolExecutionMiddlewareHooks } from '../src/activities/chat/tools/tool-calls' +import type { ToolApprovalResolution } from '../src/interrupts' +import type { + AnyTool, + JSONSchema, + ModelMessage, + SchemaInput, + StreamChunk, + TextOptions, + Tool, + ToolCall, + ToolInputRepairContext, +} from '../src/types' + +const schema = z.object({ count: z.number().int().min(1) }) + +/** Checks the schema alone: the built-in repair would fix some inputs. */ +const noRepair = { repairToolInput: false } as const + +async function runTool( + rawArguments: string, + inputSchema: SchemaInput | boolean, + options: { + client?: boolean + needsApproval?: boolean + approvals?: Map + middleware?: ToolExecutionMiddlewareHooks + clientResults?: Map + repairToolInput?: TextOptions['repairToolInput'] + } = {}, +) { + const executed: Array = [] + const tool: AnyTool = { + name: 'check', + description: 'Check input', + inputSchema, + needsApproval: options.needsApproval, + ...(options.client + ? {} + : { + execute: (input: unknown) => { + executed.push(input) + return input + }, + }), + } + const call: ToolCall = { + id: 'call-1', + type: 'function', + function: { name: 'check', arguments: rawArguments }, + } + const generator = executeToolCalls( + [call], + [tool], + options.approvals, + options.clientResults, + undefined, + options.middleware, + undefined, + undefined, + undefined, + undefined, + options.repairToolInput, + ) + while (true) { + const next = await generator.next() + if (next.done) return { ...next.value, executed } + } +} + +describe('final tool input validation', () => { + it('awaits an async Standard Schema validator', async () => { + const asyncSchema: StandardSchemaV1 = { + '~standard': { + version: 1, + vendor: 'test', + validate: async (value) => + typeof value === 'object' && + value !== null && + 'count' in value && + typeof value.count === 'number' + ? { value: { count: value.count * 10 } } + : { issues: [{ message: 'count must be a number' }] }, + }, + } + const ok = await runTool('{"count":2}', asyncSchema) + expect(ok.executed).toEqual([{ count: 20 }]) + const bad = await runTool('{"count":"2"}', asyncSchema) + expect(bad.executed).toEqual([]) + expect(bad.results[0]?.state).toBe('output-error') + expect(bad.results[0]?.result).toEqual({ + error: expect.stringContaining('count must be a number'), + }) + }) + + it('includes the received arguments in the error', async () => { + const result = await runTool('{"count":0}', schema) + expect(result.results[0]?.result).toEqual({ + error: expect.stringMatching( + /^Input validation failed for tool check: [\s\S]*\nReceived arguments: \{"count":0\}$/, + ), + }) + expect(result.executed).toEqual([]) + }) + + it('echoes only the model arguments, not values middleware added', async () => { + const result = await runTool('{"count":0}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 0, token: 'secret-token' }, + }), + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(JSON.stringify(result.results[0]?.result)).not.toContain( + 'secret-token', + ) + // The model learns that its own arguments were changed. + expect(JSON.stringify(result.results[0]?.result)).toContain( + 'An approval edit or middleware changed these arguments', + ) + }) + + it('keeps empty and null arguments as {}', async () => { + for (const raw of ['', ' ', 'null']) { + const result = await runTool(raw, z.object({})) + expect(result.executed).toEqual([{}]) + } + }) + + it.each([false, true])( + 'rejects a middleware replacement before server or client execution: client=%s', + async (client) => { + const result = await runTool('{"count":2}', schema, { + client, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + expect(result.needsClientExecution).toEqual([]) + }, + ) + + it('rejects a middleware edit made in place', async () => { + const result = await runTool('{"count":2}', schema, { + middleware: { + onBeforeToolCall: async (_call, _tool, input) => { + if (typeof input === 'object' && input !== null) + Reflect.set(input, 'count', 'wrong') + }, + }, + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + }) + + it('runs onAfterToolCall with the validation error', async () => { + const after: Array = [] + await runTool('{"count":2}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + onAfterToolCall: async (info) => { + after.push(info.ok) + }, + }, + }) + expect(after).toEqual([false]) + }) + + it('lets middleware repair invalid input', async () => { + const result = await runTool('{"count":"wrong"}', schema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 4 }, + }), + }, + }) + expect(result.executed).toEqual([{ count: 4 }]) + }) + + it('shows middleware the raw input and runs a transform once', async () => { + const seen: Array = [] + let transforms = 0 + const inputSchema = z.object({ + count: z.string().transform((value) => { + transforms++ + return Number(value) + }), + }) + const result = await runTool('{"count":"2"}', inputSchema, { + middleware: { + onBeforeToolCall: async (_call, _tool, input) => { + seen.push(structuredClone(input)) + }, + }, + }) + expect(seen).toEqual([{ count: '2' }]) + expect(result.executed).toEqual([{ count: 2 }]) + expect(transforms).toBe(1) + }) + + it('sends the checked middleware input to the client', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 5 }, + }), + }, + }) + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: { count: 5 } }, + ]) + }) + + it('checks an approval edit for a client tool', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: 'wrong' } }], + ]), + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.needsClientExecution).toEqual([]) + }) + + it('keeps a returned client result when middleware repaired the input', async () => { + let hooks = 0 + const result = await runTool('{"count":"wrong"}', schema, { + client: true, + clientResults: new Map([['call-1', 'client result']]), + middleware: { + onBeforeToolCall: async () => { + hooks++ + return { type: 'transformArgs', args: { count: 4 } } + }, + }, + }) + expect(hooks).toBe(0) + expect(result.results[0]?.state).not.toBe('output-error') + expect(result.results[0]?.result).toBe('client result') + }) + + it('rejects an invalid approval edit for a server tool', async () => { + const result = await runTool('{"count":2}', schema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: 0 } }], + ]), + }) + expect(result.results[0]?.state).toBe('output-error') + expect(result.executed).toEqual([]) + }) + + it.each([false, true])( + 'keeps the transformed approval edit: client=%s', + async (client) => { + const result = await runTool('"2"', z.string().transform(Number), { + client, + needsApproval: true, + approvals: new Map([['call-1', { approved: true, editedArgs: '4' }]]), + }) + if (client) { + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: 4 }, + ]) + } else { + expect(result.executed).toEqual([4]) + } + }, + ) + + it.each([false, true])( + 'runs no middleware and dispatches nothing while approval is pending: client=%s', + async (client) => { + let middlewareCalls = 0 + const result = await runTool('{"count":2}', schema, { + client, + needsApproval: true, + middleware: { + onBeforeToolCall: async () => { + middlewareCalls++ + }, + }, + }) + expect(result.needsApproval).toEqual([ + { + toolCallId: 'call-1', + toolName: 'check', + input: { count: 2 }, + approvalId: 'approval_call-1', + }, + ]) + expect(result.needsClientExecution).toEqual([]) + expect(middlewareCalls).toBe(0) + expect(result.executed).toEqual([]) + }, + ) + + it.each([false, true])( + 'runs no middleware for a denied tool: client=%s', + async (client) => { + let middlewareCalls = 0 + const result = await runTool('{"count":2}', schema, { + client, + needsApproval: true, + approvals: new Map([['call-1', false]]), + middleware: { + onBeforeToolCall: async () => { + middlewareCalls++ + }, + }, + }) + expect(result.results[0]?.outcome).toBe('denied') + expect(result.needsClientExecution).toEqual([]) + expect(middlewareCalls).toBe(0) + expect(result.executed).toEqual([]) + }, + ) +}) + +describe('a batch with an approval', () => { + const call = (id: string, name: string, args: string): ToolCall => ({ + id, + type: 'function', + function: { name, arguments: args }, + }) + + async function runBatch( + calls: Array, + tools: Array, + approvals?: Map, + ) { + const generator = executeToolCalls(calls, tools, approvals) + while (true) { + const next = await generator.next() + if (next.done) return next.value + } + } + + it('runs the held calls when the only approval preview fails its check', async () => { + const execute = vi.fn(() => 'ran') + const result = await runBatch( + [call('a', 'gated', '{"count":0}'), call('b', 'plain', '{}')], + [ + { + name: 'gated', + description: 'Gated', + inputSchema: schema, + needsApproval: true, + execute: vi.fn(), + }, + { name: 'plain', description: 'Plain', execute }, + ], + ) + expect(result.needsApproval).toEqual([]) + expect(execute).toHaveBeenCalledTimes(1) + expect(result.results.map((r) => [r.toolCallId, r.state])).toEqual([ + ['a', 'output-error'], + ['b', undefined], + ]) + }) + + it('sends an approved client call while another call waits for approval', async () => { + const result = await runBatch( + [call('a', 'clientGated', '{}'), call('b', 'serverGated', '{}')], + [ + { name: 'clientGated', description: 'Client', needsApproval: true }, + { + name: 'serverGated', + description: 'Server', + needsApproval: true, + execute: vi.fn(), + }, + ], + new Map([['a', true]]), + ) + expect(result.needsClientExecution.map((r) => r.toolCallId)).toEqual(['a']) + expect(result.needsApproval.map((r) => r.toolCallId)).toEqual(['b']) + }) +}) + +describe('JSON Schema tool input', () => { + const jsonSchema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer', minimum: 1 } }, + required: ['count'], + additionalProperties: false, + } + + it('runs valid input', async () => { + const result = await runTool('{"count":2}', jsonSchema) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it.each([ + '{}', + '{"count":"wrong"}', + '{"count":0}', + '{"count":2,"extra":true}', + ])('rejects invalid input without executing: %s', async (raw) => { + const result = await runTool(raw, jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Input validation failed for tool check'), + }) + }) + + it('names the failed location and shows the received arguments', async () => { + const result = await runTool('{"count":0}', jsonSchema) + expect(result.results[0]?.result).toEqual({ + error: expect.stringMatching(/#\/count: .*\n.*Received arguments:/s), + }) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Received arguments: {"count":0}'), + }) + }) + + it('checks local references and conditional fields', async () => { + const schema: JSONSchema = { + type: 'object', + $defs: { count: { type: 'integer' } }, + properties: { + mode: { type: 'string' }, + count: { $ref: '#/$defs/count' }, + }, + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { required: ['count'] }, + } + expect( + (await runTool('{"mode":"count","count":3}', schema)).executed, + ).toEqual([{ mode: 'count', count: 3 }]) + const missing = await runTool('{"mode":"count"}', schema) + expect(missing.executed).toEqual([]) + // The "then" wrapper error is left out. The nested error says why. + expect(missing.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'check: #: Instance does not have required property "count".\n', + ), + }) + expect((await runTool('{"count":"x"}', schema)).executed).toEqual([]) + }) + + it('uses the draft named by $schema', async () => { + // Draft 4 spells an exclusive bound as a boolean next to `maximum`. + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema: 'http://json-schema.org/draft-04/schema#', + type: 'object', + properties: { + count: { type: 'number', maximum: 5, exclusiveMaximum: true }, + }, + }), + ) + expect((await runTool('{"count":4}', schema)).executed).toEqual([ + { count: 4 }, + ]) + expect((await runTool('{"count":5}', schema)).executed).toEqual([]) + }) + + it('runs empty and null arguments as {} when nothing is required', async () => { + for (const raw of ['', 'null']) { + const result = await runTool(raw, { type: 'object', properties: {} }) + expect(result.executed).toEqual([{}]) + } + }) + + it('rejects empty arguments when a field is required', async () => { + const result = await runTool('', jsonSchema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.state).toBe('output-error') + }) + + it('checks a middleware rewrite and an approval edit', async () => { + const rewritten = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + }) + expect(rewritten.executed).toEqual([]) + const edited = await runTool('{"count":2}', jsonSchema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: -1 } }], + ]), + }) + expect(edited.executed).toEqual([]) + expect(edited.results[0]?.state).toBe('output-error') + }) + + it('turns an unusable schema into a tool error', async () => { + const result = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $ref: 'https://example.com/count.json' } }, + }) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Unresolved $ref'), + }) + expect(JSON.stringify(result.results[0]?.result)).not.toContain( + 'Known schemas', + ) + }) + + it('omits the validator base URI from a schema error', async () => { + const missing = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $ref: '#/$defs/missing' } }, + }) + expect(missing.results[0]?.result).toEqual({ + error: expect.stringContaining('Unresolved $ref "#/$defs/missing"'), + }) + const missingError = JSON.stringify(missing.results[0]?.result) + expect(missingError).not.toContain('Absolute URI') + expect(missingError).not.toContain('github.com') + expect(missingError).not.toContain('Known schemas') + + const duplicate = await runTool('{"count":2}', { + type: 'object', + properties: { count: { $id: 'https://github.com/cfworker' } }, + }) + const duplicateError = JSON.stringify(duplicate.results[0]?.result) + expect(duplicate.executed).toEqual([]) + expect(duplicateError).toContain('Duplicate schema URI') + expect(duplicateError).not.toContain('github.com') + }) + + it.each([false, true])( + 'a boolean schema %s runs only when the schema is true', + async (inputSchema) => { + const result = await runTool('{"count":2}', inputSchema) + expect(result.executed).toEqual(inputSchema ? [{ count: 2 }] : []) + if (!inputSchema) expect(result.results[0]?.state).toBe('output-error') + }, + ) + + it.each([ + undefined, + 'http://json-schema.org/schema#', + 'http://json-schema.org/draft-07/schema#', + 'https://json-schema.org/draft/2019-09/schema', + ])('reads a boolean exclusive bound with $schema %s', async ($schema) => { + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema, + type: 'object', + properties: { + count: { type: 'number', maximum: 5, exclusiveMaximum: true }, + }, + }), + ) + expect((await runTool('{"count":4}', schema)).executed).toEqual([ + { count: 4 }, + ]) + expect((await runTool('{"count":5}', schema)).executed).toEqual([]) + }) + + it('keeps a number bound next to a boolean bound', async () => { + const schema: JSONSchema = JSON.parse( + JSON.stringify({ + $schema: 'https://json-schema.org/draft/2020-12/schema', + type: 'object', + properties: { + low: { type: 'number', minimum: 0, exclusiveMinimum: false }, + high: { type: 'number', exclusiveMinimum: 5 }, + }, + }), + ) + expect((await runTool('{"low":0,"high":6}', schema)).executed).toEqual([ + { low: 0, high: 6 }, + ]) + expect((await runTool('{"high":5}', schema)).executed).toEqual([]) + }) + + it('lets an OpenAPI 3.0 nullable reference be null', async () => { + const schema: JSONSchema = { + type: 'object', + definitions: { User: { type: 'object', required: ['id'] } }, + properties: { + owner: { nullable: true, allOf: [{ $ref: '#/definitions/User' }] }, + editor: { nullable: true, oneOf: [{ $ref: '#/definitions/User' }] }, + }, + } + expect( + (await runTool('{"owner":null,"editor":null}', schema)).executed, + ).toEqual([{ owner: null, editor: null }]) + expect( + (await runTool('{"owner":{"id":1},"editor":{"id":2}}', schema)).executed, + ).toEqual([{ owner: { id: 1 }, editor: { id: 2 } }]) + const wrong = await runTool('{"owner":{}}', schema) + expect(wrong.executed).toEqual([]) + expect(wrong.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'check: #/owner: Instance does not have required property "id".\n', + ), + }) + }) + + it('moves a percent-encoded $ref and a ref into a nullable root', async () => { + const schema: JSONSchema = { + $id: 'https://example.com/tool', + nullable: true, + allOf: [{ $ref: '#/definitions/Input' }], + definitions: { + Input: { + type: 'object', + properties: { + 'my field': { + nullable: true, + anyOf: [{ type: 'integer' }, { type: 'string' }], + }, + 'a/b': { + $ref: '#/definitions/Input/properties/my%20field/anyOf/0', + }, + '50%': { + nullable: true, + allOf: [{ type: 'object', required: ['id'] }], + }, + copy: { $ref: '#/definitions/Input/properties/50%25/allOf/0' }, + }, + }, + }, + } + expect( + (await runTool('{"my field":null,"a/b":3}', schema)).executed, + ).toEqual([{ 'my field': null, 'a/b': 3 }]) + expect((await runTool('{"a/b":null}', schema, noRepair)).executed).toEqual( + [], + ) + expect((await runTool('{"a/b":"3"}', schema, noRepair)).executed).toEqual( + [], + ) + expect( + (await runTool('{"50%":null,"copy":{"id":1}}', schema)).executed, + ).toEqual([{ '50%': null, copy: { id: 1 } }]) + expect((await runTool('{"copy":null}', schema, noRepair)).executed).toEqual( + [], + ) + }) + + it('keeps a $ref that points into a nullable schema', async () => { + // zod-to-json-schema with `target: 'openApi3'` reuses a schema by path. + const address = { + type: 'object', + properties: { city: { type: 'string' } }, + required: ['city'], + } + const schema: JSONSchema = { + type: 'object', + properties: { + home: { anyOf: [address, { type: 'string' }], nullable: true }, + work: { $ref: '#/properties/home/anyOf/0' }, + office: { nullable: true, $ref: '#/properties/home/anyOf/0' }, + }, + } + expect( + (await runTool('{"home":null,"work":{"city":"x"},"office":null}', schema)) + .executed, + ).toEqual([{ home: null, work: { city: 'x' }, office: null }]) + expect((await runTool('{"work":null}', schema, noRepair)).executed).toEqual( + [], + ) + expect((await runTool('{"office":{}}', schema, noRepair)).executed).toEqual( + [], + ) + expect((await runTool('{"home":1}', schema, noRepair)).executed).toEqual([]) + }) + + it('moves a $ref through nested nullable schemas', async () => { + const schema: JSONSchema = { + type: 'object', + properties: { + outer: { + nullable: true, + allOf: [ + { + type: 'object', + properties: { + inner: { + nullable: true, + oneOf: [{ type: 'integer' }, { type: 'string' }], + }, + }, + }, + ], + }, + copy: { $ref: '#/properties/outer/allOf/0/properties/inner/oneOf/0' }, + }, + } + expect( + (await runTool('{"outer":{"inner":null},"copy":3}', schema)).executed, + ).toEqual([{ outer: { inner: null }, copy: 3 }]) + expect((await runTool('{"copy":"3"}', schema, noRepair)).executed).toEqual( + [], + ) + expect( + (await runTool('{"outer":{"inner":true}}', schema, noRepair)).executed, + ).toEqual([]) + }) + + it('keeps draft-07 $ref siblings ignored in a nullable schema', async () => { + const schema: JSONSchema = { + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'object', + definitions: { P: { type: 'object' } }, + properties: { + p: { $ref: '#/definitions/P', nullable: true, required: ['name'] }, + }, + } + expect((await runTool('{"p":{"id":1}}', schema)).executed).toEqual([ + { p: { id: 1 } }, + ]) + expect((await runTool('{"p":null}', schema)).executed).toEqual([ + { p: null }, + ]) + }) + + it('lets an OpenAPI 3.0 nullable field be null', async () => { + const schema: JSONSchema = { + type: 'object', + properties: { + name: { type: 'string', nullable: true }, + kind: { type: 'string', enum: ['a', 'b'], nullable: true }, + size: { type: 'integer' }, + }, + } + expect( + (await runTool('{"name":null,"kind":null}', schema)).executed, + ).toEqual([{ name: null, kind: null }]) + expect((await runTool('{"kind":"c"}', schema)).executed).toEqual([]) + expect((await runTool('{"size":null}', schema, noRepair)).executed).toEqual( + [], + ) + }) + + it('repairs a value inside an OpenAPI 3.0 nullable schema', async () => { + const schema: JSONSchema = { + type: 'object', + properties: { + count: { nullable: true, allOf: [{ type: 'integer' }] }, + size: { type: 'integer', nullable: true }, + }, + } + expect( + (await runTool('{"count":"3","size":"4"}', schema)).executed, + ).toEqual([{ count: 3, size: 4 }]) + expect( + (await runTool('{"count":null,"size":null}', schema)).executed, + ).toEqual([{ count: null, size: null }]) + }) + + it('treats format as an annotation', async () => { + const result = await runTool('{"path":"/tmp/report.txt"}', { + type: 'object', + properties: { path: { type: 'string', format: 'uri' } }, + required: ['path'], + }) + expect(result.executed).toEqual([{ path: '/tmp/report.txt' }]) + }) + + it('keeps a format key inside const and enum values', async () => { + const input = { + file: { format: 'pdf' }, + kind: { format: 'csv' }, + default: '/tmp/report.txt', + } + const result = await runTool(JSON.stringify(input), { + type: 'object', + properties: { + file: { const: { format: 'pdf' } }, + kind: { enum: [{ format: 'csv' }] }, + // A property named like a data keyword is still a schema. + default: { type: 'string', format: 'uri' }, + }, + }) + expect(result.executed).toEqual([input]) + }) + + it('lists at most 10 errors', async () => { + const result = await runTool( + JSON.stringify({ xs: Array.from({ length: 12 }, () => 'x') }), + { + type: 'object', + properties: { xs: { type: 'array', items: { type: 'integer' } } }, + }, + ) + const error = JSON.stringify(result.results[0]?.result) + expect(error).toContain('#/xs/9: ') + expect(error).not.toContain('#/xs/10: ') + expect(error).toContain(', and 2 more') + }) + + it('checks a middleware rewrite to undefined as null', async () => { + const result = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: undefined, + }), + }, + }) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining( + 'check: #: Instance type "null" is invalid. Expected "object".', + ), + }) + }) + + it('treats a key set to undefined as a missing key', async () => { + const result = await runTool('{"count":2}', jsonSchema, { + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: 2, note: undefined }, + }), + }, + }) + expect(result.executed).toEqual([{ count: 2, note: undefined }]) + }) + + it('does not change the tool schema', async () => { + const schema: JSONSchema = { + type: 'object', + $defs: { count: { type: 'integer' } }, + properties: { + count: { $ref: '#/$defs/count' }, + name: { type: 'string', format: 'uri', nullable: true }, + kind: { enum: ['a'], nullable: true }, + owner: { nullable: true, allOf: [{ $ref: '#/$defs/count' }] }, + }, + } + const before = JSON.stringify(schema) + await runTool('{"count":2}', schema) + expect(JSON.stringify(schema)).toBe(before) + expect(Object.getOwnPropertyNames(schema.properties?.count)).toEqual([ + '$ref', + ]) + }) + + it('works when runtime code generation is blocked', async () => { + vi.stubGlobal('Function', function blockedCodeGeneration() { + throw new Error('Code generation is blocked') + }) + try { + const schema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer' } }, + } + expect((await runTool('{"count":2}', schema)).executed).toEqual([ + { count: 2 }, + ]) + expect((await runTool('{"count":"2"}', schema)).executed).toEqual([ + { count: 2 }, + ]) + expect((await runTool('{"count":"x"}', schema)).executed).toEqual([]) + } finally { + vi.unstubAllGlobals() + } + }) +}) + +describe('tool input repair', () => { + const schema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer', minimum: 1 } }, + required: ['count'], + additionalProperties: false, + } + + it('coerces a numeric field before execution', async () => { + const result = await runTool('{"count":"2"}', schema) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('does not change the received arguments', async () => { + const received = { count: '2' } + expect(await validateToolInput(schema, received, 'check')).toEqual({ + count: 2, + }) + expect(received).toEqual({ count: '2' }) + }) + + it('shows the original arguments when the repaired value still fails', async () => { + const result = await runTool('{"count":"0"}', schema) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Received arguments: {"count":"0"}'), + }) + }) + + it('follows local references and drops a null optional field', async () => { + const result = await runTool('{"count":"3","label":null,"omitted":null}', { + type: 'object', + $defs: { count: { type: 'integer', minimum: 1 } }, + properties: { + count: { $ref: '#/$defs/count' }, + label: { type: ['string', 'null'] }, + omitted: { type: 'string' }, + }, + required: ['count'], + }) + expect(result.executed).toEqual([{ count: 3, label: null }]) + }) + + it.each<[unknown, JSONSchema, unknown]>([ + ['12', { anyOf: [{ type: 'integer' }, { type: 'string' }] }, '12'], + ['12', { type: ['integer', 'string'] }, '12'], + ['3', { type: 'integer' }, 3], + [['1', '2'], { type: 'array', items: { type: 'integer' } }, [1, 2]], + [ + ['1', 'x'], + { prefixItems: [{ type: 'integer' }, { type: 'string' }] }, + [1, 'x'], + ], + ['true', { type: 'boolean' }, true], + [null, { type: 'string' }, ''], + ])('repairs %j against %j', async (received, inputSchema, expected) => { + expect(await validateToolInput(inputSchema, received, 'check')).toEqual( + expected, + ) + }) + + it('rejects a primitive that cannot be repaired', async () => { + await expect( + validateToolInput({ type: 'integer' }, 'wrong', 'check'), + ).rejects.toThrow('Input validation failed for tool check') + }) + + it('keeps a valid referenced union branch', async () => { + const result = await runTool('{"count":"12"}', { + type: 'object', + $defs: { numeric: { type: 'integer' }, text: { type: 'string' } }, + properties: { + count: { + anyOf: [{ $ref: '#/$defs/numeric' }, { $ref: '#/$defs/text' }], + }, + }, + required: ['count'], + }) + expect(result.executed).toEqual([{ count: '12' }]) + }) + + it('coerces conditional branch fields', async () => { + const result = await runTool('{"mode":"count","count":"4"}', { + type: 'object', + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { + properties: { count: { type: 'integer' } }, + required: ['count'], + }, + }) + expect(result.executed).toEqual([{ mode: 'count', count: 4 }]) + }) + + it('keeps a null field that a conditional branch requires', async () => { + const result = await runTool('{"mode":"count","count":null}', { + type: 'object', + properties: { mode: { type: 'string' }, count: { type: 'integer' } }, + if: { properties: { mode: { const: 'count' } }, required: ['mode'] }, + then: { required: ['count'] }, + }) + expect(result.executed).toEqual([{ mode: 'count', count: 0 }]) + }) + + it('keeps a null field that an intersection requires', async () => { + const result = await runTool('{"count":null}', { + type: 'object', + allOf: [ + { properties: { count: { type: 'integer' } } }, + { required: ['count'] }, + ], + }) + expect(result.executed).toEqual([{ count: 0 }]) + }) + + it('keeps a null field required outside a union branch', async () => { + const result = await runTool('{"count":null}', { + type: 'object', + required: ['count'], + anyOf: [{ properties: { count: { type: 'integer' } } }], + }) + expect(result.executed).toEqual([{ count: 0 }]) + }) + + it('coerces pattern and additional properties', async () => { + const result = await runTool('{"n_1":"4","extra":"5"}', { + type: 'object', + patternProperties: { '^n_': { type: 'integer' } }, + additionalProperties: { type: 'number' }, + }) + expect(result.executed).toEqual([{ n_1: 4, extra: 5 }]) + }) + + it('coerces an approval edit before execution', async () => { + const result = await runTool('{"count":2}', schema, { + needsApproval: true, + approvals: new Map([ + ['call-1', { approved: true, editedArgs: { count: '4' } }], + ]), + }) + expect(result.executed).toEqual([{ count: 4 }]) + }) + + it('coerces a middleware rewrite before a client request', async () => { + const result = await runTool('{"count":2}', schema, { + client: true, + middleware: { + onBeforeToolCall: async () => ({ + type: 'transformArgs', + args: { count: '5' }, + }), + }, + }) + expect(result.needsClientExecution).toEqual([ + { toolCallId: 'call-1', toolName: 'check', input: { count: 5 } }, + ]) + }) + + it('repairs input for a Standard Schema that exports JSON Schema', async () => { + const result = await runTool( + '{"count":"2","label":null}', + z.object({ count: z.number(), label: z.string().optional() }), + ) + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('runs the authored transform on the repaired input', async () => { + const result = await runTool( + '{"count":"2"}', + z.object({ count: z.number().transform((value) => value * 10) }), + ) + expect(result.executed).toEqual([{ count: 20 }]) + }) + + it('reports the issues of the first run when repair does not help', async () => { + const result = await runTool( + '{"count":"wrong"}', + z.object({ count: z.number() }), + ) + expect(result.executed).toEqual([]) + expect(result.results[0]?.result).toEqual({ + error: expect.stringContaining('Received arguments: {"count":"wrong"}'), + }) + }) +}) + +const clientCheck = ( + inputSchema: SchemaInput, + needsApproval = false, +): Tool => ({ + name: 'check', + description: 'Check input', + inputSchema, + needsApproval, +}) + +const history = (input: unknown): Array => [ + { + role: 'assistant', + content: null, + toolCalls: [ + { + id: 'call', + type: 'function', + function: { name: 'check', arguments: JSON.stringify(input) }, + }, + ], + }, +] + +function descriptors(chunks: Array) { + return chunks.flatMap((chunk) => + chunk.type === EventType.RUN_FINISHED && chunk.outcome?.type === 'interrupt' + ? chunk.outcome.interrupts + : [], + ) +} + +function clientInput(chunks: Array) { + return descriptors(chunks).find((entry) => entry.id === 'client_tool_call') + ?.metadata?.input +} + +const review = defineInterrupt({ + id: 'boundary-review', + responseSchema: z.boolean(), +}) + +// Pauses at a tool boundary, so client requests come from that boundary. +const reviewAt = (phase: 'beforeTools' | 'afterTools') => + defineChatMiddleware({ + name: 'boundary', + onInterruptBoundary(ctx) { + if (ctx.phase !== phase) return undefined + return { + interrupts: [ + review.interrupt({ + key: 'review', + reason: 'review', + message: 'Review', + }), + ], + } + }, + }) + +async function initial( + input: unknown, + tool: Tool, + middleware: Array = [], + boundary?: 'beforeTools' | 'afterTools', + repairToolInput?: TextOptions['repairToolInput'], +) { + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ev.toolStart('call', 'check'), + ev.toolArgs('call', JSON.stringify(input)), + ev.runFinished('tool_calls'), + ], + ], + }) + const messages: Array = [{ role: 'user', content: 'check' }] + if (!boundary) { + return collectChunks( + chat({ adapter, messages, tools: [tool], middleware, repairToolInput }), + ) + } + return collectChunks( + chat({ + adapter, + messages, + tools: [tool], + interrupts: [review], + middleware: [reviewAt(boundary), ...middleware], + repairToolInput, + }), + ) +} + +describe('client tool input in chat()', () => { + it.each([false, true])( + 'shows middleware the raw input, then sends the checked input: boundary=%s', + async (boundary) => { + let transforms = 0 + const seen: Array = [] + const inputSchema = z.object({ + count: z.string().transform((value) => { + transforms++ + return Number(value) + }), + }) + const chunks = await initial( + { count: '2' }, + clientCheck(inputSchema), + [ + { + onBeforeToolCall(_ctx, hook) { + seen.push(structuredClone(hook.args)) + }, + }, + ], + boundary ? 'beforeTools' : undefined, + ) + expect(seen).toEqual([{ count: '2' }]) + expect(transforms).toBe(1) + expect(clientInput(chunks)).toEqual({ count: 2 }) + }, + ) + + it.each([false, true])( + 'sends no client request when middleware makes the input invalid: boundary=%s', + async (boundary) => { + const chunks = await initial( + { count: 2 }, + clientCheck(schema), + [ + { + onBeforeToolCall: () => ({ + type: 'transformArgs', + args: { count: 'wrong' }, + }), + }, + ], + boundary ? 'beforeTools' : undefined, + ) + expect(clientInput(chunks)).toBeUndefined() + expect( + chunks.some( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' && + chunk.content.includes('Input validation failed'), + ), + ).toBe(true) + }, + ) + + it('runs no middleware for the pending approval preview', async () => { + let hooks = 0 + const chunks = await initial( + { count: '2' }, + clientCheck(z.object({ count: z.string().transform(Number) }), true), + [ + { + onBeforeToolCall() { + hooks++ + }, + }, + ], + ) + expect(hooks).toBe(0) + expect( + descriptors(chunks).find((entry) => entry.id === 'approval_call') + ?.metadata?.input, + ).toEqual({ count: 2 }) + expect(clientInput(chunks)).toBeUndefined() + }) + + it('checks an approval edit after middleware on resume', async () => { + const { adapter, calls } = createMockAdapter({ iterations: [] }) + const seen: Array = [] + const inputSchema = z.object({ + count: z.number().transform((value) => value + 1), + }) + const chunks = await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 1 }), + tools: [clientCheck(inputSchema, true)], + resume: [ + { + interruptId: 'approval_call', + status: 'resolved', + payload: { approved: true, editedArgs: { count: 2 } }, + }, + ], + middleware: [ + { + onBeforeToolCall(_ctx, hook) { + seen.push(structuredClone(hook.args)) + }, + }, + ], + }), + ) + expect(calls).toHaveLength(0) + expect(seen).toEqual([{ count: 2 }]) + expect(clientInput(chunks)).toEqual({ count: 3 }) + }) + + it('keeps the client result, with no second hook run, when it comes back', async () => { + const { adapter, calls } = createMockAdapter({ + iterations: [[ev.runStarted(), ev.runFinished()]], + }) + let hooks = 0 + await collectChunks( + chat({ + adapter, + parentRunId: 'old', + messages: history({ count: 'wrong' }), + tools: [clientCheck(schema)], + resume: [ + { + interruptId: 'client_tool_call', + status: 'resolved', + payload: { ok: true }, + }, + ], + middleware: [ + { + onBeforeToolCall() { + hooks++ + }, + }, + ], + }), + ) + expect(hooks).toBe(0) + const results = calls[0]?.messages.filter( + (message) => message.role === 'tool', + ) + expect(results).toHaveLength(1) + expect(JSON.stringify(results?.[0])).not.toContain( + 'Input validation failed', + ) + }) + + it.each([false, true])( + 'runs the hook once for a client tool at an afterTools boundary: invalid=%s', + async (invalid) => { + let hooks = 0 + const chunks = await initial( + { count: 2 }, + clientCheck(schema), + [ + { + onBeforeToolCall() { + hooks++ + return invalid + ? { type: 'transformArgs', args: { count: 'wrong' } } + : undefined + }, + }, + ], + 'afterTools', + ) + expect(hooks).toBe(1) + const toolResults = chunks.filter( + (chunk) => chunk.type === EventType.TOOL_CALL_RESULT, + ) + expect(toolResults).toHaveLength(invalid ? 1 : 0) + if (invalid) { + expect(clientInput(chunks)).toBeUndefined() + } else { + expect(clientInput(chunks)).toEqual({ count: 2 }) + } + }, + ) + + it('records no execute_tool span for a client tool', async () => { + const fakeTracer = createFakeTracer() + await initial({ count: 2 }, clientCheck(schema), [ + otelMiddleware({ + tracer: fakeTracer.tracer, + meter: createFakeMeter().meter, + }), + ]) + expect( + fakeTracer.spans.some((span) => span.name.startsWith('execute_tool')), + ).toBe(false) + }) +}) + +describe('repairToolInput', () => { + const countSchema: JSONSchema = { + type: 'object', + properties: { count: { type: 'integer' } }, + required: ['count'], + } + + /** One model turn that calls `check` with each input, then a turn that answers. */ + async function run( + inputs: Array, + inputSchema: SchemaInput, + repairToolInput?: TextOptions['repairToolInput'], + ) { + const executed: Array = [] + const tool: Tool = { + name: 'check', + description: 'Check', + inputSchema, + execute: (input) => { + executed.push(input) + return 'ok' + }, + } + const { adapter } = createMockAdapter({ + iterations: [ + [ + ev.runStarted(), + ...inputs.flatMap((input, index) => [ + ev.toolStart(`call-${index}`, 'check'), + ev.toolArgs(`call-${index}`, JSON.stringify(input)), + ev.toolEnd(`call-${index}`), + ]), + ev.runFinished('tool_calls'), + ], + [ev.runStarted(), ev.textContent('Done'), ev.runFinished('stop')], + ], + }) + const chunks = await collectChunks( + chat({ + adapter, + messages: [{ role: 'user', content: 'check' }], + tools: [tool], + repairToolInput, + }), + ) + const results = chunks.flatMap((chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' + ? [chunk.content] + : [], + ) + const errors = results.filter((content) => + content.includes('Input validation failed'), + ) + return { executed, results, errors } + } + + it('turns the built-in repair off with false', async () => { + const result = await run([{ count: '3' }], countSchema, false) + expect(result.executed).toEqual([]) + expect(result.errors).toHaveLength(1) + }) + + it('turns the repair off with false for a Standard Schema', async () => { + const result = await run( + [{ count: '3' }], + z.object({ count: z.number() }), + false, + ) + expect(result.executed).toEqual([]) + expect(result.errors).toHaveLength(1) + }) + + it.each([false, true])( + 'turns the repair off with false for a client tool: boundary=%s', + async (boundary) => { + const chunks = await initial( + { count: '3' }, + clientCheck(schema), + [], + boundary ? 'beforeTools' : undefined, + false, + ) + expect(clientInput(chunks)).toBeUndefined() + expect( + chunks.some( + (chunk) => + chunk.type === EventType.TOOL_CALL_RESULT && + typeof chunk.content === 'string' && + chunk.content.includes('Input validation failed'), + ), + ).toBe(true) + }, + ) + + it('replaces the built-in repair with a function', async () => { + const repair = vi.fn(({ input }: ToolInputRepairContext) => + typeof input === 'object' && + input !== null && + 'count' in input && + input.count === 'many' + ? { count: 7 } + : undefined, + ) + const result = await run( + [{ count: 'many' }, { count: '3' }], + countSchema, + repair, + ) + // `'3'` is not repaired: the function said `undefined`. + expect(result.executed).toEqual([{ count: 7 }]) + expect(result.errors).toHaveLength(1) + expect(result.errors[0]).toContain( + String.raw`Received arguments: {\"count\":\"3\"}`, + ) + const context = repair.mock.calls[0]?.[0] + expect(context?.tool.name).toBe('check') + expect(context?.toolCall.id).toBe('call-0') + expect(context?.input).toEqual({ count: 'many' }) + expect(context?.error.message).toContain('Input validation failed') + }) + + it('does not call the function for a valid input', async () => { + const repair = vi.fn(() => ({ count: 9 })) + const result = await run([{ count: 2 }], countSchema, repair) + expect(repair).not.toHaveBeenCalled() + expect(result.executed).toEqual([{ count: 2 }]) + }) + + it('reports the original error when the repaired input still fails', async () => { + const result = await run([{ count: 'many' }], countSchema, () => ({ + count: 'still wrong', + })) + expect(result.executed).toEqual([]) + expect(result.errors[0]).toContain( + String.raw`Received arguments: {\"count\":\"many\"}`, + ) + }) + + it('runs a Standard Schema, with its transform, on the repaired input', async () => { + const result = await run( + [{ count: 'two' }], + z.object({ count: z.number().transform((value) => value * 10) }), + () => ({ count: 2 }), + ) + expect(result.executed).toEqual([{ count: 20 }]) + }) + + it('makes the thrown error the tool error', async () => { + const result = await run([{ count: 'many' }], countSchema, () => { + throw new Error('cannot repair') + }) + expect(result.executed).toEqual([]) + // The model's own arguments follow the thrown message. + expect(result.results).toEqual([ + JSON.stringify({ + error: 'cannot repair\nReceived arguments: {"count":"many"}', + }), + ]) + }) +}) diff --git a/packages/openai-base/src/adapters/chat-completions-text.ts b/packages/openai-base/src/adapters/chat-completions-text.ts index ae0e6172f8..43048659ba 100644 --- a/packages/openai-base/src/adapters/chat-completions-text.ts +++ b/packages/openai-base/src/adapters/chat-completions-text.ts @@ -220,6 +220,7 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolCallId, toolCallName: rejectedToolCall.toolName, toolName: rejectedToolCall.toolName, + args: rejectedToolCall.arguments, ...(rejectedToolCall.input !== undefined && { input: rejectedToolCall.input, }), @@ -1166,7 +1167,8 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolName: toolCall.name, model: chunk.model || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } emittedAnyToolCallEnd = true } @@ -1236,7 +1238,8 @@ export abstract class OpenAIBaseChatCompletionsTextAdapter< toolName: toolCall.name, model: lastModel || options.model, timestamp: Date.now(), - input: parsedInput, + args: toolCall.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } pendingToolCount += 1 emittedAnyToolCallEnd = true diff --git a/packages/openai-base/src/adapters/responses-text.ts b/packages/openai-base/src/adapters/responses-text.ts index cb8c05b562..b59154da5b 100644 --- a/packages/openai-base/src/adapters/responses-text.ts +++ b/packages/openai-base/src/adapters/responses-text.ts @@ -2074,7 +2074,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: chunk.arguments, }, ) - parsedInput = {} + parsedInput = undefined } } @@ -2085,7 +2085,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: chunk.arguments, + ...(parsedInput !== undefined && { input: parsedInput }), } } @@ -2171,7 +2172,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -2181,7 +2182,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined @@ -2342,7 +2344,7 @@ export abstract class OpenAIBaseResponsesTextAdapter< rawArguments: rawArgs, }, ) - parsedInput = {} + parsedInput = undefined } } yield { @@ -2352,7 +2354,8 @@ export abstract class OpenAIBaseResponsesTextAdapter< toolName: name, model: model || options.model, timestamp: Date.now(), - input: parsedInput, + args: rawArgs, + ...(parsedInput !== undefined && { input: parsedInput }), } metadata.ended = true metadata.pendingArguments = undefined diff --git a/packages/openai-base/tests/chat-completions-text.test.ts b/packages/openai-base/tests/chat-completions-text.test.ts index 1fe050309a..8b16063ece 100644 --- a/packages/openai-base/tests/chat-completions-text.test.ts +++ b/packages/openai-base/tests/chat-completions-text.test.ts @@ -1240,7 +1240,8 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { const toolEnd = chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') expect(toolEnd).toBeDefined() - expect(toolEnd?.input).toBeUndefined() + expect(toolEnd).toMatchObject({ args: '{"location":' }) + expect(toolEnd).not.toHaveProperty('input') expect( chunks.find((chunk) => chunk.type === 'TOOL_CALL_ARGS'), ).toMatchObject({ @@ -1263,6 +1264,64 @@ describe('OpenAIBaseChatCompletionsTextAdapter', () => { ) }) + describe('raw tool arguments', () => { + it.each([true, false])( + 'sends the accumulated argument string as args on TOOL_CALL_END (finish_reason=%s)', + async (withFinishReason) => { + setupMockSdkClient([ + { + id: 'chatcmpl-raw', + model: 'test-model', + choices: [ + { + delta: { + tool_calls: [ + { + index: 0, + id: 'call_raw', + type: 'function', + function: { name: 'lookup_weather', arguments: '{"n": ' }, + }, + ], + }, + finish_reason: null, + }, + ], + }, + { + id: 'chatcmpl-raw', + model: 'test-model', + choices: [ + { + delta: { + tool_calls: [{ index: 0, function: { arguments: '1.50}' } }], + }, + finish_reason: withFinishReason ? 'tool_calls' : null, + }, + ], + }, + ]) + const adapter = new TestChatCompletionsAdapter(testConfig, 'test-model') + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'Weather?' }], + tools: [weatherTool], + })) { + chunks.push(chunk) + } + + const toolEnd = chunks.find((chunk) => chunk.type === 'TOOL_CALL_END') + expect(toolEnd).toMatchObject({ + toolCallId: 'call_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + }) + describe('subclassing', () => { it('allows subclassing with custom name', () => { class MyProviderAdapter extends OpenAIBaseChatCompletionsTextAdapter { diff --git a/packages/openai-base/tests/responses-text.test.ts b/packages/openai-base/tests/responses-text.test.ts index 404d41239a..6dc49e6bcd 100644 --- a/packages/openai-base/tests/responses-text.test.ts +++ b/packages/openai-base/tests/responses-text.test.ts @@ -1869,6 +1869,108 @@ describe('OpenAIBaseResponsesTextAdapter', () => { }) }) + describe('raw tool arguments', () => { + type EndSite = 'arguments.done' | 'output_item.done' | 'response.completed' + + // Each site emits TOOL_CALL_END from a different event. Without a name on + // output_item.added, args.done is deferred to a later backfill. + function rawArgsStream( + site: EndSite, + raw: string, + ): Array> { + const named = { + type: 'function_call', + id: 'fc_raw', + name: 'lookup_weather', + } + return [ + { + type: 'response.created', + response: { + id: 'resp-raw', + model: 'test-model', + status: 'in_progress', + }, + }, + { + type: 'response.output_item.added', + output_index: 0, + item: + site === 'arguments.done' + ? named + : { type: 'function_call', id: 'fc_raw' }, + }, + { + type: 'response.function_call_arguments.done', + item_id: 'fc_raw', + arguments: raw, + }, + ...(site === 'output_item.done' + ? [ + { + type: 'response.output_item.done', + output_index: 0, + item: { ...named, arguments: raw }, + }, + ] + : []), + { + type: 'response.completed', + response: { + id: 'resp-raw', + model: 'test-model', + status: 'completed', + output: [{ ...named, arguments: raw }], + usage: { input_tokens: 1, output_tokens: 1, total_tokens: 2 }, + }, + }, + ] + } + + async function toolEnds(site: EndSite, raw: string) { + setupMockResponsesClient(rawArgsStream(site, raw)) + const adapter = new TestResponsesAdapter(testConfig, 'test-model') + const chunks: Array = [] + for await (const chunk of adapter.chatStream({ + logger: testLogger, + model: 'test-model', + messages: [{ role: 'user', content: 'q' }], + tools: [weatherTool], + })) { + chunks.push(chunk) + } + return chunks.filter((chunk) => chunk.type === 'TOOL_CALL_END') + } + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])( + 'sends the raw argument string as args on TOOL_CALL_END (%s)', + async (site) => { + const ends = await toolEnds(site, '{"n": 1.50}') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ + toolCallId: 'fc_raw', + args: '{"n": 1.50}', + input: { n: 1.5 }, + }) + }, + ) + + it.each([ + 'arguments.done', + 'output_item.done', + 'response.completed', + ])('sends malformed arguments as args with no input (%s)', async (site) => { + const ends = await toolEnds(site, '{"a":') + expect(ends).toHaveLength(1) + expect(ends[0]).toMatchObject({ toolCallId: 'fc_raw', args: '{"a":' }) + expect(ends[0]).not.toHaveProperty('input') + }) + }) + describe('content_part events', () => { it('emits TEXT_MESSAGE_START on content_part.added with output_text', async () => { const streamChunks = [ diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index aa9e45eb08..9c11f88a32 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1072,10 +1072,10 @@ importers: version: 2.68.0 '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(supports-color@7.2.0)(zod@4.2.1) + version: 1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1) '@ngrok/ngrok': specifier: ^1.7.0 version: 1.7.0 @@ -2011,6 +2011,9 @@ importers: '@ag-ui/core': specifier: 1.0.0 version: 1.0.0(zod@4.2.1) + '@cfworker/json-schema': + specifier: ^4.1.1 + version: 4.1.1 '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -2497,7 +2500,7 @@ importers: dependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai-utils': specifier: workspace:^ version: link:../ai-utils @@ -2976,7 +2979,7 @@ importers: devDependencies: '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3007,7 +3010,7 @@ importers: devDependencies: '@mcp-ui/client': specifier: ^7.1.1 - version: 7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3) + version: 7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3) '@standard-schema/spec': specifier: ^1.1.0 version: 1.1.0 @@ -3100,7 +3103,7 @@ importers: dependencies: '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(zod@4.3.6) + version: 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) '@tanstack/ai-skills': specifier: workspace:^ version: link:../ai-skills @@ -3125,7 +3128,7 @@ importers: dependencies: '@blaxel/core': specifier: ^0.3.10 - version: 0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) + version: 0.3.11(@cfworker/json-schema@4.1.1)(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) devDependencies: '@tanstack/ai-sandbox': specifier: workspace:* @@ -3450,7 +3453,7 @@ importers: devDependencies: '@google/genai': specifier: ^2.10.0 - version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)) + version: 2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)) '@tanstack/ai': specifier: workspace:* version: link:../ai @@ -3641,7 +3644,7 @@ importers: version: 1.34.0(vitest@4.1.11(@opentelemetry/api@1.9.1)(@types/node@24.10.3)(@vitest/coverage-v8@4.1.10(vitest@4.1.10))(happy-dom@20.11.2)(jsdom@27.4.0(@noble/hashes@2.3.0))(vite@8.2.1(@types/node@24.10.3)(esbuild@0.28.1)(jiti@2.7.0)(less@4.6.6)(sass@1.101.0)(terser@5.44.1)(tsx@4.23.12)(yaml@2.9.0))) '@modelcontextprotocol/sdk': specifier: ^1.29.0 - version: 1.29.0(zod@4.3.6) + version: 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) '@openrouter/sdk': specifier: 0.13.20 version: 0.13.20 @@ -5106,6 +5109,9 @@ packages: '@bufbuild/protobuf@2.15.0': resolution: {integrity: sha512-DAheWUkVr/SJTWCc+lg9dhY0eN4SaWlf4+bG1KzHeXbnqt0AfB/NX0Z+VunGlM1ki1B4zVvye27MpKh/svySUA==} + '@cfworker/json-schema@4.1.1': + resolution: {integrity: sha512-gAmrUZSGtKc3AiBL71iNWxDsyUC5uMaKKGdvzYsBoTW/xi42JQHl7eKV2OYzCUqvc+D2RCcf7EXY2iCyFIk6og==} + '@changesets/apply-release-plan@7.1.0': resolution: {integrity: sha512-yq8ML3YS7koKQ/9bk1PqO0HMzApIFNwjlwCnwFEXMzNe8NpzeeYYKCmnhWJGkN8g7E51MnWaSbqRcTcdIxUgnQ==} @@ -20988,10 +20994,10 @@ snapshots: '@bcoe/v8-coverage@1.0.2': {} - '@blaxel/core@0.3.11(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': + '@blaxel/core@0.3.11(@cfworker/json-schema@4.1.1)(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2))': dependencies: '@hey-api/client-fetch': 0.10.2(@hey-api/openapi-ts@0.99.0(magicast@0.5.2)(typescript@7.0.2)) - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) archiver: 7.0.1 axios: 1.18.1(debug@4.4.3(supports-color@7.2.0))(supports-color@7.2.0) dockerfile-ast: 0.7.1 @@ -21022,6 +21028,8 @@ snapshots: '@bufbuild/protobuf@2.15.0': {} + '@cfworker/json-schema@4.1.1': {} + '@changesets/apply-release-plan@7.1.0': dependencies: '@changesets/config': 3.1.3 @@ -22223,27 +22231,27 @@ snapshots: '@shikijs/types': 3.23.0 '@shikijs/vscode-textmate': 10.0.2 - '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.3.6))': + '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6))': dependencies: google-auth-library: 10.5.0 p-retry: 4.6.2 protobufjs: 7.6.4 ws: 8.21.0 optionalDependencies: - '@modelcontextprotocol/sdk': 1.29.0(zod@4.3.6) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6) transitivePeerDependencies: - bufferutil - supports-color - utf-8-validate - '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(zod@4.6.5))': + '@google/genai@2.10.0(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5))': dependencies: google-auth-library: 10.5.0 p-retry: 4.6.2 protobufjs: 7.6.4 ws: 8.21.0 optionalDependencies: - '@modelcontextprotocol/sdk': 1.29.0(zod@4.6.5) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5) transitivePeerDependencies: - bufferutil - supports-color @@ -22637,10 +22645,10 @@ snapshots: - encoding - supports-color - '@mcp-ui/client@7.1.1(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': + '@mcp-ui/client@7.1.1(@cfworker/json-schema@4.1.1)(react-dom@19.2.3(react@19.2.3))(react@19.2.3)': dependencies: - '@modelcontextprotocol/ext-apps': 1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76) - '@modelcontextprotocol/sdk': 1.29.0(zod@3.25.76) + '@modelcontextprotocol/ext-apps': 1.7.4(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(zod@3.25.76) react: 19.2.3 react-dom: 19.2.3(react@19.2.3) zod: 3.25.76 @@ -22738,16 +22746,16 @@ snapshots: dependencies: zod: 4.3.6 - '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': + '@modelcontextprotocol/ext-apps@1.7.4(@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1))(react-dom@19.2.3(react@19.2.3))(react@19.2.3)(zod@3.25.76)': dependencies: - '@modelcontextprotocol/sdk': 1.29.0(supports-color@7.2.0)(zod@4.2.1) + '@modelcontextprotocol/sdk': 1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1) '@standard-schema/spec': 1.1.0 zod: 3.25.76 optionalDependencies: react: 19.2.3 react-dom: 19.2.3(react@19.2.3) - '@modelcontextprotocol/sdk@1.29.0(supports-color@7.2.0)(zod@4.2.1)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(supports-color@7.2.0)(zod@4.2.1)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22766,10 +22774,12 @@ snapshots: raw-body: 3.0.2 zod: 4.2.1 zod-to-json-schema: 3.25.2(zod@4.2.1) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@3.25.76)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@3.25.76)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22788,10 +22798,12 @@ snapshots: raw-body: 3.0.2 zod: 3.25.76 zod-to-json-schema: 3.25.2(zod@3.25.76) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@4.3.6)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.3.6)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22810,10 +22822,12 @@ snapshots: raw-body: 3.0.2 zod: 4.3.6 zod-to-json-schema: 3.25.2(zod@4.3.6) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color - '@modelcontextprotocol/sdk@1.29.0(zod@4.6.5)': + '@modelcontextprotocol/sdk@1.29.0(@cfworker/json-schema@4.1.1)(zod@4.6.5)': dependencies: '@hono/node-server': 1.19.17(hono@4.13.8) ajv: 8.20.0 @@ -22832,6 +22846,8 @@ snapshots: raw-body: 3.0.2 zod: 4.6.5 zod-to-json-schema: 3.25.2(zod@4.6.5) + optionalDependencies: + '@cfworker/json-schema': 4.1.1 transitivePeerDependencies: - supports-color optional: true diff --git a/testing/e2e/fixtures/middleware-test/raw-tool-args.json b/testing/e2e/fixtures/middleware-test/raw-tool-args.json new file mode 100644 index 0000000000..97e159b49d --- /dev/null +++ b/testing/e2e/fixtures/middleware-test/raw-tool-args.json @@ -0,0 +1,27 @@ +{ + "fixtures": [ + { + "match": { + "userMessage": "[raw-tool-args] run test", + "sequenceIndex": 0 + }, + "response": { + "toolCalls": [ + { + "name": "get_weather", + "arguments": "{\"city\": \"NYC\", \"days\": 1.50}" + } + ] + } + }, + { + "match": { + "userMessage": "[raw-tool-args] run test", + "sequenceIndex": 1 + }, + "response": { + "content": "The weather is sunny." + } + } + ] +} diff --git a/testing/e2e/src/routeTree.gen.ts b/testing/e2e/src/routeTree.gen.ts index 6fb7e92c61..cbbd7a9f08 100644 --- a/testing/e2e/src/routeTree.gen.ts +++ b/testing/e2e/src/routeTree.gen.ts @@ -96,6 +96,7 @@ import { Route as ApiGenerationPersistenceServerRouteImport } from './routes/api import { Route as ApiImageRouteImport } from './routes/api.image' import { Route as ApiInterruptsTestRouteImport } from './routes/api.interrupts-test' import { Route as ApiJoinRunClientToolRouteImport } from './routes/api.join-run-client-tool' +import { Route as ApiJsonSchemaToolInputRouteImport } from './routes/api.json-schema-tool-input' import { Route as ApiJsonTransportRouteImport } from './routes/api.json-transport' import { Route as ApiKeyedAdapterRouteImport } from './routes/api.keyed-adapter' import { Route as ApiLazyToolsWireRouteImport } from './routes/api.lazy-tools-wire' @@ -636,6 +637,11 @@ const ApiJoinRunClientToolRoute = ApiJoinRunClientToolRouteImport.update({ path: '/api/join-run-client-tool', getParentRoute: () => rootRouteImport, } as any) +const ApiJsonSchemaToolInputRoute = ApiJsonSchemaToolInputRouteImport.update({ + id: '/api/json-schema-tool-input', + path: '/api/json-schema-tool-input', + getParentRoute: () => rootRouteImport, +} as any) const ApiJsonTransportRoute = ApiJsonTransportRouteImport.update({ id: '/api/json-transport', path: '/api/json-transport', @@ -1160,6 +1166,7 @@ export interface FileRoutesByFullPath { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1332,6 +1339,7 @@ export interface FileRoutesByTo { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1505,6 +1513,7 @@ export interface FileRoutesById { '/api/image': typeof ApiImageRouteWithChildren '/api/interrupts-test': typeof ApiInterruptsTestRoute '/api/join-run-client-tool': typeof ApiJoinRunClientToolRoute + '/api/json-schema-tool-input': typeof ApiJsonSchemaToolInputRoute '/api/json-transport': typeof ApiJsonTransportRoute '/api/keyed-adapter': typeof ApiKeyedAdapterRoute '/api/lazy-tools-wire': typeof ApiLazyToolsWireRoute @@ -1679,6 +1688,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -1851,6 +1861,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -2023,6 +2034,7 @@ export interface FileRouteTypes { | '/api/image' | '/api/interrupts-test' | '/api/join-run-client-tool' + | '/api/json-schema-tool-input' | '/api/json-transport' | '/api/keyed-adapter' | '/api/lazy-tools-wire' @@ -2196,6 +2208,7 @@ export interface RootRouteChildren { ApiImageRoute: typeof ApiImageRouteWithChildren ApiInterruptsTestRoute: typeof ApiInterruptsTestRoute ApiJoinRunClientToolRoute: typeof ApiJoinRunClientToolRoute + ApiJsonSchemaToolInputRoute: typeof ApiJsonSchemaToolInputRoute ApiJsonTransportRoute: typeof ApiJsonTransportRoute ApiKeyedAdapterRoute: typeof ApiKeyedAdapterRoute ApiLazyToolsWireRoute: typeof ApiLazyToolsWireRoute @@ -2888,6 +2901,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof ApiJoinRunClientToolRouteImport parentRoute: typeof rootRouteImport } + '/api/json-schema-tool-input': { + id: '/api/json-schema-tool-input' + path: '/api/json-schema-tool-input' + fullPath: '/api/json-schema-tool-input' + preLoaderRoute: typeof ApiJsonSchemaToolInputRouteImport + parentRoute: typeof rootRouteImport + } '/api/json-transport': { id: '/api/json-transport' path: '/api/json-transport' @@ -3622,6 +3642,7 @@ const rootRouteChildren: RootRouteChildren = { ApiImageRoute: ApiImageRouteWithChildren, ApiInterruptsTestRoute: ApiInterruptsTestRoute, ApiJoinRunClientToolRoute: ApiJoinRunClientToolRoute, + ApiJsonSchemaToolInputRoute: ApiJsonSchemaToolInputRoute, ApiJsonTransportRoute: ApiJsonTransportRoute, ApiKeyedAdapterRoute: ApiKeyedAdapterRoute, ApiLazyToolsWireRoute: ApiLazyToolsWireRoute, diff --git a/testing/e2e/src/routes/api.json-schema-tool-input.ts b/testing/e2e/src/routes/api.json-schema-tool-input.ts new file mode 100644 index 0000000000..5cba53fc80 --- /dev/null +++ b/testing/e2e/src/routes/api.json-schema-tool-input.ts @@ -0,0 +1,87 @@ +import { createFileRoute } from '@tanstack/react-router' +import { chat, toolDefinition } from '@tanstack/ai' +import { fakeText } from '@tanstack/ai/testing' + +// No network: `fakeText` answers from a script. The first answer calls a +// JSON Schema tool twice: once with a value that cannot be repaired, once +// with a number sent as a string. It also calls an OpenAPI 3.0 tool with +// `null` for a `nullable` field. The second answer repeats the tool results. +// MCP tools have the same raw JSON Schema input. With `{ repair: 'custom' }`, +// a `repairToolInput` function replaces the built-in repair: it maps +// `'many'` to 2 and repairs nothing else. +export const Route = createFileRoute('/api/json-schema-tool-input')({ + server: { + handlers: { + POST: async ({ request }) => { + const body: unknown = await request.json().catch(() => ({})) + const custom = + typeof body === 'object' && + body !== null && + 'repair' in body && + body.repair === 'custom' + + const fake = fakeText() + fake.setResponses([ + { + toolCalls: [ + { name: 'count_items', input: { count: 'many' } }, + { name: 'count_items', input: { count: '3' } }, + { name: 'label_item', input: { label: null } }, + ], + }, + ({ request }) => ({ + text: request.messages + .filter((message) => message.role === 'tool') + .map((message) => `Tool said: ${String(message.content)}`) + .join('\n'), + }), + ]) + + const runs: Array = [] + const countItems = toolDefinition({ + name: 'count_items', + description: 'Count the items', + inputSchema: { + $schema: 'http://json-schema.org/draft-07/schema#', + type: 'object', + properties: { count: { type: 'integer' } }, + required: ['count'], + }, + }).server((input) => { + runs.push(input) + return 'counted' + }) + const labelItem = toolDefinition({ + name: 'label_item', + description: 'Label the item', + inputSchema: { + type: 'object', + properties: { label: { type: 'string', nullable: true } }, + required: ['label'], + }, + }).server((input) => { + runs.push(input) + return 'labeled' + }) + + const { text } = await chat({ + adapter: fake, + messages: [{ role: 'user', content: 'Count the items' }], + tools: [countItems, labelItem], + stream: false, + ...(custom && { + repairToolInput: ({ input }) => + typeof input === 'object' && + input !== null && + 'count' in input && + input.count === 'many' + ? { count: 2 } + : undefined, + }), + }) + + return Response.json({ text, runs }) + }, + }, + }, +}) diff --git a/testing/e2e/src/routes/api.middleware-test.ts b/testing/e2e/src/routes/api.middleware-test.ts index 8c7213fe5a..f5d7cb672a 100644 --- a/testing/e2e/src/routes/api.middleware-test.ts +++ b/testing/e2e/src/routes/api.middleware-test.ts @@ -177,6 +177,16 @@ const toolReplaceMiddleware: ChatMiddleware = { }, } +// Rewrites the args to a value the tool's inputSchema rejects. The final +// input check runs after this hook, so the tool must not run. +const toolBadArgsMiddleware: ChatMiddleware = { + name: 'tool-bad-args', + onBeforeToolCall(_ctx, hookCtx) { + if (hookCtx.toolName !== 'get_weather') return undefined + return { type: 'transformArgs' as const, args: { city: { name: 'NYC' } } } + }, +} + /** * Capability provide/consume flow (`capability` mode). * @@ -579,6 +589,8 @@ export const Route = createFileRoute('/api/middleware-test')({ middleware.push(toolSkipMiddleware) if (middlewareMode === 'tool-replace') middleware.push(toolReplaceMiddleware) + if (middlewareMode === 'tool-bad-args') + middleware.push(toolBadArgsMiddleware) if (middlewareMode === 'capability') { // Order matters: the provider's setup() must run before the // consumer reads the capability. Array order is preserved. @@ -663,7 +675,7 @@ export const Route = createFileRoute('/api/middleware-test')({ const tools = genericScenario && testId ? genericTools(testId, genericScenario) - : scenario === 'with-tool' + : scenario === 'with-tool' || scenario === 'raw-tool-args' ? [weatherTool] : scenario === 'structured-client-tool-wait' ? [clientContextToolDefinition] diff --git a/testing/e2e/src/routes/api.tools-test.ts b/testing/e2e/src/routes/api.tools-test.ts index 3e40bc8957..43f5815264 100644 --- a/testing/e2e/src/routes/api.tools-test.ts +++ b/testing/e2e/src/routes/api.tools-test.ts @@ -65,9 +65,10 @@ function createProviderFreeAdapter(scenario: string): AnyTextAdapter { : scenario === 'client-tool-input-error' || scenario === 'invalid-client-tool-retry' ? { - arguments: '{"message":42,"type":"info"}', + // An object cannot be repaired into the string `message`. + arguments: '{"message":{"text":"hi"},"type":"info"}', initialText: 'Showing a notification.', - input: { message: 42, type: 'info' }, + input: { message: { text: 'hi' }, type: 'info' }, name: scenario === 'invalid-client-tool-retry' ? 'invalid-client-tool-retry-test' diff --git a/testing/e2e/src/routes/middleware-test.tsx b/testing/e2e/src/routes/middleware-test.tsx index 6bde171f32..fc326cc99b 100644 --- a/testing/e2e/src/routes/middleware-test.tsx +++ b/testing/e2e/src/routes/middleware-test.tsx @@ -19,6 +19,7 @@ const MIDDLEWARE_MODES = [ { id: 'chunk-transform', label: 'Chunk Transform (prefix text)' }, { id: 'tool-skip', label: 'Tool Skip (skip with custom result)' }, { id: 'tool-replace', label: 'Tool Replace (replace result after call)' }, + { id: 'tool-bad-args', label: 'Tool Bad Args (invalid args rewrite)' }, { id: 'capability', label: 'Capability (provide/consume prefix)' }, { id: 'phase-recorder', label: 'Phase Recorder (capture phase + chunks)' }, { id: 'otel', label: 'OpenTelemetry (capture spans/metrics)' }, @@ -253,6 +254,7 @@ function MiddlewareTestPage() { +