@tanstack/ai 0.23.1 → 0.24.0

package/src/activities/chat/middleware/types.ts

@@ -34,7 +34,7 @@ export type ChatMiddlewarePhase =
  * Stable context object passed to all middleware hooks.
  * Created once per chat() invocation and shared across all hooks.
  */
-export interface ChatMiddlewareContext {
+export interface ChatMiddlewareContext<TContext = unknown> {
   /** Unique identifier for this chat request */
   requestId: string
   /** Unique identifier for this stream */
@@ -64,8 +64,8 @@ export interface ChatMiddlewareContext {
   signal?: AbortSignal
   /** Abort the chat run with a reason */
   abort: (reason?: string) => void
-  /** Opaque user-provided value from chat() options */
-  context: unknown
+  /** Runtime context provided by chat() options */
+  context: TContext
   /**
    * Defer a non-blocking side-effect promise.
    * Deferred promises do not block streaming and are awaited
@@ -343,7 +343,7 @@ export interface ErrorInfo {
  * }
  * ```
  */
-export interface ChatMiddleware {
+export interface ChatMiddleware<TContext = unknown> {
   /** Optional name for debugging and identification */
   name?: string
 
@@ -355,7 +355,7 @@ export interface ChatMiddleware {
    * Only the fields you return are overwritten — everything else is preserved.
    */
   onConfig?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     config: ChatMiddlewareConfig,
   ) =>
     | void
@@ -379,7 +379,7 @@ export interface ChatMiddleware {
    * outputSchema or apply structured-output-specific behavior.
    */
   onStructuredOutputConfig?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     config: StructuredOutputMiddlewareConfig,
   ) =>
     | void
@@ -390,14 +390,14 @@ export interface ChatMiddleware {
   /**
    * Called when the chat run starts (after initial onConfig).
    */
-  onStart?: (ctx: ChatMiddlewareContext) => void | Promise<void>
+  onStart?: (ctx: ChatMiddlewareContext<TContext>) => void | Promise<void>
 
   /**
    * Called at the start of each agent loop iteration, after a new assistant message ID
    * is created. Use this to observe iteration boundaries.
    */
   onIteration?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     info: IterationInfo,
   ) => void | Promise<void>
 
@@ -408,7 +408,7 @@ export interface ChatMiddleware {
    * @returns void (pass through), chunk (replace), chunk[] (expand), null (drop)
    */
   onChunk?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     chunk: StreamChunk,
   ) =>
     | void
@@ -422,7 +422,7 @@ export interface ChatMiddleware {
    * Can observe, transform args, skip execution, or abort the run.
    */
   onBeforeToolCall?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     hookCtx: ToolCallHookContext,
   ) => BeforeToolCallDecision | Promise<BeforeToolCallDecision>
 
@@ -430,7 +430,7 @@ export interface ChatMiddleware {
    * Called after a tool execution completes (success or failure).
    */
   onAfterToolCall?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     info: AfterToolCallInfo,
   ) => void | Promise<void>
 
@@ -439,7 +439,7 @@ export interface ChatMiddleware {
    * Provides aggregate data about tool execution results, approvals, and client tools.
    */
   onToolPhaseComplete?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     info: ToolPhaseCompleteInfo,
   ) => void | Promise<void>
 
@@ -448,7 +448,7 @@ export interface ChatMiddleware {
    * Called once per model iteration that reports usage.
    */
   onUsage?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     usage: UsageInfo,
   ) => void | Promise<void>
 
@@ -457,7 +457,7 @@ export interface ChatMiddleware {
    * Exactly one of onFinish/onAbort/onError will be called per run.
    */
   onFinish?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     info: FinishInfo,
   ) => void | Promise<void>
 
@@ -466,7 +466,7 @@ export interface ChatMiddleware {
    * Exactly one of onFinish/onAbort/onError will be called per run.
    */
   onAbort?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     info: AbortInfo,
   ) => void | Promise<void>
 
@@ -475,7 +475,7 @@ export interface ChatMiddleware {
    * Exactly one of onFinish/onAbort/onError will be called per run.
    */
   onError?: (
-    ctx: ChatMiddlewareContext,
+    ctx: ChatMiddlewareContext<TContext>,
     info: ErrorInfo,
   ) => void | Promise<void>
 }

package/src/activities/chat/runtime-context-types.ts

@@ -0,0 +1,68 @@
+import type { ChatMiddleware } from './middleware/types'
+
+/**
+ * Shared type-level helpers for inferring the runtime `context` requirement
+ * from typed tools and middleware.
+ *
+ * These primitives are consumed by both the chat activity options
+ * (`./index.ts`, which merges tool + middleware requirements) and the tool
+ * execution layer (`./tools/tool-calls.ts`, which only sees tools). They live
+ * here so the two call sites share one definition instead of maintaining
+ * divergent copies.
+ */
+
+/** True only when `T` is exactly `unknown`. */
+type IsUnknown<T> = unknown extends T
+  ? [T] extends [unknown]
+    ? true
+    : false
+  : false
+
+/**
+ * Drops an `unknown` context requirement to `never` so that untyped tools and
+ * middleware (which default `TContext` to `unknown`) contribute no requirement
+ * to the merged context.
+ */
+type KnownContext<T> = IsUnknown<T> extends true ? never : T
+
+/**
+ * Merge two inferred context requirements, treating `never` as "no
+ * requirement". Using this instead of a raw intersection keeps a `never`
+ * (untyped) contributor from collapsing the whole merge to `never`.
+ */
+export type MergeContext<TLeft, TRight> = [TLeft] extends [never]
+  ? TRight
+  : [TRight] extends [never]
+    ? TLeft
+    : TLeft & TRight
+
+/** Collapse a union of context requirements into their intersection. */
+export type UnionToIntersection<T> = [T] extends [never]
+  ? never
+  : (T extends unknown ? (value: T) => void : never) extends (
+        value: infer TIntersection,
+      ) => void
+    ? TIntersection
+    : never
+
+/** Strip `undefined` from a context requirement. */
+export type DefinedContext<T> = Exclude<T, undefined>
+
+/**
+ * Extract the `context` requirement declared by a tool execute function's
+ * second argument, dropping `unknown` (untyped) contexts to `never`.
+ */
+type ContextFromExecute<T> = T extends (...args: any) => any
+  ? NonNullable<Parameters<T>[1]> extends { context: infer TUserContext }
+    ? KnownContext<TUserContext>
+    : never
+  : never
+
+/** Extract the context requirement declared by a single tool. */
+export type ContextFromTool<T> = T extends { execute?: infer TExecute }
+  ? ContextFromExecute<TExecute>
+  : never
+
+/** Extract the context requirement declared by a single middleware. */
+export type ContextFromMiddleware<T> =
+  T extends ChatMiddleware<infer TContext> ? KnownContext<TContext> : never

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

@@ -7,6 +7,7 @@
 
 import { parsePartialJSON } from './json-parser'
 import type {
+  ContentPart,
   StructuredOutputPart,
   ThinkingPart,
   ToolCallPart,
@@ -107,7 +108,7 @@ export function updateToolResultPart(
   messages: Array<UIMessage>,
   messageId: string,
   toolCallId: string,
-  content: string,
+  content: string | Array<ContentPart>,
   state: ToolResultState,
   error?: string,
 ): Array<UIMessage> {

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

@@ -18,6 +18,7 @@
  *   adapter contract, single-shot flows, and expected UIMessage output.
  */
 import { generateMessageId, uiMessageToModelMessages } from '../messages.js'
+import { normalizeToolResult } from '../../../utilities/tool-result'
 import { defaultJSONParser } from './json-parser'
 import {
   appendStructuredOutputDelta,
@@ -321,7 +322,7 @@ export class StreamProcessor {
     )
 
     // Step 2: Create a tool-result part (for LLM conversation history)
-    const content = typeof output === 'string' ? output : JSON.stringify(output)
+    const content = normalizeToolResult(output)
     const toolResultState: ToolResultState = error ? 'error' : 'complete'
 
     updatedMessages = updateToolResultPart(
@@ -1170,30 +1171,51 @@ export class StreamProcessor {
       // Step 1: Update the tool-call part's output field (for UI consistency
       // with client tools — see GitHub issue #176)
       let output: unknown
-      try {
-        output = JSON.parse(chunk.result)
-      } catch {
+      if (Array.isArray(chunk.result)) {
         output = chunk.result
+      } else {
+        try {
+          output = JSON.parse(chunk.result)
+        } catch {
+          output = chunk.result
+        }
       }
       this.messages = updateToolCallWithOutput(
         this.messages,
         chunk.toolCallId,
         output,
+        chunk.state === 'output-error' ? 'input-complete' : undefined,
       )
 
       // Step 2: Create/update the tool-result part (for LLM conversation history)
-      const resultState: ToolResultState = 'complete'
+      const resultState: ToolResultState =
+        chunk.state === 'output-error' ? 'error' : 'complete'
       this.messages = updateToolResultPart(
         this.messages,
         messageId,
         chunk.toolCallId,
         chunk.result,
         resultState,
+        resultState === 'error'
+          ? this.extractToolResultError(output)
+          : undefined,
       )
       this.emitMessagesChange()
     }
   }
 
+  private extractToolResultError(output: unknown): string {
+    if (
+      output &&
+      typeof output === 'object' &&
+      'error' in output &&
+      typeof output.error === 'string'
+    ) {
+      return output.error
+    }
+    return typeof output === 'string' ? output : 'Tool execution failed'
+  }
+
   /**
    * Handle TOOL_CALL_RESULT event (AG-UI spec).
    *
@@ -1218,16 +1240,19 @@ export class StreamProcessor {
       this.messages,
       chunk.toolCallId,
       output,
+      chunk.state === 'output-error' ? 'input-complete' : undefined,
     )
 
     // Step 2: Create/update the tool-result part
-    const resultState: ToolResultState = 'complete'
+    const resultState: ToolResultState =
+      chunk.state === 'output-error' ? 'error' : 'complete'
     this.messages = updateToolResultPart(
       this.messages,
       messageId,
       chunk.toolCallId,
       chunk.content,
       resultState,
+      resultState === 'error' ? this.extractToolResultError(output) : undefined,
     )
     this.emitMessagesChange()
   }
@@ -1274,7 +1299,10 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
   ): void {
     this.hasError = true
-    const runId = (chunk as any).runId as string | undefined
+    const runId =
+      'runId' in chunk && typeof chunk.runId === 'string'
+        ? chunk.runId
+        : undefined
     if (runId) {
       this.activeRuns.delete(runId)
     } else {
@@ -1305,7 +1333,19 @@ export class StreamProcessor {
       this.emitMessagesChange()
     }
 
-    this.events.onError?.(new Error(errorMessage))
+    // Attach the provider's structured error body (`rawEvent`) and `code` to
+    // the surfaced Error so consumers can recover the upstream detail that the
+    // RUN_ERROR's `message` alone discards. Both are optional and added only
+    // when present, keeping the Error backward compatible.
+    const error = new Error(errorMessage)
+    const code = chunk.code ?? chunk.error?.code
+    if (code !== undefined) {
+      Object.assign(error, { code })
+    }
+    if (chunk.rawEvent !== undefined) {
+      Object.assign(error, { rawEvent: chunk.rawEvent })
+    }
+    this.events.onError?.(error)
   }
 
   /**

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

@@ -1,5 +1,8 @@
+import { normalizeToolResult } from '../../../utilities/tool-result'
 import { isStandardSchema, parseWithStandardSchema } from './schema-converter'
 import type {
+  AnyTool,
+  ContentPart,
   CustomEvent,
   ModelMessage,
   RunFinishedEvent,
@@ -9,11 +12,18 @@ import type {
   ToolCallEndEvent,
   ToolCallStartEvent,
   ToolExecutionContext,
+  ToolOutputState,
 } from '../../../types'
 import type {
   AfterToolCallInfo,
   BeforeToolCallDecision,
 } from '../middleware/types'
+import type {
+  ContextFromTool,
+  DefinedContext,
+  MergeContext,
+  UnionToIntersection,
+} from '../runtime-context-types'
 
 function safeJsonParse(value: string): unknown {
   try {
@@ -46,6 +56,36 @@ export class MiddlewareAbortError extends Error {
   }
 }
 
+// The leaf context-inference primitives (ContextFromTool, MergeContext,
+// UnionToIntersection, DefinedContext) are shared with the chat activity
+// options layer — see ../runtime-context-types.
+type RequiredContextFromToolUnion<T> = T extends unknown
+  ? undefined extends ContextFromTool<T>
+    ? never
+    : ContextFromTool<T>
+  : never
+
+type ContextFromToolUnion<T> = [
+  UnionToIntersection<DefinedContext<ContextFromTool<T>>>,
+] extends [never]
+  ? unknown
+  : [RequiredContextFromToolUnion<T>] extends [never]
+    ? UnionToIntersection<DefinedContext<ContextFromTool<T>>> | undefined
+    : UnionToIntersection<DefinedContext<ContextFromTool<T>>>
+
+type ContextFromTools<TTools> = TTools extends readonly [
+  infer THead,
+  ...infer TTail,
+]
+  ? MergeContext<ContextFromTool<THead>, ContextFromTools<TTail>>
+  : TTools extends ReadonlyArray<infer TTool>
+    ? ContextFromToolUnion<TTool>
+    : unknown
+
+type ExecuteToolsContextArgs<TContext> = undefined extends TContext
+  ? [userContext?: TContext]
+  : [userContext: TContext]
+
 /**
  * Manages tool call accumulation and execution for the chat() method's automatic tool execution loop.
  *
@@ -80,11 +120,22 @@ export class MiddlewareAbortError extends Error {
  * }
  * ```
  */
-export class ToolCallManager {
+export class ToolCallManager<
+  TToolsOrContext = ReadonlyArray<AnyTool>,
+  TContext = TToolsOrContext extends ReadonlyArray<AnyTool>
+    ? ContextFromTools<TToolsOrContext>
+    : TToolsOrContext,
+> {
   private readonly toolCallsMap = new Map<number, ToolCall>()
-  private readonly tools: ReadonlyArray<Tool>
-
-  constructor(tools: ReadonlyArray<Tool>) {
+  private readonly tools: TToolsOrContext extends ReadonlyArray<AnyTool>
+    ? TToolsOrContext
+    : ReadonlyArray<AnyTool>
+
+  constructor(
+    tools: TToolsOrContext extends ReadonlyArray<AnyTool>
+      ? TToolsOrContext
+      : ReadonlyArray<AnyTool>,
+  ) {
     this.tools = tools
   }
 
@@ -161,14 +212,18 @@ export class ToolCallManager {
    */
   async *executeTools(
     finishEvent: RunFinishedEvent,
+    ...contextArgs: ExecuteToolsContextArgs<TContext>
   ): AsyncGenerator<ToolCallEndEvent, Array<ModelMessage>, void> {
     const toolCallsArray = this.getToolCalls()
     const toolResults: Array<ModelMessage> = []
+    const hasRuntimeContext = contextArgs.length > 0
+    const userContext = contextArgs[0]
 
     for (const toolCall of toolCallsArray) {
       const tool = this.tools.find((t) => t.name === toolCall.function.name)
 
-      let toolResultContent: string
+      let toolResultContent: string | Array<ContentPart>
+      let toolResultState: ToolOutputState | undefined
       if (tool?.execute) {
         try {
           // Parse arguments (normalize null/non-object to {} for empty tool_use blocks)
@@ -199,15 +254,21 @@ export class ToolCallManager {
           }
 
           // Execute the tool
-          let result = await tool.execute(args)
-
-          // Validate output against outputSchema if provided (for Standard Schema compliant schemas)
-          if (
-            tool.outputSchema &&
-            isStandardSchema(tool.outputSchema) &&
-            result !== undefined &&
-            result !== null
-          ) {
+          const executionContext = {
+            toolCallId: toolCall.id,
+            context: userContext,
+            emitCustomEvent: () => {},
+          } as ToolExecutionContext<TContext>
+          let result = hasRuntimeContext
+            ? await tool.execute(args, executionContext)
+            : await tool.execute(args)
+
+          // Validate output against outputSchema if provided (for Standard
+          // Schema compliant schemas). Unlike the previous implementation we
+          // intentionally validate `undefined`/`null` results too, so a tool
+          // whose schema forbids them surfaces a validation error instead of
+          // silently passing — the schema itself decides whether they're valid.
+          if (tool.outputSchema && isStandardSchema(tool.outputSchema)) {
             try {
               result = parseWithStandardSchema(tool.outputSchema, result)
             } catch (validationError: unknown) {
@@ -221,13 +282,13 @@ export class ToolCallManager {
             }
           }
 
-          toolResultContent =
-            typeof result === 'string' ? result : JSON.stringify(result)
+          toolResultContent = normalizeToolResult(result)
         } catch (error: unknown) {
           // If tool execution fails, add error message
           const message =
             error instanceof Error ? error.message : 'Unknown error'
           toolResultContent = `Error executing tool: ${message}`
+          toolResultState = 'output-error'
         }
       } else {
         // Tool doesn't have execute function, add placeholder
@@ -243,6 +304,7 @@ export class ToolCallManager {
         model: finishEvent.model,
         timestamp: Date.now(),
         result: toolResultContent,
+        ...(toolResultState !== undefined && { state: toolResultState }),
       } as ToolCallEndEvent
 
       // Add tool result message
@@ -370,7 +432,7 @@ async function applyBeforeToolCallDecision(
       result:
         typeof skipResult === 'string'
           ? safeJsonParse(skipResult)
-          : skipResult || null,
+          : (skipResult ?? null),
       duration: 0,
     })
     if (middlewareHooks.onAfterToolCall) {
@@ -394,12 +456,12 @@ async function applyBeforeToolCallDecision(
  * Execute a server-side tool with event polling, output validation, and middleware hooks.
  * Yields CustomEvent chunks during execution and pushes the result to the results array.
  */
-async function* executeServerTool(
+async function* executeServerTool<TContext = unknown>(
   toolCall: ToolCall,
-  tool: Tool,
+  tool: AnyTool,
   toolName: string,
   input: unknown,
-  context: ToolExecutionContext,
+  context: ToolExecutionContext<TContext>,
   pendingEvents: Array<CustomEvent>,
   results: Array<ToolResult>,
   middlewareHooks?: ToolExecutionMiddlewareHooks,
@@ -419,18 +481,14 @@ async function* executeServerTool(
       yield pendingEvent
     }
 
-    // Validate output against outputSchema if provided
-    if (
-      tool.outputSchema &&
-      isStandardSchema(tool.outputSchema) &&
-      result !== undefined &&
-      result !== null
-    ) {
+    // Validate output against outputSchema if provided. Validates
+    // `undefined`/`null` too — the schema decides whether they're valid.
+    if (tool.outputSchema && isStandardSchema(tool.outputSchema)) {
       result = parseWithStandardSchema(tool.outputSchema, result)
     }
 
     const finalResult =
-      typeof result === 'string' ? safeJsonParse(result) : result || null
+      typeof result === 'string' ? safeJsonParse(result) : (result ?? null)
 
     results.push({
       toolCallId: toolCall.id,
@@ -486,6 +544,35 @@ async function* executeServerTool(
   }
 }
 
+function buildClientToolResult(
+  toolCallId: string,
+  toolName: string,
+  tool: AnyTool,
+  rawResult: unknown,
+): ToolResult {
+  try {
+    let result = rawResult
+    if (tool.outputSchema && isStandardSchema(tool.outputSchema)) {
+      result = parseWithStandardSchema(tool.outputSchema, result)
+    }
+
+    return {
+      toolCallId,
+      toolName,
+      result:
+        typeof result === 'string' ? safeJsonParse(result) : (result ?? null),
+    }
+  } catch (error: unknown) {
+    const message = error instanceof Error ? error.message : 'Validation failed'
+    return {
+      toolCallId,
+      toolName,
+      result: { error: message },
+      state: 'output-error',
+    }
+  }
+}
+
 /**
  * Execute tool calls based on their configuration.
  * Yields CustomEvent chunks during tool execution for real-time progress updates.
@@ -501,9 +588,9 @@ async function* executeServerTool(
  * @param clientResults - Map of client-side execution results (toolCallId -> result)
  * @param createCustomEventChunk - Factory to create CustomEvent chunks (optional)
  */
-export async function* executeToolCalls(
+export async function* executeToolCalls<TContext = unknown>(
   toolCalls: Array<ToolCall>,
-  tools: ReadonlyArray<Tool>,
+  tools: ReadonlyArray<AnyTool>,
   approvals: Map<string, boolean> = new Map(),
   clientResults: Map<string, any> = new Map(),
   createCustomEventChunk?: (
@@ -511,13 +598,14 @@ export async function* executeToolCalls(
     value: Record<string, any>,
   ) => CustomEvent,
   middlewareHooks?: ToolExecutionMiddlewareHooks,
+  userContext?: TContext,
 ): AsyncGenerator<CustomEvent, ExecuteToolCallsResult, void> {
   const results: Array<ToolResult> = []
   const needsApproval: Array<ApprovalRequest> = []
   const needsClientExecution: Array<ClientToolRequest> = []
 
   // Create tool lookup map
-  const toolMap = new Map<string, Tool>()
+  const toolMap = new Map<string, AnyTool>()
   for (const tool of tools) {
     toolMap.set(tool.name, tool)
   }
@@ -588,8 +676,9 @@ export async function* executeToolCalls(
 
     // Create a ToolExecutionContext for this tool call with event emission
     const pendingEvents: Array<CustomEvent> = []
-    const context: ToolExecutionContext = {
+    const context = {
       toolCallId: toolCall.id,
+      context: userContext,
       emitCustomEvent: (eventName: string, value: Record<string, any>) => {
         if (createCustomEventChunk) {
           pendingEvents.push(
@@ -600,7 +689,7 @@ export async function* executeToolCalls(
           )
         }
       },
-    }
+    } as ToolExecutionContext<TContext>
 
     // CASE 1: Client-side tool (no execute function)
     if (!tool.execute) {
@@ -615,11 +704,14 @@ export async function* executeToolCalls(
           if (approved) {
             // Approved - check if client has executed
             if (clientResults.has(toolCall.id)) {
-              results.push({
-                toolCallId: toolCall.id,
-                toolName,
-                result: clientResults.get(toolCall.id),
-              })
+              results.push(
+                buildClientToolResult(
+                  toolCall.id,
+                  toolName,
+                  tool,
+                  clientResults.get(toolCall.id),
+                ),
+              )
             } else {
               // Approved but not executed yet - request client execution
               needsClientExecution.push({
@@ -649,11 +741,14 @@ export async function* executeToolCalls(
       } else {
         // No approval needed - check if client has executed
         if (clientResults.has(toolCall.id)) {
-          results.push({
-            toolCallId: toolCall.id,
-            toolName,
-            result: clientResults.get(toolCall.id),
-          })
+          results.push(
+            buildClientToolResult(
+              toolCall.id,
+              toolName,
+              tool,
+              clientResults.get(toolCall.id),
+            ),
+          )
         } else {
           // Request client execution
           needsClientExecution.push({