@tanstack/ai 0.23.1 → 0.25.0

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

@@ -2,9 +2,9 @@ import type {
   JSONSchema,
   ModelMessage,
   StreamChunk,
+  TokenUsage,
   Tool,
   ToolCall,
-  UsageTotals,
 } from '../../../types'
 import type { SystemPrompt } from '../../../system-prompts'
 
@@ -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
@@ -266,11 +266,11 @@ export interface ToolPhaseCompleteInfo {
  * Token usage statistics passed to the onUsage hook.
  * Extracted from the RUN_FINISHED chunk when usage data is present.
  *
- * Includes optional provider-reported `cost`/`costDetails` (see {@link UsageTotals}).
- * Kept as an interface extending `UsageTotals` to preserve declaration merging for
+ * Includes optional provider-reported `cost`/`costDetails` (see {@link TokenUsage}).
+ * Kept as an interface extending `TokenUsage` to preserve declaration merging for
  * this publicly exported type.
  */
-export interface UsageInfo extends UsageTotals {}
+export interface UsageInfo extends TokenUsage {}
 
 // ===========================
 // Terminal Hook Info
@@ -287,7 +287,7 @@ export interface FinishInfo {
   /** Final accumulated text content */
   content: string
   /** Final usage totals, if available (optionally including provider-reported cost) */
-  usage?: UsageTotals | undefined
+  usage?: TokenUsage | undefined
 }
 
 /**
@@ -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)
   }
 
   /**