From 5036284fe8d62084992eb3fc6fa75e1252837298 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Tue, 21 Jul 2026 16:15:01 +1000 Subject: [PATCH 1/3] feat(ai): rework tool-call budgets as middleware hooks Replace the unreleased maxToolCalls strategy and chat({ maxToolCallsPerTurn }) with an onShouldContinue middleware hook so apps can stop further agent turns without aborting. Keep toolCallCount / lastTurnToolCallCount on AgentLoopState for strategies and middleware. Document an app-owned tool-budget recipe instead of shipping a built-in maxToolCallsMiddleware. --- .changeset/max-tool-calls-middleware.md | 53 +++++++ docs/advanced/built-in-middleware.md | 2 + docs/advanced/middleware.md | 21 +++ docs/api/ai.md | 35 +---- docs/chat/agentic-cycle.md | 90 +++++++++--- docs/config.json | 9 +- docs/reference/functions/combineStrategies.md | 1 - docs/reference/functions/maxIterations.md | 6 +- docs/reference/functions/maxToolCalls.md | 53 ------- docs/reference/index.md | 1 - docs/reference/interfaces/AgentLoopState.md | 6 +- docs/reference/interfaces/TextOptions.md | 26 ---- .../ai/skills/ai-core/middleware/SKILL.md | 47 ++++++ .../ai/skills/ai-core/tool-calling/SKILL.md | 4 +- .../activities/chat/agent-loop-strategies.ts | 44 +----- packages/ai/src/activities/chat/index.ts | 137 ++++++------------ .../src/activities/chat/middleware/compose.ts | 42 +++++- .../src/activities/chat/middleware/types.ts | 20 +++ packages/ai/src/index.ts | 1 - packages/ai/src/types.ts | 28 +--- .../ai/tests/agent-loop-strategies.test.ts | 37 ----- packages/ai/tests/chat.test.ts | 83 +++++++---- .../e2e/src/routes/api.max-tool-calls-wire.ts | 66 +++++++-- testing/e2e/tests/max-tool-calls.spec.ts | 12 +- 24 files changed, 436 insertions(+), 388 deletions(-) create mode 100644 .changeset/max-tool-calls-middleware.md delete mode 100644 docs/reference/functions/maxToolCalls.md diff --git a/.changeset/max-tool-calls-middleware.md b/.changeset/max-tool-calls-middleware.md new file mode 100644 index 000000000..1a9013eb8 --- /dev/null +++ b/.changeset/max-tool-calls-middleware.md @@ -0,0 +1,53 @@ +--- +'@tanstack/ai': minor +--- + +Rework tool-call fan-out budgets as middleware hooks (unreleased #965 API). + +- **Remove** (never released): `maxToolCalls()` strategy and `chat({ maxToolCallsPerTurn })` +- **Add** `onShouldContinue` middleware hook so policies can stop further agent turns without aborting +- **Keep** `AgentLoopState.toolCallCount` / `lastTurnToolCallCount` for strategies and middleware +- Tool-call budgets are an **app-owned middleware recipe** (docs), not a built-in export + +```ts +import { chat, maxIterations, type ChatMiddleware } from '@tanstack/ai' + +function toolCallBudget({ + max, + maxPerTurn, +}: { + max?: number + maxPerTurn?: number +}): ChatMiddleware { + let perTurn = 0 + return { + onIteration: () => { + perTurn = 0 + }, + onToolPhaseComplete: () => { + perTurn = 0 + }, + onBeforeToolCall: () => { + if (maxPerTurn == null) return + if (++perTurn > maxPerTurn) { + return { + type: 'skip', + result: { + error: `Skipped: exceeded maxToolCallsPerTurn (${maxPerTurn})`, + }, + } + } + }, + onShouldContinue: (_ctx, state) => + max != null && state.toolCallCount >= max ? false : undefined, + } +} + +chat({ + adapter, + messages, + tools, + agentLoopStrategy: maxIterations(20), + middleware: [toolCallBudget({ maxPerTurn: 10, max: 20 })], +}) +``` diff --git a/docs/advanced/built-in-middleware.md b/docs/advanced/built-in-middleware.md index b2c888c49..87cd59eec 100644 --- a/docs/advanced/built-in-middleware.md +++ b/docs/advanced/built-in-middleware.md @@ -23,6 +23,8 @@ TanStack AI ships ready-made middleware so you don't have to hand-roll the commo > `toolCacheMiddleware` and `contentGuardMiddleware` are exported from the main `@tanstack/ai/middlewares` barrel. `otelMiddleware` lives on its own subpath (`@tanstack/ai/middlewares/otel`) so that importing the barrel never eagerly pulls in `@opentelemetry/api` (an optional peer dependency). +For app-owned policies (for example tool-call budgets), see the [tool-call budget recipe](../chat/agentic-cycle#tool-call-budgets-middleware-recipe) — those stay in your code, not in `@tanstack/ai/middlewares`. + ## toolCacheMiddleware Caches tool call results based on tool name and arguments. When a tool is called with the same name and arguments as a previous call, the cached result is returned immediately without re-executing the tool. diff --git a/docs/advanced/middleware.md b/docs/advanced/middleware.md index a416e8f70..80fe5fed4 100644 --- a/docs/advanced/middleware.md +++ b/docs/advanced/middleware.md @@ -312,6 +312,26 @@ const redactStructuredOutput: ChatMiddleware = { > Why is there `onStructuredOutputConfig` but no `onStructuredOutputChunk`? Because the **config** shape genuinely differs at the structured-output boundary — it carries an `outputSchema` field that plain `ChatMiddlewareConfig` doesn't (see [onStructuredOutputConfig](#onstructuredoutputconfig)). **Chunks** are all just `StreamChunk` regardless of phase, so one `onChunk` plus `ctx.phase` (or the CUSTOM event name) covers every case — a parallel chunk hook would be redundant. +### onShouldContinue + +Called when the engine is deciding whether to start another agent-loop iteration (after a tool phase or between model turns). Combined with AND semantics across middleware **and** with `agentLoopStrategy` — any explicit `false` stops the loop. Return `true`, `void`, or `undefined` to allow continuation. + +Does **not** abort the run: the stream finishes normally with the current messages. Use `ctx.abort()` only for a hard abort. + +```typescript +import { type ChatMiddleware } from "@tanstack/ai"; + +const budget: ChatMiddleware = { + name: "tool-budget", + onShouldContinue: (_ctx, state) => { + // Stop further turns once 20 tool calls have been emitted + if (state.toolCallCount >= 20) return false; + }, +}; +``` + +For a full per-turn + cumulative tool budget recipe, see [Tool-call budgets](../chat/agentic-cycle#tool-call-budgets-middleware-recipe). + ### 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. @@ -596,6 +616,7 @@ const stream = chat({ | `onStart` | Sequential | All run in order | | `onChunk` | **Piped** — chunks flow through each middleware | If first drops a chunk, later middleware never see it | | `onBeforeToolCall` | **First-win** — first non-void decision wins | Earlier middleware has priority | +| `onShouldContinue` | **AND** — any explicit `false` stops the loop | Order only affects which middleware runs first when short-circuiting | | `onAfterToolCall` | Sequential | All run in order | | `onUsage` | Sequential | All run in order | | `onFinish/onAbort/onError` | Sequential | All run in order | diff --git a/docs/api/ai.md b/docs/api/ai.md index ad4418bad..c70aeb861 100644 --- a/docs/api/ai.md +++ b/docs/api/ai.md @@ -46,8 +46,8 @@ const stream = chat({ - `tools?` - Array of tools for function calling - `context?` - Typed runtime context passed to server tools and middleware. If a tool or middleware declares a concrete context type, `chat()` requires a compatible value here - `systemPrompts?` - System prompts to prepend to messages -- `agentLoopStrategy?` - Strategy for agent loops (default: `maxIterations(5)`). Strategies receive `{ iterationCount, finishReason, messages, toolCallCount, lastTurnToolCallCount }` and run between model turns. Prefer `maxToolCalls(n)` when you need a tool-call budget — iterations are model turns, not tool calls. `maxToolCalls` stops *further* turns once cumulative **emitted** tool calls are `>= n`; the turn that crosses the limit is not truncated. -- `maxToolCallsPerTurn?` - Cap how many tool calls from a single model turn (or pending/resume batch) are executed. Excess calls receive error results so message history stays consistent. Unset means no per-turn execution cap; `0` skips all execution for that batch. Pair with `maxToolCalls(n)` for a cumulative **emitted**-call budget (skipped calls still count). +- `agentLoopStrategy?` - Strategy for agent loops (default: `maxIterations(5)`). Strategies receive `{ iterationCount, finishReason, messages, toolCallCount, lastTurnToolCallCount }` and run between model turns. Iterations are model turns, not tool calls — for tool-call budgets use middleware (`onBeforeToolCall` + `onShouldContinue`); see [Tool-call budgets](../chat/agentic-cycle#tool-call-budgets-middleware-recipe). +- `middleware?` - Array of chat middleware. Use `onShouldContinue` / `onBeforeToolCall` for app-owned tool budgets. - `abortController?` - AbortController for cancellation - `modelOptions?` - Provider-native model options. This is where sampling parameters live — `temperature`, `top_p`/`topP`, and the provider's token-limit key (`max_output_tokens`, `max_tokens`, `maxOutputTokens`, …) — under each provider's canonical name, rather than as generic root-level props. See [Moving Sampling Options into modelOptions](../migration/sampling-options-to-model-options). (Renamed from `providerOptions`.) - `threadId?` - AG-UI thread identifier propagated into `RUN_STARTED` events for run correlation @@ -315,7 +315,7 @@ A merged tool record suitable for `chat({ tools })`. ## `maxIterations(count)` -Creates an agent loop strategy that limits **model turns** (iterations), not tool calls. One turn can still emit many parallel tool calls — use `maxToolCalls` and/or `maxToolCallsPerTurn` for tool-call budgets. +Creates an agent loop strategy that limits **model turns** (iterations), not tool calls. One turn can still emit many parallel tool calls — use middleware for tool-call budgets ([recipe](../chat/agentic-cycle#tool-call-budgets-middleware-recipe)). ```typescript import { chat, maxIterations } from "@tanstack/ai"; @@ -336,35 +336,6 @@ const stream = chat({ An `AgentLoopStrategy` function. -## `maxToolCalls(count)` - -Creates an agent loop strategy that continues while cumulative **emitted** tool calls are below `count` (including calls skipped by `maxToolCallsPerTurn`). - -Strategies only run between turns — the turn that crosses the limit is not truncated, so the final count may exceed `count` unless you also set `maxToolCallsPerTurn`. Pair with `chat({ maxToolCallsPerTurn })` to cap parallel fan-out inside a single turn. - -```typescript -import { chat, combineStrategies, maxIterations, maxToolCalls } from "@tanstack/ai"; -import { openaiText } from "@tanstack/ai-openai"; - -const stream = chat({ - adapter: openaiText("gpt-5.2"), - messages: [{ role: "user", content: "Hello!" }], - maxToolCallsPerTurn: 10, - agentLoopStrategy: combineStrategies([ - maxIterations(20), - maxToolCalls(20), - ]), -}); -``` - -### Parameters - -- `count` - Maximum cumulative tool calls across the run - -### Returns - -An `AgentLoopStrategy` function. - ## Types ### `ModelMessage` diff --git a/docs/chat/agentic-cycle.md b/docs/chat/agentic-cycle.md index 70ebd10f6..2b3ddd07d 100644 --- a/docs/chat/agentic-cycle.md +++ b/docs/chat/agentic-cycle.md @@ -159,17 +159,17 @@ The loop continues only while the model's finish reason is `tool_calls` (with pe By default the loop is bounded by `maxIterations(5)` — after five **model turns** it stops even if the model would keep calling tools. Override this with the `agentLoopStrategy` option. -> **Iterations ≠ tool calls.** One model turn can emit many parallel tool calls. `maxIterations` only bounds turns. Use `maxToolCalls(n)` for a cumulative **emitted**-call budget (stops further turns once the count is reached; the crossing turn is not truncated), and `maxToolCallsPerTurn` to cap how many of those parallel calls **execute** in a single turn or pending/resume batch (strategies only run *between* turns, so without a per-turn cap a single runaway turn can still fan out unbounded). Skipped calls still count toward `maxToolCalls`. +Other built-in strategies: + +- **`untilFinishReason([...])`** — continue until the model returns one of the given finish reasons (e.g. `untilFinishReason(["stop", "length"])`). +- **`combineStrategies([...])`** — combine multiple strategies with AND logic; the loop continues only while every strategy agrees. + +A strategy is just a function that receives `{ iterationCount, finishReason, messages, toolCallCount, lastTurnToolCallCount }` and returns `true` to allow another iteration or `false` to stop, so you can also write your own: ```typescript -import { - chat, - combineStrategies, - maxIterations, - maxToolCalls, - toServerSentEventsResponse, -} from "@tanstack/ai"; +import { chat, combineStrategies, maxIterations, toServerSentEventsResponse } from "@tanstack/ai"; import { openaiText } from "@tanstack/ai-openai"; +import type { AgentLoopState } from "@tanstack/ai"; import { getWeather, getClothingAdvice } from "./tools"; export async function POST(request: Request) { @@ -178,41 +178,87 @@ export async function POST(request: Request) { adapter: openaiText("gpt-5.5"), messages, tools: [getWeather, getClothingAdvice], - maxToolCallsPerTurn: 10, // cap parallel fan-out inside one turn agentLoopStrategy: combineStrategies([ - maxIterations(20), // max model turns - maxToolCalls(20), // max tool calls across the run + maxIterations(10), + ({ messages }: AgentLoopState) => messages.length < 100, ]), }); return toServerSentEventsResponse(stream); } ``` -Other built-in strategies: +### Tool-call budgets (middleware recipe) -- **`maxToolCalls(n)`** — continue while cumulative emitted tool calls are below `n` (including per-turn skips; not model turns). -- **`untilFinishReason([...])`** — continue until the model returns one of the given finish reasons (e.g. `untilFinishReason(["stop", "length"])`). -- **`combineStrategies([...])`** — combine multiple strategies with AND logic; the loop continues only while every strategy agrees. +> **Iterations ≠ tool calls.** One model turn can emit many parallel tool calls. `maxIterations` only bounds **model turns**. Strategies run *between* turns, so without a per-turn cap a single runaway turn can still fan out unbounded. -A strategy is just a function that receives `{ iterationCount, finishReason, messages, toolCallCount, lastTurnToolCallCount }` and returns `true` to allow another iteration or `false` to stop, so you can also write your own: +There is no built-in `maxToolCalls` strategy. Cap tools with middleware: + +- **`onBeforeToolCall`** — skip excess calls inside one turn (`maxPerTurn`) +- **`onShouldContinue`** — stop further turns once cumulative **emitted** tools hit a budget (`max`); skipped calls still count toward `toolCallCount` ```typescript -import { chat, combineStrategies, maxIterations, toServerSentEventsResponse } from "@tanstack/ai"; +import { + chat, + maxIterations, + toServerSentEventsResponse, + type ChatMiddleware, +} from "@tanstack/ai"; import { openaiText } from "@tanstack/ai-openai"; -import type { AgentLoopState } from "@tanstack/ai"; import { getWeather, getClothingAdvice } from "./tools"; +/** App-owned policy — not a library export. */ +function toolCallBudget(options: { + max?: number; + maxPerTurn?: number; +}): ChatMiddleware { + const { max, maxPerTurn } = options; + let perTurn = 0; + + return { + name: "tool-call-budget", + onIteration() { + perTurn = 0; + }, + // Fresh per-turn budget for pending/resume batches (no onIteration). + onToolPhaseComplete() { + perTurn = 0; + }, + onBeforeToolCall() { + if (maxPerTurn == null) return undefined; + perTurn += 1; + if (perTurn > maxPerTurn) { + return { + type: "skip", + result: { + error: `Skipped: exceeded maxToolCallsPerTurn (${maxPerTurn})`, + }, + }; + } + return undefined; + }, + onShouldContinue(_ctx, state) { + if (max != null && state.toolCallCount >= max) return false; + return undefined; + }, + }; +} + export async function POST(request: Request) { const { messages } = await request.json(); const stream = chat({ adapter: openaiText("gpt-5.5"), messages, tools: [getWeather, getClothingAdvice], - agentLoopStrategy: combineStrategies([ - maxIterations(10), - ({ messages }: AgentLoopState) => messages.length < 100, - ]), + agentLoopStrategy: maxIterations(20), // model turns + middleware: [ + toolCallBudget({ + maxPerTurn: 10, // cap parallel fan-out inside one turn + max: 20, // stop further turns once cumulative emitted tools hit 20 + }), + ], }); return toServerSentEventsResponse(stream); } ``` + +Place this **before** `toolCacheMiddleware` so over-budget skips win over cache hits. See [`onShouldContinue`](../advanced/middleware#onshouldcontinue) for the hook contract. diff --git a/docs/config.json b/docs/config.json index 55a365282..ce7dee3a9 100644 --- a/docs/config.json +++ b/docs/config.json @@ -156,7 +156,7 @@ "label": "Agentic Cycle", "to": "chat/agentic-cycle", "addedAt": "2026-04-15", - "updatedAt": "2026-07-20" + "updatedAt": "2026-07-21" }, { "label": "Streaming", @@ -337,12 +337,13 @@ "label": "Middleware", "to": "advanced/middleware", "addedAt": "2026-04-15", - "updatedAt": "2026-06-15" + "updatedAt": "2026-07-21" }, { "label": "Built-in Middleware", "to": "advanced/built-in-middleware", - "addedAt": "2026-06-03" + "addedAt": "2026-06-03", + "updatedAt": "2026-07-21" }, { "label": "OpenTelemetry", @@ -502,7 +503,7 @@ "label": "@tanstack/ai", "to": "api/ai", "addedAt": "2026-04-15", - "updatedAt": "2026-07-20" + "updatedAt": "2026-07-21" }, { "label": "@tanstack/ai-client", diff --git a/docs/reference/functions/combineStrategies.md b/docs/reference/functions/combineStrategies.md index 5d62a5413..7796fa58f 100644 --- a/docs/reference/functions/combineStrategies.md +++ b/docs/reference/functions/combineStrategies.md @@ -38,7 +38,6 @@ const stream = chat({ tools: [weatherTool], agentLoopStrategy: combineStrategies([ maxIterations(10), - maxToolCalls(20), ({ messages }) => messages.length < 100, ]), }); diff --git a/docs/reference/functions/maxIterations.md b/docs/reference/functions/maxIterations.md index d31183c07..5d1cdb739 100644 --- a/docs/reference/functions/maxIterations.md +++ b/docs/reference/functions/maxIterations.md @@ -14,9 +14,9 @@ Defined in: [packages/ai/src/activities/chat/agent-loop-strategies.ts:25](https: Creates a strategy that continues for a maximum number of **model turns** (iterations), not tool calls. -One iteration can still emit many parallel tool calls. Prefer -[maxToolCalls](maxToolCalls.md) (and optionally `maxToolCallsPerTurn` on `chat()`) -when you need a tool-call budget. +One iteration can still emit many parallel tool calls. Prefer middleware +with `onBeforeToolCall` / `onShouldContinue` when you need a tool-call +budget (see the Agentic Cycle docs recipe). ## Parameters diff --git a/docs/reference/functions/maxToolCalls.md b/docs/reference/functions/maxToolCalls.md deleted file mode 100644 index d8e0c2b12..000000000 --- a/docs/reference/functions/maxToolCalls.md +++ /dev/null @@ -1,53 +0,0 @@ ---- -id: maxToolCalls -title: maxToolCalls ---- - -# Function: maxToolCalls() - -```ts -function maxToolCalls(max): AgentLoopStrategy; -``` - -Defined in: [packages/ai/src/activities/chat/agent-loop-strategies.ts:59](https://github.com/TanStack/ai/blob/main/packages/ai/src/activities/chat/agent-loop-strategies.ts#L59) - -Creates a strategy that continues while `toolCallCount < max`. - -Unlike [maxIterations](maxIterations.md) (which counts model turns), this bounds -**emitted** tool calls counted during the run (including ones skipped by -`maxToolCallsPerTurn`). Strategies only run between turns, so the turn that -crosses `max` is not truncated — the final count (and executions, unless -`maxToolCallsPerTurn` is set) may exceed `max`. Pair with -`chat({ maxToolCallsPerTurn })` to also cap parallel fan-out inside a single -turn. - -## Parameters - -### max - -`number` - -Maximum cumulative emitted tool calls before stopping further turns - -## Returns - -[`AgentLoopStrategy`](../type-aliases/AgentLoopStrategy.md) - -AgentLoopStrategy that returns true while `toolCallCount < max` - -## Example - -```typescript -import { chat, combineStrategies, maxIterations, maxToolCalls } from '@tanstack/ai' - -const stream = chat({ - adapter: openaiText('gpt-4o'), - messages: [...], - tools: [weatherTool], - maxToolCallsPerTurn: 10, - agentLoopStrategy: combineStrategies([ - maxIterations(20), - maxToolCalls(20), - ]), -}) -``` diff --git a/docs/reference/index.md b/docs/reference/index.md index 898a2702a..73f19eec0 100644 --- a/docs/reference/index.md +++ b/docs/reference/index.md @@ -294,7 +294,6 @@ title: "@tanstack/ai" - [isProviderExecutedToolCall](functions/isProviderExecutedToolCall.md) - [isStandardSchema](functions/isStandardSchema.md) - [maxIterations](functions/maxIterations.md) -- [maxToolCalls](functions/maxToolCalls.md) - [mergeAgentTools](functions/mergeAgentTools.md) - [modelMessagesToUIMessages](functions/modelMessagesToUIMessages.md) - [modelMessageToUIMessage](functions/modelMessageToUIMessage.md) diff --git a/docs/reference/interfaces/AgentLoopState.md b/docs/reference/interfaces/AgentLoopState.md index da5c9d732..ff5c224ce 100644 --- a/docs/reference/interfaces/AgentLoopState.md +++ b/docs/reference/interfaces/AgentLoopState.md @@ -72,6 +72,6 @@ toolCallCount: number; Defined in: [packages/ai/src/types.ts:846](https://github.com/TanStack/ai/blob/main/packages/ai/src/types.ts#L846) Cumulative tool calls counted so far in this run (model-emitted during the -agent loop, including ones skipped by `maxToolCallsPerTurn`, and pending -tools from the inbound message list when resumed). Not a recount of full -message history; not model turns. +agent loop, including ones skipped by middleware, and pending tools from +the inbound message list when resumed). Not a recount of full message +history; not model turns. diff --git a/docs/reference/interfaces/TextOptions.md b/docs/reference/interfaces/TextOptions.md index 08364f1f7..b1da1eaea 100644 --- a/docs/reference/interfaces/TextOptions.md +++ b/docs/reference/interfaces/TextOptions.md @@ -157,32 +157,6 @@ chunk received, and `logger.errors()` in catch blocks. *** -### maxToolCallsPerTurn? - -```ts -optional maxToolCallsPerTurn: number; -``` - -Defined in: [packages/ai/src/types.ts:919](https://github.com/TanStack/ai/blob/main/packages/ai/src/types.ts#L919) - -Maximum number of tool calls to **execute** from a single model turn (or -pending/resume batch). `0` skips all execution for that batch. - -Models can emit many parallel tool calls in one turn. `agentLoopStrategy` -(including `maxIterations` / `maxToolCalls`) is only evaluated between -turns, so without this cap a single runaway turn can still execute an -unbounded fan-out. - -When set, only the first `maxToolCallsPerTurn` calls are executed; the -remainder receive error tool results so the message history stays -consistent. Unset means no per-turn execution cap. Must be a non-negative -finite number when set. - -Pair with the `maxToolCalls(n)` strategy for a cumulative **emitted**-call -budget across the run (skipped calls still count toward that budget). - -*** - ### messages ```ts diff --git a/packages/ai/skills/ai-core/middleware/SKILL.md b/packages/ai/skills/ai-core/middleware/SKILL.md index 65c6bf7b8..2d90ccd74 100644 --- a/packages/ai/skills/ai-core/middleware/SKILL.md +++ b/packages/ai/skills/ai-core/middleware/SKILL.md @@ -57,6 +57,7 @@ Every hook receives a `ChatMiddlewareContext` as its first argument, which provi | `onStructuredOutputConfig` | Once at the structured-output boundary (only when `chat({ outputSchema })`) | `StructuredOutputMiddlewareConfig` (return partial) | | `onStart` | Once after initial `onConfig` | none | | `onIteration` | Start of each agent loop iteration | `IterationInfo` | +| `onShouldContinue` | Whether to start another agent-loop iteration (AND with strategy; `false` stops) | `AgentLoopState` | | `onChunk` | Every streamed chunk | `StreamChunk` (return void/chunk/chunk[]/null) | | `onBeforeToolCall` | Before each tool executes | `ToolCallHookContext` (return decision or void) | | `onAfterToolCall` | After each tool executes | `AfterToolCallInfo` | @@ -346,6 +347,52 @@ const stream = chat({ | `onUsage` | Sequential | All run in order | | `onFinish/onAbort/onError` | Sequential | All run in order | +## Pattern: tool-call budget (app-owned) + +Not a built-in. Cap fan-out with `onBeforeToolCall` skip + `onShouldContinue`. +See `docs/chat/agentic-cycle.md` ("Tool-call budgets"). + +```typescript +import { chat, maxIterations, type ChatMiddleware } from '@tanstack/ai' + +function toolCallBudget(opts: { + max?: number + maxPerTurn?: number +}): ChatMiddleware { + let perTurn = 0 + return { + onIteration: () => { + perTurn = 0 + }, + onToolPhaseComplete: () => { + perTurn = 0 + }, + onBeforeToolCall: () => { + if (opts.maxPerTurn == null) return undefined + if (++perTurn > opts.maxPerTurn) { + return { + type: 'skip', + result: { + error: `Skipped: exceeded maxToolCallsPerTurn (${opts.maxPerTurn})`, + }, + } + } + return undefined + }, + onShouldContinue: (_ctx, state) => + opts.max != null && state.toolCallCount >= opts.max ? false : undefined, + } +} + +chat({ + adapter, + messages, + tools: [weatherTool], + agentLoopStrategy: maxIterations(20), + middleware: [toolCallBudget({ maxPerTurn: 10, max: 20 })], +}) +``` + ## Built-in: toolCacheMiddleware Caches tool call results by name + arguments. Import from `@tanstack/ai/middlewares`: diff --git a/packages/ai/skills/ai-core/tool-calling/SKILL.md b/packages/ai/skills/ai-core/tool-calling/SKILL.md index 62c605de6..9346c2bee 100644 --- a/packages/ai/skills/ai-core/tool-calling/SKILL.md +++ b/packages/ai/skills/ai-core/tool-calling/SKILL.md @@ -375,8 +375,8 @@ export async function POST(request: Request) { adapter: openaiText('gpt-5.5'), messages, tools: [getProducts, compareProducts], - // maxIterations bounds model turns, not tool calls. Prefer maxToolCalls - // (and maxToolCallsPerTurn) when you need a tool-call budget. + // maxIterations bounds model turns, not tool calls. For tool budgets, + // use middleware onBeforeToolCall + onShouldContinue (see agentic-cycle docs). agentLoopStrategy: maxIterations(20), }) return toServerSentEventsResponse(stream) diff --git a/packages/ai/src/activities/chat/agent-loop-strategies.ts b/packages/ai/src/activities/chat/agent-loop-strategies.ts index 442ebe01c..5777710ad 100644 --- a/packages/ai/src/activities/chat/agent-loop-strategies.ts +++ b/packages/ai/src/activities/chat/agent-loop-strategies.ts @@ -4,9 +4,10 @@ import type { AgentLoopStrategy } from '../../types' * Creates a strategy that continues for a maximum number of **model turns** * (iterations), not tool calls. * - * One iteration can still emit many parallel tool calls. Prefer - * {@link maxToolCalls} (and optionally `maxToolCallsPerTurn` on `chat()`) - * when you need a tool-call budget. + * One iteration can still emit many parallel tool calls. For a tool-call + * budget, use middleware with `onBeforeToolCall` (per-turn cap) and + * `onShouldContinue` (cumulative run budget) — see the docs recipe under + * Agentic Cycle. * * @param max - Maximum number of model turns to allow * @returns AgentLoopStrategy that stops after max iterations @@ -26,40 +27,6 @@ export function maxIterations(max: number): AgentLoopStrategy { return ({ iterationCount }) => iterationCount < max } -/** - * Creates a strategy that continues while `toolCallCount < max`. - * - * Unlike {@link maxIterations} (which counts model turns), this bounds - * **emitted** tool calls counted during the run (including ones skipped by - * `maxToolCallsPerTurn`). Strategies only run between turns, so the turn that - * crosses `max` is not truncated — the final count (and executions, unless - * `maxToolCallsPerTurn` is set) may exceed `max`. Pair with - * `chat({ maxToolCallsPerTurn })` to also cap parallel fan-out inside a single - * turn. - * - * @param max - Maximum cumulative emitted tool calls before stopping further turns - * @returns AgentLoopStrategy that returns true while `toolCallCount < max` - * - * @example - * ```typescript - * import { chat, combineStrategies, maxIterations, maxToolCalls } from '@tanstack/ai' - * - * const stream = chat({ - * adapter: openaiText('gpt-4o'), - * messages: [...], - * tools: [weatherTool], - * maxToolCallsPerTurn: 10, - * agentLoopStrategy: combineStrategies([ - * maxIterations(20), - * maxToolCalls(20), - * ]), - * }) - * ``` - */ -export function maxToolCalls(max: number): AgentLoopStrategy { - return ({ toolCallCount }) => toolCallCount < max -} - /** * Creates a strategy that continues until a specific finish reason is encountered * @@ -99,7 +66,7 @@ export function untilFinishReason( * All strategies must return true to continue * * @param strategies - Array of strategies to combine - * @returns AgentLoopStrategy that continues only if all strategies return true + * @returns AgentLoopStrategy that continues only if all strategies agree * * @example * ```typescript @@ -110,7 +77,6 @@ export function untilFinishReason( * tools: [weatherTool], * agentLoopStrategy: combineStrategies([ * maxIterations(10), - * maxToolCalls(20), * ({ messages }) => messages.length < 100, * ]), * }); diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index deee12e8f..2d0edb220 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -241,11 +241,6 @@ export interface TextActivityOptions< abortController?: TextOptions['abortController'] /** Strategy for controlling the agent loop */ agentLoopStrategy?: TextOptions['agentLoopStrategy'] - /** - * Cap how many tool calls from a single model turn are executed. - * Excess calls receive error results. See {@link TextOptions.maxToolCallsPerTurn}. - */ - maxToolCallsPerTurn?: TextOptions['maxToolCallsPerTurn'] /** * Optional configuration for lazy-tool discovery (tools marked `lazy: true`). * Tunes how much of each lazy tool's description appears in the discovery @@ -498,23 +493,6 @@ interface TextEngineConfig< type ToolPhaseResult = 'continue' | 'stop' | 'wait' type CyclePhase = 'processText' | 'executeToolCalls' -/** - * Validate and normalize `maxToolCallsPerTurn`. - * Unset → unlimited. `0` → execute none. Negatives / non-finite → throw - * (Array#slice treats negatives as "from end", which is not a useful cap). - */ -function resolveMaxToolCallsPerTurn( - cap: number | undefined, -): number | undefined { - if (cap == null) return undefined - if (!Number.isFinite(cap) || cap < 0) { - throw new Error( - `maxToolCallsPerTurn must be a non-negative finite number, got ${cap}`, - ) - } - return Math.floor(cap) -} - /** * Combine two optional AbortSignals into one that aborts when either does. * Returns the other signal directly when one is absent or already aborted. @@ -582,7 +560,6 @@ class TextEngine< private earlyTermination = false private toolPhase: ToolPhaseResult = 'continue' private cyclePhase: CyclePhase = 'processText' - private readonly maxToolCallsPerTurn: number | undefined // Client state extracted from initial messages (before conversion to ModelMessage) private readonly initialApprovals: Map private readonly initialClientToolResults: Map @@ -648,9 +625,6 @@ class TextEngine< this.systemPrompts = config.params.systemPrompts || [] this.loopStrategy = config.params.agentLoopStrategy || maxIterationsStrategy(5) - this.maxToolCallsPerTurn = resolveMaxToolCallsPerTurn( - config.params.maxToolCallsPerTurn, - ) this.initialMessageCount = config.params.messages.length // Extract client state (approvals, client tool results) from original messages BEFORE conversion @@ -902,7 +876,7 @@ class TextEngine< } this.endCycle() - } while (this.shouldContinue()) + } while (await this.shouldContinue()) } this.logger.agentLoop('run finished', { @@ -1354,14 +1328,13 @@ class TextEngine< const finishEvent = this.createSyntheticFinishedEvent() - // Same fan-out budget as live model turns (seeded history / resume). // Count is deduped so wait→resume after a live turn does not double-count. - const { toExecute: budgetedToolCalls, skippedResults } = - this.applyToolCallBudget(pendingToolCalls) + // Per-turn execution caps are app middleware via onBeforeToolCall skip. + this.recordToolCalls(pendingToolCalls) // Handle undiscovered lazy tool calls with self-correcting error messages const undiscoveredLazyResults: Array = [] - const executablePendingCalls = budgetedToolCalls.filter((tc) => { + const executablePendingCalls = pendingToolCalls.filter((tc) => { if (this.lazyToolManager.isUndiscoveredLazyTool(tc.function.name)) { undiscoveredLazyResults.push({ toolCallId: tc.id, @@ -1378,9 +1351,10 @@ class TextEngine< return true }) - // Non-executed outcomes (undiscovered lazy + per-turn fan-out skips). - // Emitted after executed results so the stream prefers real results first. - const deferredErrorResults = [...undiscoveredLazyResults, ...skippedResults] + // Non-executed outcomes (undiscovered lazy). Emitted after executed + // results so the stream prefers real results first. Per-turn skips are + // produced by middleware via onBeforeToolCall and appear in execution results. + const deferredErrorResults = [...undiscoveredLazyResults] // Build args lookup so buildToolResultChunks can emit TOOL_CALL_START + // TOOL_CALL_ARGS before TOOL_CALL_END during continuation re-executions. @@ -1524,15 +1498,15 @@ class TextEngine< return } - // Count every model-emitted tool call (including ones we may skip). - const { toExecute: budgetedToolCalls, skippedResults } = - this.applyToolCallBudget(toolCalls) + // Count every model-emitted tool call. Per-turn execution caps are app + // middleware via onBeforeToolCall skip. + this.recordToolCalls(toolCalls) this.addAssistantToolCallMessage(toolCalls) // Handle undiscovered lazy tool calls with self-correcting error messages const undiscoveredLazyResults: Array = [] - const executableToolCalls = budgetedToolCalls.filter((tc) => { + const executableToolCalls = toolCalls.filter((tc) => { if (this.lazyToolManager.isUndiscoveredLazyTool(tc.function.name)) { undiscoveredLazyResults.push({ toolCallId: tc.id, @@ -1549,13 +1523,13 @@ class TextEngine< return true }) - // Non-executed outcomes (undiscovered lazy + per-turn fan-out skips). - // Emitted after executed results so the stream prefers real results first. - const deferredErrorResults = [...undiscoveredLazyResults, ...skippedResults] + // Non-executed outcomes (undiscovered lazy). Per-turn skips come from + // middleware and appear in execution results. + const deferredErrorResults = [...undiscoveredLazyResults] if (executableToolCalls.length === 0) { - // All tool calls were undiscovered lazy tools and/or skipped by the - // per-turn fan-out cap — errors emitted, continue loop (strategy may stop). + // All tool calls were undiscovered lazy tools — errors emitted, continue + // loop (strategy / onShouldContinue may stop). if (deferredErrorResults.length > 0) { for (const chunk of this.buildToolResultChunks( deferredErrorResults, @@ -2010,35 +1984,48 @@ class TextEngine< } as RunFinishedEvent } - private shouldContinue(): boolean { + private async shouldContinue(): Promise { + // Always enter the tool-execution half-cycle after a model turn. if (this.cyclePhase === 'executeToolCalls') { return true } - return ( - this.loopStrategy({ - iterationCount: this.iterationCount, - messages: this.messages, - finishReason: this.lastFinishReason, - toolCallCount: this.toolCallCount, - lastTurnToolCallCount: this.lastTurnToolCallCount, - }) && this.toolPhase === 'continue' - ) + const state = { + iterationCount: this.iterationCount, + messages: this.messages, + finishReason: this.lastFinishReason, + toolCallCount: this.toolCallCount, + lastTurnToolCallCount: this.lastTurnToolCallCount, + } + + // Evaluate strategy and middleware even when toolPhase is 'stop' so + // observers can see final counters; AND with toolPhase at the end. + if (!this.loopStrategy(state)) { + return false + } + + if ( + !(await this.middlewareRunner.runOnShouldContinue( + this.middlewareCtx, + state, + )) + ) { + return false + } + + return this.toolPhase === 'continue' } /** - * Record tool calls (deduped by id) and return the subset that should be - * executed after applying `maxToolCallsPerTurn`. Excess calls get synthetic - * error results so every tool_call still has a matching result. + * Record tool calls (deduped by id) toward `toolCallCount` / + * `lastTurnToolCallCount` for strategies and middleware `onShouldContinue`. * * Used for both live model turns and pending/resume batches. IDs already * counted in this run (e.g. wait→resume after a live turn) are not - * re-added to `toolCallCount`. + * re-added to `toolCallCount`. Per-turn execution caps are app middleware + * (`onBeforeToolCall` skip), not engine policy. */ - private applyToolCallBudget(toolCalls: Array): { - toExecute: Array - skippedResults: Array - } { + private recordToolCalls(toolCalls: Array): void { this.lastTurnToolCallCount = toolCalls.length let newlyCounted = 0 for (const tc of toolCalls) { @@ -2048,34 +2035,6 @@ class TextEngine< } } this.toolCallCount += newlyCounted - - const cap = this.maxToolCallsPerTurn - if (cap == null || toolCalls.length <= cap) { - return { toExecute: toolCalls, skippedResults: [] } - } - - this.logger.agentLoop( - `maxToolCallsPerTurn=${cap} skipped=${toolCalls.length - cap}`, - { - maxToolCallsPerTurn: cap, - emitted: toolCalls.length, - skipped: toolCalls.length - cap, - }, - ) - - const toExecute = toolCalls.slice(0, cap) - const skippedResults: Array = toolCalls - .slice(cap) - .map((tc) => ({ - toolCallId: tc.id, - toolName: tc.function.name, - result: { - error: `Skipped: exceeded maxToolCallsPerTurn (${cap})`, - }, - state: 'output-error' as const, - })) - - return { toExecute, skippedResults } } private isAborted(): boolean { diff --git a/packages/ai/src/activities/chat/middleware/compose.ts b/packages/ai/src/activities/chat/middleware/compose.ts index 9d449ad00..8badc579b 100644 --- a/packages/ai/src/activities/chat/middleware/compose.ts +++ b/packages/ai/src/activities/chat/middleware/compose.ts @@ -1,5 +1,5 @@ import { aiEventClient } from '@tanstack/ai-event-client' -import type { StreamChunk } from '../../../types' +import type { AgentLoopState, StreamChunk } from '../../../types' import type { InternalLogger } from '../../../logger/internal-logger' import type { AbortInfo, @@ -595,6 +595,46 @@ export class MiddlewareRunner { } } + /** + * Run onShouldContinue through middleware in order (AND semantics). + * Any explicit `false` stops further iterations; `true` / void / undefined pass. + * Called after `agentLoopStrategy` has already approved continuation. + */ + async runOnShouldContinue( + ctx: ChatMiddlewareContext, + state: AgentLoopState, + ): Promise { + for (const mw of this.middlewares) { + if (mw.onShouldContinue) { + const skip = shouldSkipInstrumentation(mw) + const start = Date.now() + const result = await mw.onShouldContinue(ctx, state) + if (!skip) { + this.logger.middleware( + `hook=onShouldContinue middleware=${mw.name ?? 'unnamed'}`, + { + middleware: mw.name ?? 'unnamed', + hook: 'onShouldContinue', + result, + }, + ) + aiEventClient.emit('middleware:hook:executed', { + ...instrumentCtx(ctx), + middlewareName: mw.name || 'unnamed', + hookName: 'onShouldContinue', + iteration: ctx.iteration, + duration: Date.now() - start, + hasTransform: result === false, + }) + } + if (result === false) { + return false + } + } + } + return true + } + /** * Run onToolPhaseComplete on all middleware in order. * Called after all tool calls in an iteration have been processed. diff --git a/packages/ai/src/activities/chat/middleware/types.ts b/packages/ai/src/activities/chat/middleware/types.ts index 74502b5c7..c908ea2fa 100644 --- a/packages/ai/src/activities/chat/middleware/types.ts +++ b/packages/ai/src/activities/chat/middleware/types.ts @@ -1,4 +1,5 @@ import type { + AgentLoopState, JSONSchema, ModelMessage, StreamChunk, @@ -506,6 +507,25 @@ export interface ChatMiddleware { info: IterationInfo, ) => void | Promise + /** + * Called when the engine is deciding whether to start another agent-loop + * iteration (after a tool phase or between model turns). + * + * Return `false` to stop further iterations. Return `true`, `void`, or + * `undefined` to allow continuation. Combined with AND semantics across + * middleware and with `agentLoopStrategy` — any `false` stops the loop. + * + * Does not abort the run: the stream finishes normally with the current + * messages. Use `ctx.abort()` only when you need a hard abort. + * + * Receives the same {@link AgentLoopState} passed to strategies + * (`iterationCount`, `toolCallCount`, `lastTurnToolCallCount`, etc.). + */ + onShouldContinue?: ( + ctx: ChatMiddlewareContext, + state: AgentLoopState, + ) => boolean | void | Promise + /** * Called for every chunk yielded by chat(). * Can observe, transform, expand, or drop chunks. diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 16a07fede..9f4c48655 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -99,7 +99,6 @@ export { brandProviderTool } from './tools/provider-tool' // Agent loop strategies export { maxIterations, - maxToolCalls, untilFinishReason, combineStrategies, } from './activities/chat/agent-loop-strategies' diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 5a655ed20..a56594ec2 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -846,13 +846,13 @@ export interface AgentLoopState { finishReason: string | null /** * Cumulative tool calls counted so far in this run (model-emitted during the - * agent loop, including ones skipped by `maxToolCallsPerTurn`, and pending - * tools from the inbound message list when resumed). Not a recount of full - * message history; not model turns. + * agent loop, including ones skipped by middleware, and pending tools from + * the inbound message list when resumed). Not a recount of full message + * history; not model turns. */ toolCallCount: number /** - * Tool calls in the most recent budgeted batch — a live model turn or a + * Tool calls in the most recent batch — a live model turn or a * pending/resume batch (0 when the last phase produced no tool calls). */ lastTurnToolCallCount: number @@ -868,7 +868,7 @@ export interface AgentLoopState { * ```typescript * // Continue for up to 5 iterations (model turns, not tool calls) * const strategy: AgentLoopStrategy = ({ iterationCount }) => iterationCount < 5; - * // Cap total tool calls across the run + * // Cap total tool calls across the run (or use middleware onShouldContinue) * const byTools: AgentLoopStrategy = ({ toolCallCount }) => toolCallCount < 20; * ``` */ @@ -906,24 +906,6 @@ export interface TextOptions< */ systemPrompts?: Array agentLoopStrategy?: AgentLoopStrategy - /** - * Maximum number of tool calls to **execute** from a single model turn (or - * pending/resume batch). `0` skips all execution for that batch. - * - * Models can emit many parallel tool calls in one turn. `agentLoopStrategy` - * (including `maxIterations` / `maxToolCalls`) is only evaluated between - * turns, so without this cap a single runaway turn can still execute an - * unbounded fan-out. - * - * When set, only the first `maxToolCallsPerTurn` calls are executed; the - * remainder receive error tool results so the message history stays - * consistent. Unset means no per-turn execution cap. Must be a non-negative - * finite number when set. - * - * Pair with the `maxToolCalls(n)` strategy for a cumulative **emitted**-call - * budget across the run (skipped calls still count toward that budget). - */ - maxToolCallsPerTurn?: number /** * Optional configuration for lazy-tool discovery (tools marked `lazy: true`). * Tunes how much of each lazy tool's description appears in the discovery diff --git a/packages/ai/tests/agent-loop-strategies.test.ts b/packages/ai/tests/agent-loop-strategies.test.ts index baea1c0fe..d71f384b7 100644 --- a/packages/ai/tests/agent-loop-strategies.test.ts +++ b/packages/ai/tests/agent-loop-strategies.test.ts @@ -1,7 +1,6 @@ import { describe, it, expect } from 'vitest' import { maxIterations, - maxToolCalls, untilFinishReason, combineStrategies, } from '../src/activities/chat/agent-loop-strategies' @@ -49,42 +48,6 @@ describe('Agent Loop Strategies', () => { }) }) - describe('maxToolCalls', () => { - it('should continue when below max tool calls', () => { - const strategy = maxToolCalls(20) - - expect(strategy(createState({ toolCallCount: 0 }))).toBe(true) - expect(strategy(createState({ toolCallCount: 10 }))).toBe(true) - expect(strategy(createState({ toolCallCount: 19 }))).toBe(true) - }) - - it('should stop when reaching max tool calls', () => { - const strategy = maxToolCalls(20) - - expect(strategy(createState({ toolCallCount: 20 }))).toBe(false) - expect(strategy(createState({ toolCallCount: 21 }))).toBe(false) - }) - - it('should ignore iterationCount (counts tools, not turns)', () => { - const strategy = maxToolCalls(5) - - // Many iterations but few tools — still continue - expect( - strategy(createState({ iterationCount: 100, toolCallCount: 4 })), - ).toBe(true) - // Few iterations but many tools — stop - expect( - strategy(createState({ iterationCount: 1, toolCallCount: 5 })), - ).toBe(false) - }) - - it('should work with max = 0 (never allow tool calls)', () => { - const strategy = maxToolCalls(0) - - expect(strategy(createState({ toolCallCount: 0 }))).toBe(false) - }) - }) - describe('untilFinishReason', () => { it('should continue on first iteration even with matching finish reason', () => { const strategy = untilFinishReason(['stop']) diff --git a/packages/ai/tests/chat.test.ts b/packages/ai/tests/chat.test.ts index 6214e2e2f..e7b877dde 100644 --- a/packages/ai/tests/chat.test.ts +++ b/packages/ai/tests/chat.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it, vi } from 'vitest' import { chat, createChatOptions } from '../src/activities/chat/index' import { DISCOVERY_TOOL_NAME } from '../src/activities/chat/tools/lazy-tool-manager' import { EventType } from '../src/types' +import type { ChatMiddleware } from '../src/activities/chat/middleware/types' import type { StreamChunk, Tool } from '../src/types' import { chunk, @@ -12,6 +13,44 @@ import { clientTool, } from './test-utils' +/** + * App-owned tool budget middleware (docs recipe). Not a library export — + * used here to lock hook behavior for #964 fan-out budgets. + */ +function toolCallBudget(options: { + max?: number + maxPerTurn?: number +}): ChatMiddleware { + const { max, maxPerTurn } = options + let perTurn = 0 + return { + name: 'tool-call-budget', + onIteration() { + perTurn = 0 + }, + onToolPhaseComplete() { + perTurn = 0 + }, + onBeforeToolCall() { + if (maxPerTurn == null) return undefined + perTurn += 1 + if (perTurn > maxPerTurn) { + return { + type: 'skip' as const, + result: { + error: `Skipped: exceeded maxToolCallsPerTurn (${maxPerTurn})`, + }, + } + } + return undefined + }, + onShouldContinue(_ctx, state) { + if (max != null && state.toolCallCount >= max) return false + return undefined + }, + } +} + /** Lazy server tool (has execute, lazy: true). */ function lazyServerTool(name: string, executeFn: (args: any) => any): Tool { return { @@ -1121,7 +1160,7 @@ describe('chat()', () => { expect(calls.length).toBe(2) }) - it('should respect maxToolCalls strategy (counts tools, not turns)', async () => { + it('should respect toolCallCount on strategy state (counts tools, not turns)', async () => { const executeSpy = vi.fn().mockReturnValue({ data: 'result' }) let callCount = 0 @@ -1201,7 +1240,7 @@ describe('chat()', () => { expect(seenLastTurn[seenLastTurn.length - 1]).toBe(0) }) - it('should cap parallel fan-out with maxToolCallsPerTurn', async () => { + it('should cap parallel fan-out with onBeforeToolCall budget middleware', async () => { const executeSpy = vi.fn().mockReturnValue({ data: 'result' }) const { adapter, calls } = createMockAdapter({ @@ -1228,7 +1267,7 @@ describe('chat()', () => { adapter, messages: [{ role: 'user', content: 'Fan out' }], tools: [serverTool('repeater', executeSpy)], - maxToolCallsPerTurn: 2, + middleware: [toolCallBudget({ maxPerTurn: 2 })], // Allow a follow-up turn so we can inspect tool results on the adapter agentLoopStrategy: (state) => state.iterationCount < 2, }) @@ -1260,14 +1299,14 @@ describe('chat()', () => { expect(toolMessages).toHaveLength(5) }) - it('counts skipped emissions toward maxToolCalls and stops further turns', async () => { + it('counts skipped emissions toward max and stops further turns via onShouldContinue', async () => { const executeSpy = vi.fn().mockReturnValue({ data: 'result' }) const seenCounts: Array = [] const { adapter, calls } = createMockAdapter({ chatStreamFn: () => { // Fat turn: 8 parallel tools. If we only counted executed (3), - // maxToolCalls(5) would allow another model turn. + // max: 5 would allow another model turn. return (async function* () { yield ev.runStarted() for (let i = 0; i < 8; i++) { @@ -1283,10 +1322,11 @@ describe('chat()', () => { adapter, messages: [{ role: 'user', content: 'Fan out' }], tools: [serverTool('repeater', executeSpy)], - maxToolCallsPerTurn: 3, + middleware: [toolCallBudget({ maxPerTurn: 3, max: 5 })], agentLoopStrategy: (state) => { seenCounts.push(state.toolCallCount) - return state.toolCallCount < 5 + // Strategy alone would continue (high iteration budget); middleware stops + return state.iterationCount < 10 }, }) @@ -1295,11 +1335,11 @@ describe('chat()', () => { expect(executeSpy).toHaveBeenCalledTimes(3) // All 8 emissions counted (not just 3 executed) expect(seenCounts).toContain(8) - // Strategy stops further model turns once count >= 5 + // Middleware onShouldContinue stops further model turns once count >= 5 expect(calls.length).toBe(1) }) - it('maxToolCallsPerTurn: 0 skips all tool execution', async () => { + it('maxPerTurn: 0 skips all tool execution', async () => { const executeSpy = vi.fn().mockReturnValue({ data: 'result' }) const { adapter, calls } = createMockAdapter({ @@ -1320,7 +1360,7 @@ describe('chat()', () => { adapter, messages: [{ role: 'user', content: 'None' }], tools: [serverTool('repeater', executeSpy)], - maxToolCallsPerTurn: 0, + middleware: [toolCallBudget({ maxPerTurn: 0 })], agentLoopStrategy: (state) => state.iterationCount < 2, }) @@ -1340,26 +1380,7 @@ describe('chat()', () => { expect(calls.length).toBe(2) }) - it('rejects negative maxToolCallsPerTurn', async () => { - const { adapter } = createMockAdapter({ - iterations: [[ev.runStarted(), ev.textContent('x'), ev.runFinished()]], - }) - - // TextEngine is constructed when the stream is consumed, not at chat() call. - const stream = chat({ - adapter, - messages: [{ role: 'user', content: 'Hi' }], - maxToolCallsPerTurn: -1, - }) - - await expect( - collectChunks(stream as AsyncIterable), - ).rejects.toThrow( - /maxToolCallsPerTurn must be a non-negative finite number/, - ) - }) - - it('applies maxToolCallsPerTurn to pending tool calls on resume', async () => { + it('applies maxPerTurn to pending tool calls on resume', async () => { const executeSpy = vi.fn().mockReturnValue({ ok: true }) let strategyToolCount = 0 @@ -1407,7 +1428,7 @@ describe('chat()', () => { }, ], tools: [serverTool('repeater', executeSpy)], - maxToolCallsPerTurn: 2, + middleware: [toolCallBudget({ maxPerTurn: 2 })], agentLoopStrategy: (state) => { strategyToolCount = state.toolCallCount return state.iterationCount < 1 diff --git a/testing/e2e/src/routes/api.max-tool-calls-wire.ts b/testing/e2e/src/routes/api.max-tool-calls-wire.ts index d2354dd4d..fbf4c4e69 100644 --- a/testing/e2e/src/routes/api.max-tool-calls-wire.ts +++ b/testing/e2e/src/routes/api.max-tool-calls-wire.ts @@ -2,11 +2,10 @@ import { createFileRoute } from '@tanstack/react-router' import { EventType, chat, - combineStrategies, createChatOptions, maxIterations, - maxToolCalls, toolDefinition, + type ChatMiddleware, } from '@tanstack/ai' import { z } from 'zod' import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' @@ -16,11 +15,48 @@ import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' * * `maxIterations` counts model turns, not tool calls. A single turn can emit * unbounded parallel tool calls. This route emits a fat parallel turn (8 - * calls) and asserts: - * 1. `maxToolCallsPerTurn` caps how many execute - * 2. `maxToolCalls` stops further model turns once the cumulative count hits - * 3. Skipped calls still get error tool results (message history stays consistent) + * calls) and asserts an app-owned tool-budget middleware can: + * 1. Cap how many execute per turn (`onBeforeToolCall` skip) + * 2. Stop further model turns once cumulative emitted tools hit a budget + * (`onShouldContinue`) + * 3. Still emit error tool results for skipped calls (history stays consistent) + * + * The budget middleware is the docs recipe — not a library export. */ +function toolCallBudget(options: { + max?: number + maxPerTurn?: number +}): ChatMiddleware { + const { max, maxPerTurn } = options + let perTurn = 0 + return { + name: 'tool-call-budget', + onIteration() { + perTurn = 0 + }, + onToolPhaseComplete() { + perTurn = 0 + }, + onBeforeToolCall() { + if (maxPerTurn == null) return undefined + perTurn += 1 + if (perTurn > maxPerTurn) { + return { + type: 'skip' as const, + result: { + error: `Skipped: exceeded maxToolCallsPerTurn (${maxPerTurn})`, + }, + } + } + return undefined + }, + onShouldContinue(_ctx, state) { + if (max != null && state.toolCallCount >= max) return false + return undefined + }, + } +} + function createFatParallelAdapter(callCount: number): AnyTextAdapter { return { kind: 'text', @@ -148,16 +184,14 @@ export const Route = createFileRoute('/api/max-tool-calls-wire')({ }), tools: [ping], messages: [{ role: 'user', content: 'Call ping many times' }], - // Cap execution fan-out inside the fat turn - maxToolCallsPerTurn: 3, - // Cumulative tool-call budget: after the first turn (8 emitted), - // toolCallCount >= 5 so the strategy must not request another model - // turn that would keep looping. We also allow maxIterations high - // enough that only maxToolCalls can be the stopping reason. - agentLoopStrategy: combineStrategies([ - maxIterations(10), - maxToolCalls(5), - ]), + // maxIterations high enough that only middleware can stop the loop + agentLoopStrategy: maxIterations(10), + middleware: [ + toolCallBudget({ + maxPerTurn: 3, + max: 5, + }), + ], })) { chunks.push(chunk) } diff --git a/testing/e2e/tests/max-tool-calls.spec.ts b/testing/e2e/tests/max-tool-calls.spec.ts index 947d8b001..1a1206394 100644 --- a/testing/e2e/tests/max-tool-calls.spec.ts +++ b/testing/e2e/tests/max-tool-calls.spec.ts @@ -2,9 +2,13 @@ import { expect, test } from '@playwright/test' /** * Regression for issue #964: maxIterations counts model turns, not tool calls. - * maxToolCalls + maxToolCallsPerTurn must bound fan-out from a single fat turn. + * App-owned tool-budget middleware (onBeforeToolCall + onShouldContinue) must + * bound fan-out from a single fat turn. + * + * Uses a synthetic in-process adapter (no real provider HTTP call), so aimock + * wiring is not applicable here. */ -test.describe('chat() maxToolCalls / maxToolCallsPerTurn (#964)', () => { +test.describe('chat() tool-call budget middleware (#964)', () => { test('caps parallel fan-out and cumulative tool calls', async ({ request, }) => { @@ -19,7 +23,7 @@ test.describe('chat() maxToolCalls / maxToolCallsPerTurn (#964)', () => { expect(body.error).toBeNull() - // 8 parallel calls emitted; maxToolCallsPerTurn=3 → only 3 execute + // 8 parallel calls emitted; maxPerTurn=3 → only 3 execute expect(body.executeCount).toBe(3) const toolResults = body.chunks.filter((c) => c.type === 'TOOL_CALL_RESULT') @@ -35,7 +39,7 @@ test.describe('chat() maxToolCalls / maxToolCallsPerTurn (#964)', () => { }) expect(skipped.length).toBe(5) - // maxToolCalls(5): after the fat turn toolCallCount=8 (>= 5), so no further + // max: 5 — after the fat turn toolCallCount=8 (>= 5), so no further // model turn runs. executeCount stays at 3 either way. const runFinished = body.chunks.filter((c) => c.type === 'RUN_FINISHED') expect(runFinished.length).toBe(1) From c3f61548d161c9935a8543163be93127efccd4e4 Mon Sep 17 00:00:00 2001 From: Tom Beckenham <34339192+tombeckenham@users.noreply.github.com> Date: Tue, 21 Jul 2026 17:47:46 +1000 Subject: [PATCH 2/3] fix(ai): always run onShouldContinue when evaluating the agent loop Evaluate agentLoopStrategy and middleware onShouldContinue together so middleware still observes the final state when the strategy is what stops the loop. Also use a type-only import for ChatMiddleware in the E2E wire route. --- packages/ai/src/activities/chat/index.ts | 26 ++++++++----------- .../e2e/src/routes/api.max-tool-calls-wire.ts | 7 +++-- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/packages/ai/src/activities/chat/index.ts b/packages/ai/src/activities/chat/index.ts index 2d0edb220..de4f0c053 100644 --- a/packages/ai/src/activities/chat/index.ts +++ b/packages/ai/src/activities/chat/index.ts @@ -1998,22 +1998,18 @@ class TextEngine< lastTurnToolCallCount: this.lastTurnToolCallCount, } - // Evaluate strategy and middleware even when toolPhase is 'stop' so - // observers can see final counters; AND with toolPhase at the end. - if (!this.loopStrategy(state)) { - return false - } - - if ( - !(await this.middlewareRunner.runOnShouldContinue( - this.middlewareCtx, - state, - )) - ) { - return false - } + // Evaluate strategy and middleware unconditionally (even when the + // strategy already says stop) so every onShouldContinue observer still + // sees the final counters; AND all three at the end. + const strategyContinues = this.loopStrategy(state) + const middlewareContinues = await this.middlewareRunner.runOnShouldContinue( + this.middlewareCtx, + state, + ) - return this.toolPhase === 'continue' + return ( + strategyContinues && middlewareContinues && this.toolPhase === 'continue' + ) } /** diff --git a/testing/e2e/src/routes/api.max-tool-calls-wire.ts b/testing/e2e/src/routes/api.max-tool-calls-wire.ts index fbf4c4e69..7226bbb4d 100644 --- a/testing/e2e/src/routes/api.max-tool-calls-wire.ts +++ b/testing/e2e/src/routes/api.max-tool-calls-wire.ts @@ -5,10 +5,13 @@ import { createChatOptions, maxIterations, toolDefinition, - type ChatMiddleware, } from '@tanstack/ai' import { z } from 'zod' -import type { AnyTextAdapter, StreamChunk } from '@tanstack/ai' +import type { + AnyTextAdapter, + ChatMiddleware, + StreamChunk, +} from '@tanstack/ai' /** * Wire-format regression for issue #964. From 18f47a4eefd7670e310c63280c1076b06f5d565f Mon Sep 17 00:00:00 2001 From: "autofix-ci[bot]" <114827586+autofix-ci[bot]@users.noreply.github.com> Date: Tue, 21 Jul 2026 07:49:22 +0000 Subject: [PATCH 3/3] ci: apply automated fixes --- testing/e2e/src/routes/api.max-tool-calls-wire.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/testing/e2e/src/routes/api.max-tool-calls-wire.ts b/testing/e2e/src/routes/api.max-tool-calls-wire.ts index 7226bbb4d..a77516d05 100644 --- a/testing/e2e/src/routes/api.max-tool-calls-wire.ts +++ b/testing/e2e/src/routes/api.max-tool-calls-wire.ts @@ -7,11 +7,7 @@ import { toolDefinition, } from '@tanstack/ai' import { z } from 'zod' -import type { - AnyTextAdapter, - ChatMiddleware, - StreamChunk, -} from '@tanstack/ai' +import type { AnyTextAdapter, ChatMiddleware, StreamChunk } from '@tanstack/ai' /** * Wire-format regression for issue #964.