@tanstack/ai 0.40.0 → 0.42.0

package/src/activities/chat/index.ts

@@ -240,6 +240,11 @@ 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
@@ -473,6 +478,23 @@ 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.
@@ -519,6 +541,12 @@ class TextEngine<
 
   private messages: Array<ModelMessage>
   private iterationCount = 0
+  /** Cumulative tool calls counted in this run (emitted + pending resume). */
+  private toolCallCount = 0
+  /** Tool calls in the most recent budgeted batch (0 when none). */
+  private lastTurnToolCallCount = 0
+  /** Tool call IDs already counted toward `toolCallCount` (avoids double-count on resume). */
+  private readonly countedToolCallIds = new Set<string>()
   private lastFinishReason: string | null = null
   private streamStartTime = 0
   private totalChunkCount = 0
@@ -534,6 +562,7 @@ 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<string, boolean>
   private readonly initialClientToolResults: Map<string, any>
@@ -599,6 +628,9 @@ 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
@@ -1302,9 +1334,14 @@ 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)
+
     // Handle undiscovered lazy tool calls with self-correcting error messages
     const undiscoveredLazyResults: Array<ToolResult> = []
-    const executablePendingCalls = pendingToolCalls.filter((tc) => {
+    const executablePendingCalls = budgetedToolCalls.filter((tc) => {
       if (this.lazyToolManager.isUndiscoveredLazyTool(tc.function.name)) {
         undiscoveredLazyResults.push({
           toolCallId: tc.id,
@@ -1321,16 +1358,27 @@ class TextEngine<
       return true
     })
 
-    if (undiscoveredLazyResults.length > 0) {
-      for (const chunk of this.buildToolResultChunks(
-        undiscoveredLazyResults,
-        finishEvent,
-      )) {
-        yield* this.pipeThroughMiddleware(chunk)
-      }
+    // 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]
+
+    // Build args lookup so buildToolResultChunks can emit TOOL_CALL_START +
+    // TOOL_CALL_ARGS before TOOL_CALL_END during continuation re-executions.
+    const argsMap = new Map<string, string>()
+    for (const tc of pendingToolCalls) {
+      argsMap.set(tc.id, tc.function.arguments)
     }
 
     if (executablePendingCalls.length === 0) {
+      if (deferredErrorResults.length > 0) {
+        for (const chunk of this.buildToolResultChunks(
+          deferredErrorResults,
+          finishEvent,
+          argsMap,
+        )) {
+          yield* this.pipeThroughMiddleware(chunk)
+        }
+      }
       return 'continue'
     }
 
@@ -1384,28 +1432,23 @@ class TextEngine<
       return 'stop'
     }
 
+    const allResults = [...executionResult.results, ...deferredErrorResults]
+
     // Notify middleware of tool phase completion (devtools emits aggregate events here)
     await this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, {
       toolCalls: pendingToolCalls,
-      results: executionResult.results,
+      results: allResults,
       needsApproval: executionResult.needsApproval,
       needsClientExecution: executionResult.needsClientExecution,
     })
 
-    // Build args lookup so buildToolResultChunks can emit TOOL_CALL_START +
-    // TOOL_CALL_ARGS before TOOL_CALL_END during continuation re-executions.
-    const argsMap = new Map<string, string>()
-    for (const tc of pendingToolCalls) {
-      argsMap.set(tc.id, tc.function.arguments)
-    }
-
     if (
       executionResult.needsApproval.length > 0 ||
       executionResult.needsClientExecution.length > 0
     ) {
-      if (executionResult.results.length > 0) {
+      if (allResults.length > 0) {
         for (const chunk of this.buildToolResultChunks(
-          executionResult.results,
+          allResults,
           finishEvent,
           argsMap,
         )) {
@@ -1432,7 +1475,7 @@ class TextEngine<
     }
 
     const toolResultChunks = this.buildToolResultChunks(
-      executionResult.results,
+      allResults,
       finishEvent,
       argsMap,
     )
@@ -1446,6 +1489,8 @@ class TextEngine<
 
   private async *processToolCalls(): AsyncGenerator<StreamChunk, void, void> {
     if (!this.shouldExecuteToolPhase()) {
+      // Text-only turn — clear per-turn count so strategies see 0 tools.
+      this.lastTurnToolCallCount = 0
       this.setToolPhase('stop')
       return
     }
@@ -1454,15 +1499,20 @@ class TextEngine<
     const finishEvent = this.finishedEvent
 
     if (!finishEvent || toolCalls.length === 0) {
+      this.lastTurnToolCallCount = 0
       this.setToolPhase('stop')
       return
     }
 
+    // Count every model-emitted tool call (including ones we may skip).
+    const { toExecute: budgetedToolCalls, skippedResults } =
+      this.applyToolCallBudget(toolCalls)
+
     this.addAssistantToolCallMessage(toolCalls)
 
     // Handle undiscovered lazy tool calls with self-correcting error messages
     const undiscoveredLazyResults: Array<ToolResult> = []
-    const executableToolCalls = toolCalls.filter((tc) => {
+    const executableToolCalls = budgetedToolCalls.filter((tc) => {
       if (this.lazyToolManager.isUndiscoveredLazyTool(tc.function.name)) {
         undiscoveredLazyResults.push({
           toolCallId: tc.id,
@@ -1479,17 +1529,21 @@ class TextEngine<
       return true
     })
 
-    if (undiscoveredLazyResults.length > 0 && this.finishedEvent) {
-      for (const chunk of this.buildToolResultChunks(
-        undiscoveredLazyResults,
-        this.finishedEvent,
-      )) {
-        yield* this.pipeThroughMiddleware(chunk)
-      }
-    }
+    // 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]
 
     if (executableToolCalls.length === 0) {
-      // All tool calls were undiscovered lazy tools — errors emitted, continue loop
+      // All tool calls were undiscovered lazy tools and/or skipped by the
+      // per-turn fan-out cap — errors emitted, continue loop (strategy may stop).
+      if (deferredErrorResults.length > 0) {
+        for (const chunk of this.buildToolResultChunks(
+          deferredErrorResults,
+          finishEvent,
+        )) {
+          yield* this.pipeThroughMiddleware(chunk)
+        }
+      }
       this.toolCallManager.clear()
       this.setToolPhase('continue')
       return
@@ -1548,10 +1602,13 @@ class TextEngine<
       return
     }
 
+    // Executed results first, then deferred errors (fan-out skips / undiscovered)
+    const allResults = [...executionResult.results, ...deferredErrorResults]
+
     // Notify middleware of tool phase completion (devtools emits aggregate events here)
     await this.middlewareRunner.runOnToolPhaseComplete(this.middlewareCtx, {
       toolCalls,
-      results: executionResult.results,
+      results: allResults,
       needsApproval: executionResult.needsApproval,
       needsClientExecution: executionResult.needsClientExecution,
     })
@@ -1560,9 +1617,9 @@ class TextEngine<
       executionResult.needsApproval.length > 0 ||
       executionResult.needsClientExecution.length > 0
     ) {
-      if (executionResult.results.length > 0) {
+      if (allResults.length > 0) {
         for (const chunk of this.buildToolResultChunks(
-          executionResult.results,
+          allResults,
           finishEvent,
         )) {
           yield* this.pipeThroughMiddleware(chunk)
@@ -1587,10 +1644,7 @@ class TextEngine<
       return
     }
 
-    const toolResultChunks = this.buildToolResultChunks(
-      executionResult.results,
-      finishEvent,
-    )
+    const toolResultChunks = this.buildToolResultChunks(allResults, finishEvent)
 
     for (const chunk of toolResultChunks) {
       yield* this.pipeThroughMiddleware(chunk)
@@ -1943,10 +1997,64 @@ class TextEngine<
         iterationCount: this.iterationCount,
         messages: this.messages,
         finishReason: this.lastFinishReason,
+        toolCallCount: this.toolCallCount,
+        lastTurnToolCallCount: this.lastTurnToolCallCount,
       }) && 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.
+   *
+   * 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`.
+   */
+  private applyToolCallBudget(toolCalls: Array<ToolCall>): {
+    toExecute: Array<ToolCall>
+    skippedResults: Array<ToolResult>
+  } {
+    this.lastTurnToolCallCount = toolCalls.length
+    let newlyCounted = 0
+    for (const tc of toolCalls) {
+      if (!this.countedToolCallIds.has(tc.id)) {
+        this.countedToolCallIds.add(tc.id)
+        newlyCounted++
+      }
+    }
+    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<ToolResult> = 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 {
     return !!this.effectiveSignal?.aborted
   }

package/src/activities/chat/messages.ts

@@ -444,12 +444,21 @@ export function modelMessageToUIMessage(
   // Handle tool calls
   if (modelMessage.toolCalls && modelMessage.toolCalls.length > 0) {
     for (const toolCall of modelMessage.toolCalls) {
+      // Model-message arguments are complete, so surface the parsed input.
+      // A malformed arguments string just leaves `input` undefined.
+      let input: unknown
+      try {
+        input = JSON.parse(toolCall.function.arguments)
+      } catch {
+        input = undefined
+      }
       parts.push({
         type: 'tool-call',
         id: toolCall.id,
         name: toolCall.function.name,
         arguments: toolCall.function.arguments,
         state: 'input-complete', // Model messages have complete arguments
+        ...(input !== undefined && { input }),
         ...(toolCall.metadata !== undefined && { metadata: toolCall.metadata }),
       })
     }

package/src/activities/chat/stream/message-updaters.ts

@@ -58,6 +58,8 @@ export function updateToolCallPart(
     name: string
     arguments: string
     state: ToolCallState
+    /** Parsed input — set when the arguments are complete. */
+    input?: unknown
     metadata?: Record<string, unknown>
   },
 ): Array<UIMessage> {
@@ -76,6 +78,9 @@ export function updateToolCallPart(
     // Gemini's thoughtSignature on TOOL_CALL_START) we must not lose it on
     // subsequent updates that don't re-supply it.
     const metadata = toolCall.metadata ?? existing?.metadata
+    // Same for the parsed input: it's supplied once at completion, so
+    // subsequent arg-less updates (approval, etc.) must not drop it.
+    const input = toolCall.input ?? existing?.input
 
     const toolCallPart: ToolCallPart = {
       type: 'tool-call',
@@ -83,9 +88,10 @@ export function updateToolCallPart(
       name: toolCall.name,
       arguments: toolCall.arguments,
       state: toolCall.state,
-      // Carry forward approval and output from the existing part
+      // Carry forward approval, output and parsed input from the existing part
       ...(existing?.approval && { approval: { ...existing.approval } }),
       ...(existing?.output !== undefined && { output: existing.output }),
+      ...(input !== undefined && { input }),
       ...(metadata !== undefined && { metadata }),
     }
 

package/src/activities/chat/stream/processor.ts

@@ -1298,15 +1298,35 @@ export class StreamProcessor {
       // received, back-fill the arguments string so the UIMessage ToolCallPart
       // carries the correct value (defensive against adapters that skip ARGS).
       if (chunk.input !== undefined && !existingToolCall.arguments) {
-        existingToolCall.arguments = JSON.stringify(chunk.input)
+        try {
+          existingToolCall.arguments = JSON.stringify(chunk.input)
+        } catch {
+          // circular refs, BigInt, etc. — leave arguments empty rather than
+          // aborting stream processing
+        }
       }
 
       const index = msgState.toolCallOrder.indexOf(chunk.toolCallId)
       this.completeToolCall(messageId, index, existingToolCall)
       // If TOOL_CALL_END provides parsed input, use it as the canonical parsed
       // arguments (overrides the accumulated string parse from completeToolCall)
+      // and refresh the rendered part's `input` so it reflects the canonical
+      // value rather than the possibly-divergent accumulated-args parse that
+      // completeToolCall wrote (e.g. an adapter that coerces values differently
+      // between the streamed args and the final structured input).
       if (chunk.input !== undefined) {
         existingToolCall.parsedArguments = chunk.input
+        this.messages = updateToolCallPart(this.messages, messageId, {
+          id: existingToolCall.id,
+          name: existingToolCall.name,
+          arguments: existingToolCall.arguments,
+          state: 'input-complete',
+          input: chunk.input,
+          ...(existingToolCall.metadata !== undefined && {
+            metadata: existingToolCall.metadata,
+          }),
+        })
+        this.emitMessagesChange()
       }
     }
 
@@ -1890,12 +1910,20 @@ export class StreamProcessor {
       return
     }
 
-    // Update UIMessage
+    // Update UIMessage. The arguments are complete now, so surface the parsed
+    // input on the part. For adapters that skip TOOL_CALL_ARGS the arguments
+    // string was back-filled from TOOL_CALL_END.input, so this parse matches
+    // the canonical input. If a TOOL_CALL_END.input diverges from the
+    // accumulated args, handleToolCallEndEvent re-updates the part with the
+    // canonical value after this call.
     this.messages = updateToolCallPart(this.messages, messageId, {
       id: toolCall.id,
       name: toolCall.name,
       arguments: toolCall.arguments,
       state: 'input-complete',
+      ...(toolCall.parsedArguments !== undefined && {
+        input: toolCall.parsedArguments,
+      }),
       ...(toolCall.metadata !== undefined && { metadata: toolCall.metadata }),
     })
     this.emitMessagesChange()

package/src/activities/chat/tools/tool-definition.ts

@@ -26,6 +26,10 @@ export interface ClientTool<
   TOutput extends SchemaInput = SchemaInput,
   TName extends string = string,
   TContext = unknown,
+  // Captured as a literal (`true` / `false`) so downstream types — notably
+  // the tool-call part's `approval` field — can be gated on it. Defaults to
+  // `false` when the tool config omits `needsApproval`.
+  TNeedsApproval extends boolean = false,
 > {
   __toolSide: 'client'
   name: TName
@@ -37,7 +41,7 @@ export interface ClientTool<
   // because `undefined` doesn't extend the schema constraint.
   inputSchema?: TInput
   outputSchema?: TOutput
-  needsApproval?: boolean
+  needsApproval?: TNeedsApproval
   lazy?: boolean
   metadata?: Record<string, unknown>
   execute?: ToolExecuteFunction<TInput, TOutput, TContext>
@@ -51,18 +55,22 @@ export interface ToolDefinitionInstance<
   TOutput extends SchemaInput = SchemaInput,
   TName extends string = string,
   TContext = unknown,
+  TNeedsApproval extends boolean = false,
 > extends Tool<TInput, TOutput, TName, TContext> {
   __toolSide: 'definition'
+  // Narrow the base `needsApproval?: boolean` to the captured literal so it
+  // survives into `ToolCallPartForTool`'s approval gate.
+  needsApproval?: TNeedsApproval
 }
 
 /**
  * Union type for any kind of client-side tool (client tool or definition)
  */
 export type AnyClientTool =
-  | (Omit<ClientTool<any, any, string, any>, 'execute'> & {
+  | (Omit<ClientTool<any, any, string, any, boolean>, 'execute'> & {
       execute?: ((args: any, context?: any) => any) | undefined
     })
-  | (Omit<ToolDefinitionInstance<any, any, string, any>, 'execute'> & {
+  | (Omit<ToolDefinitionInstance<any, any, string, any, boolean>, 'execute'> & {
       execute?: ((args: any, context?: any) => any) | undefined
     })
 
@@ -100,12 +108,13 @@ export interface ToolDefinitionConfig<
   TInput extends SchemaInput = SchemaInput,
   TOutput extends SchemaInput = SchemaInput,
   TName extends string = string,
+  TNeedsApproval extends boolean = false,
 > {
   name: TName
   description: string
   inputSchema?: TInput
   outputSchema?: TOutput
-  needsApproval?: boolean
+  needsApproval?: TNeedsApproval
   lazy?: boolean
   metadata?: Record<string, unknown>
 }
@@ -117,7 +126,14 @@ export interface ToolDefinition<
   TInput extends SchemaInput = SchemaInput,
   TOutput extends SchemaInput = SchemaInput,
   TName extends string = string,
-> extends ToolDefinitionInstance<TInput, TOutput, TName> {
+  TNeedsApproval extends boolean = false,
+> extends ToolDefinitionInstance<
+  TInput,
+  TOutput,
+  TName,
+  unknown,
+  TNeedsApproval
+> {
   /**
    * Create a server-side tool with execute function
    */
@@ -126,11 +142,13 @@ export interface ToolDefinition<
   ) => ServerTool<TInput, TOutput, TName, TContext>
 
   /**
-   * Create a client-side tool with optional execute function
+   * Create a client-side tool with optional execute function.
+   * Carries the definition's `needsApproval` literal through to the client
+   * tool so the tool-call part's `approval` field stays gated on it.
    */
   client: <TContext = unknown>(
     execute?: ToolExecuteFunction<TInput, TOutput, TContext>,
-  ) => ClientTool<TInput, TOutput, TName, TContext>
+  ) => ClientTool<TInput, TOutput, TName, TContext, TNeedsApproval>
 }
 
 /**
@@ -192,10 +210,14 @@ export function toolDefinition<
   TInput extends SchemaInput = SchemaInput,
   TOutput extends SchemaInput = SchemaInput,
   TName extends string = string,
+  // `const` forces the literal (`true` / `false`) to be captured from the
+  // config's optional `needsApproval` — without it TS widens to `boolean`,
+  // which collapses the approval gate in `ToolCallPartForTool`.
+  const TNeedsApproval extends boolean = false,
 >(
-  config: ToolDefinitionConfig<TInput, TOutput, TName>,
-): ToolDefinition<TInput, TOutput, TName> {
-  const definition: ToolDefinition<TInput, TOutput, TName> = {
+  config: ToolDefinitionConfig<TInput, TOutput, TName, TNeedsApproval>,
+): ToolDefinition<TInput, TOutput, TName, TNeedsApproval> {
+  const definition: ToolDefinition<TInput, TOutput, TName, TNeedsApproval> = {
     __toolSide: 'definition',
     ...config,
     server<TContext = unknown>(
@@ -210,7 +232,7 @@ export function toolDefinition<
 
     client<TContext = unknown>(
       execute?: ToolExecuteFunction<TInput, TOutput, TContext>,
-    ): ClientTool<TInput, TOutput, TName, TContext> {
+    ): ClientTool<TInput, TOutput, TName, TContext, TNeedsApproval> {
       return {
         __toolSide: 'client',
         ...config,

package/src/activities/generateAudio/index.ts

@@ -14,10 +14,10 @@ import {
   runGenerationFinish,
   runGenerationStart,
   runGenerationUsage,
-} from '../middleware'
+} from '../middleware/run'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
-import type { GenerationMiddleware } from '../middleware'
+import type { GenerationMiddleware } from '../middleware/types'
 import type { AudioAdapter } from './adapter'
 import type { AudioGenerationResult, StreamChunk } from '../../types'
 

package/src/activities/generateImage/index.ts

@@ -14,11 +14,11 @@ import {
   runGenerationFinish,
   runGenerationStart,
   runGenerationUsage,
-} from '../middleware'
+} from '../middleware/run'
 import { resolveMediaPrompt } from '../../utilities/media-prompt'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
-import type { GenerationMiddleware } from '../middleware'
+import type { GenerationMiddleware } from '../middleware/types'
 import type { ImageAdapter } from './adapter'
 import type {
   ImageGenerationResult,

package/src/activities/generateSpeech/index.ts

@@ -14,10 +14,10 @@ import {
   runGenerationFinish,
   runGenerationStart,
   runGenerationUsage,
-} from '../middleware'
+} from '../middleware/run'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
-import type { GenerationMiddleware } from '../middleware'
+import type { GenerationMiddleware } from '../middleware/types'
 import type { TTSAdapter } from './adapter'
 import type { StreamChunk, TTSResult } from '../../types'
 

package/src/activities/generateTranscription/index.ts

@@ -14,10 +14,10 @@ import {
   runGenerationFinish,
   runGenerationStart,
   runGenerationUsage,
-} from '../middleware'
+} from '../middleware/run'
 import type { InternalLogger } from '../../logger/internal-logger'
 import type { DebugOption } from '../../logger/types'
-import type { GenerationMiddleware } from '../middleware'
+import type { GenerationMiddleware } from '../middleware/types'
 import type { TranscriptionAdapter } from './adapter'
 import type {
   StreamChunk,