@tanstack/ai 0.23.1 → 0.25.0

package/src/types.ts

@@ -4,6 +4,16 @@ import type {
 } from '@standard-schema/spec'
 import type { InternalLogger } from './logger/internal-logger'
 import type { SystemPrompt } from './system-prompts'
+// The canonical usage types live in the leaf `@tanstack/ai-event-client`
+// package (which `@tanstack/ai` already depends on) so there is a single source
+// of truth without a dependency cycle. They are re-exported below.
+import type {
+  CompletionTokensDetails,
+  PromptTokensDetails,
+  ProviderUsageDetails,
+  TokenUsage,
+  UsageCostBreakdown,
+} from '@tanstack/ai-event-client'
 import type {
   BaseEvent as AGUIBaseEvent,
   CustomEvent as AGUICustomEvent,
@@ -50,6 +60,8 @@ export type ToolResultState =
   | 'complete' // Result is complete
   | 'error' // Error occurred
 
+export type ToolOutputState = 'output-available' | 'output-error'
+
 /**
  * JSON Schema type for defining tool input/output schemas as raw JSON Schema objects.
  * This allows tools to be defined without schema libraries when you have JSON Schema definitions available.
@@ -355,7 +367,7 @@ export interface ToolCallPart<TMetadata = unknown> {
 export interface ToolResultPart {
   type: 'tool-result'
   toolCallId: string
-  content: string
+  content: string | Array<ContentPart>
   state: ToolResultState
   error?: string // Error message if state is "error"
 }
@@ -443,33 +455,75 @@ export type ConstrainedModelMessage<
   content: ConstrainedContent<TInputModalitiesTypes>
 }
 
+type IsUnknown<T> = unknown extends T
+  ? [T] extends [unknown]
+    ? true
+    : false
+  : false
+
+type RuntimeContextField<TContext> =
+  IsUnknown<TContext> extends true
+    ? {
+        /**
+         * Runtime context provided by the caller.
+         *
+         * This is request-local application state for tool and middleware
+         * implementations, not the AG-UI `Context[]` protocol field.
+         */
+        context?: TContext
+      }
+    : {
+        /**
+         * Runtime context provided by the caller.
+         *
+         * This is request-local application state for tool and middleware
+         * implementations, not the AG-UI `Context[]` protocol field.
+         */
+        context: TContext
+      }
+
 /**
  * Context passed to tool execute functions, providing capabilities like
  * emitting custom events during execution.
  */
-export interface ToolExecutionContext {
-  /** The ID of the tool call being executed */
-  toolCallId?: string
-  /**
-   * Emit a custom event during tool execution.
-   * Events are streamed to the client in real-time as AG-UI CUSTOM events.
-   *
-   * @param eventName - Name of the custom event
-   * @param value - Event payload value
-   *
-   * @example
-   * ```ts
-   * const tool = toolDefinition({ ... }).server(async (args, context) => {
-   *   context?.emitCustomEvent('progress', { step: 1, total: 3 })
-   *   // ... do work ...
-   *   context?.emitCustomEvent('progress', { step: 2, total: 3 })
-   *   // ... do more work ...
-   *   return result
-   * })
-   * ```
-   */
-  emitCustomEvent: (eventName: string, value: Record<string, any>) => void
-}
+export type ToolExecutionContext<TContext = unknown> =
+  RuntimeContextField<TContext> & {
+    /** The ID of the tool call being executed */
+    toolCallId?: string
+    /**
+     * Emit a custom event during tool execution.
+     * Events are streamed to the client in real-time as AG-UI CUSTOM events.
+     *
+     * @param eventName - Name of the custom event
+     * @param value - Event payload value
+     *
+     * @example
+     * ```ts
+     * const tool = toolDefinition({ ... }).server(async (args, context) => {
+     *   context?.emitCustomEvent('progress', { step: 1, total: 3 })
+     *   // ... do work ...
+     *   context?.emitCustomEvent('progress', { step: 2, total: 3 })
+     *   // ... do more work ...
+     *   return result
+     * })
+     * ```
+     */
+    emitCustomEvent: (eventName: string, value: Record<string, any>) => void
+  }
+
+export type ToolExecuteFunction<
+  TInput extends SchemaInput = SchemaInput,
+  TOutput extends SchemaInput = SchemaInput,
+  TContext = unknown,
+> = undefined extends TContext
+  ? (
+      args: InferSchemaType<TInput>,
+      context?: ToolExecutionContext<TContext>,
+    ) => Promise<InferSchemaType<TOutput>> | InferSchemaType<TOutput>
+  : (
+      args: InferSchemaType<TInput>,
+      context: ToolExecutionContext<TContext>,
+    ) => Promise<InferSchemaType<TOutput>> | InferSchemaType<TOutput>
 
 /**
  * Tool/Function definition for function calling.
@@ -488,6 +542,7 @@ export interface Tool<
   TInput extends SchemaInput = SchemaInput,
   TOutput extends SchemaInput = SchemaInput,
   TName extends string = string,
+  TContext = unknown,
 > {
   /**
    * Unique name of the tool (used by the model to call it).
@@ -587,9 +642,7 @@ export interface Tool<
    *   return weather; // Can return object or string
    * }
    */
-  execute?:
-    | ((args: any, context?: ToolExecutionContext) => Promise<any> | any)
-    | undefined
+  execute?: ToolExecuteFunction<TInput, TOutput, TContext> | undefined
 
   /** If true, tool execution requires user approval before running. Works with both server and client tools. */
   needsApproval?: boolean
@@ -601,6 +654,10 @@ export interface Tool<
   metadata?: Record<string, any> | undefined
 }
 
+export type AnyTool = Omit<Tool<any, any, any, any>, 'execute'> & {
+  execute?: ((args: any, context?: any) => any) | undefined
+}
+
 export interface ToolConfig {
   [key: string]: Tool
 }
@@ -729,10 +786,16 @@ export type AgentLoopStrategy = (state: AgentLoopState) => boolean
 export interface TextOptions<
   TProviderOptionsSuperset extends Record<string, any> = Record<string, any>,
   TProviderOptionsForModel = TProviderOptionsSuperset,
+  TContext = unknown,
 > {
   model: string
   messages: Array<ModelMessage>
-  tools?: Array<Tool<any, any, any>> | undefined
+  tools?: Array<AnyTool> | undefined
+  /**
+   * Runtime context provided by the caller and passed to middleware and
+   * server-side tool implementations.
+   */
+  context?: TContext
   /**
    * System prompts to include with the request.
    *
@@ -925,38 +988,22 @@ export interface RunStartedEvent extends AGUIRunStartedEvent {
   model?: string
 }
 
-/**
- * Provider-reported cost breakdown for a single request, normalized onto a
- * canonical shape so consumer code is portable across gateways. Each adapter's
- * extractor maps its provider-specific wire keys (e.g. OpenRouter's
- * `upstream_inference_prompt_cost`, `upstream_inference_input_cost`) onto these
- * fields at runtime.
- */
-export interface UsageCostBreakdown {
-  /** Total cost the gateway paid the upstream provider. */
-  upstreamCost?: number
-  /** Upstream cost for input (prompt) tokens. */
-  upstreamInputCost?: number
-  /** Upstream cost for output (completion) tokens. */
-  upstreamOutputCost?: number
+// Re-export the canonical usage types (defined in `@tanstack/ai-event-client`)
+// so `@tanstack/ai` consumers keep importing them from here unchanged.
+export type {
+  CompletionTokensDetails,
+  PromptTokensDetails,
+  ProviderUsageDetails,
+  TokenUsage,
+  UsageCostBreakdown,
 }
 
 /**
- * Token usage totals for a run, optionally including provider-reported cost.
- *
- * `cost` and `costDetails` are populated only by adapters whose provider returns
- * authoritative per-request cost (e.g. OpenRouter). They are absent for adapters
- * that do not report cost, so consumers must treat them as optional.
+ * @deprecated Renamed to {@link TokenUsage}. Kept as an alias for backward
+ * compatibility with `@tanstack/ai@0.23` and earlier; will be removed in a
+ * future release.
  */
-export interface UsageTotals {
-  promptTokens: number
-  completionTokens: number
-  totalTokens: number
-  /** Provider-reported cost for the request, when available. */
-  cost?: number
-  /** Provider-reported cost breakdown, when available. */
-  costDetails?: UsageCostBreakdown
-}
+export type UsageTotals = TokenUsage
 
 /**
  * Emitted when a run completes successfully.
@@ -969,8 +1016,8 @@ export interface RunFinishedEvent extends AGUIRunFinishedEvent {
   model?: string
   /** Why the generation stopped */
   finishReason?: 'stop' | 'length' | 'content_filter' | 'tool_calls' | null
-  /** Token usage statistics, optionally including provider-reported cost. */
-  usage?: UsageTotals
+  /** Token usage statistics with optional detailed breakdowns and provider-reported cost. */
+  usage?: TokenUsage
 }
 
 /**
@@ -1084,7 +1131,9 @@ export interface ToolCallEndEvent extends AGUIToolCallEndEvent {
   /** Final parsed input arguments (TanStack AI internal) */
   input?: unknown
   /** Tool execution result (TanStack AI internal) */
-  result?: string
+  result?: string | Array<ContentPart>
+  /** Tool execution output state (TanStack AI internal) */
+  state?: ToolOutputState
 }
 
 /**
@@ -1096,6 +1145,8 @@ export interface ToolCallEndEvent extends AGUIToolCallEndEvent {
 export interface ToolCallResultEvent extends AGUIToolCallResultEvent {
   /** Model identifier for multi-model support */
   model?: string
+  /** Tool execution output state (TanStack AI internal) */
+  state?: ToolOutputState
 }
 
 /**
@@ -1416,11 +1467,7 @@ export interface TextCompletionChunk {
   content: string
   role?: 'assistant'
   finishReason?: 'stop' | 'length' | 'content_filter' | null
-  usage?: {
-    promptTokens: number
-    completionTokens: number
-    totalTokens: number
-  }
+  usage?: TokenUsage
 }
 
 export interface SummarizationOptions<
@@ -1444,11 +1491,7 @@ export interface SummarizationResult {
   id: string
   model: string
   summary: string
-  usage: {
-    promptTokens: number
-    completionTokens: number
-    totalTokens: number
-  }
+  usage: TokenUsage
 }
 
 // ============================================================================
@@ -1517,11 +1560,7 @@ export interface ImageGenerationResult {
   /** Array of generated images */
   images: Array<GeneratedImage>
   /** Token usage information (if available) */
-  usage?: {
-    inputTokens?: number
-    outputTokens?: number
-    totalTokens?: number
-  }
+  usage?: TokenUsage
 }
 
 // ============================================================================
@@ -1572,11 +1611,7 @@ export interface AudioGenerationResult {
   /** The generated audio */
   audio: GeneratedAudio
   /** Token usage information (if available) */
-  usage?: {
-    inputTokens?: number
-    outputTokens?: number
-    totalTokens?: number
-  }
+  usage?: TokenUsage
 }
 
 // ============================================================================
@@ -1697,6 +1732,8 @@ export interface TTSResult {
   duration?: number
   /** Content type of the audio (e.g., 'audio/mp3') */
   contentType?: string
+  /** Token usage information (if provided by the adapter) */
+  usage?: TokenUsage
 }
 
 // ============================================================================
@@ -1778,6 +1815,8 @@ export interface TranscriptionResult {
   segments?: Array<TranscriptionSegment>
   /** Word-level timestamps, if available */
   words?: Array<TranscriptionWord>
+  /** Token usage information (if provided by the adapter) */
+  usage?: TokenUsage
 }
 
 /**

package/src/utilities/ag-ui-wire.ts

@@ -103,7 +103,10 @@ export function uiMessagesToWire(
           role: 'tool',
           id: deriveToolMessageId(part.toolCallId),
           toolCallId: part.toolCallId,
-          content: part.content,
+          content:
+            typeof part.content === 'string'
+              ? part.content
+              : JSON.stringify(part.content),
           ...(part.error !== undefined && { error: part.error }),
         })
       }

package/src/utilities/chat-params.ts

@@ -1,6 +1,12 @@
 import { AGUIError, RunAgentInputSchema } from '@ag-ui/core'
 import type { Context as AGUIContext } from '@ag-ui/core'
-import type { JSONSchema, ModelMessage, Tool, UIMessage } from '../types'
+import type {
+  JSONSchema,
+  ModelMessage,
+  SchemaInput,
+  Tool,
+  UIMessage,
+} from '../types'
 
 const KNOWN_PART_TYPES = new Set([
   'text',
@@ -43,7 +49,12 @@ export function chatParamsFromRequestBody(body: unknown): Promise<{
   tools: Array<{ name: string; description: string; parameters: JSONSchema }>
   forwardedProps: Record<string, unknown>
   state: unknown
+  /**
+   * @deprecated Use `aguiContext` instead. This alias will be removed in a
+   * future release.
+   */
   context: Array<AGUIContext>
+  aguiContext: Array<AGUIContext>
 }> {
   const parseResult = RunAgentInputSchema.safeParse(body)
   if (!parseResult.success) {
@@ -58,6 +69,7 @@ export function chatParamsFromRequestBody(body: unknown): Promise<{
   }
 
   const parsed = parseResult.data
+  const aguiContext = parsed.context
 
   // AG-UI Zod uses `.strip()` so extra fields like `parts` on messages are
   // dropped during parse. We re-attach them from the original body so the
@@ -89,7 +101,8 @@ export function chatParamsFromRequestBody(body: unknown): Promise<{
     }>,
     forwardedProps: (parsed.forwardedProps ?? {}) as Record<string, unknown>,
     state: parsed.state,
-    context: parsed.context,
+    context: aguiContext,
+    aguiContext,
   })
 }
 
@@ -171,16 +184,18 @@ export async function chatParamsFromRequest(
  *   `chatParamsFromRequest(...)` / `chatParamsFromRequestBody(...)`.
  * @returns A merged array suitable for `chat({ tools })`.
  */
-export function mergeAgentTools(
-  serverTools: ReadonlyArray<Tool>,
+export function mergeAgentTools<TContext = unknown>(
+  serverTools: ReadonlyArray<Tool<SchemaInput, SchemaInput, string, TContext>>,
   clientTools: ReadonlyArray<{
     name: string
     description: string
     parameters: JSONSchema
   }>,
-): Array<Tool> {
+): Array<Tool<SchemaInput, SchemaInput, string, TContext>> {
   const seen = new Set(serverTools.map((t) => t.name))
-  const merged: Array<Tool> = [...serverTools]
+  const merged: Array<Tool<SchemaInput, SchemaInput, string, TContext>> = [
+    ...serverTools,
+  ]
   for (const ct of clientTools) {
     if (seen.has(ct.name)) {
       // Server wins on name collision.
@@ -193,7 +208,7 @@ export function mergeAgentTools(
       inputSchema: ct.parameters,
       // No `execute` — runtime treats this as a client-side tool and
       // emits ClientToolRequest events.
-    } as Tool)
+    } as Tool<SchemaInput, SchemaInput, string, TContext>)
   }
   return merged
 }

package/src/utilities/tool-result.ts

@@ -0,0 +1,60 @@
+import type { ContentPart } from '../types'
+
+const CONTENT_PART_TYPES = new Set([
+  'text',
+  'image',
+  'audio',
+  'video',
+  'document',
+])
+
+/**
+ * Structural check for a single `ContentPart`. A text part must carry a string
+ * `content`; every other modality must carry a `source` with `type` of
+ * `'url' | 'data'` and a string `value`.
+ */
+export function isContentPart(value: unknown): value is ContentPart {
+  if (typeof value !== 'object' || value === null) return false
+  const part = value as Record<string, unknown>
+  if (typeof part.type !== 'string' || !CONTENT_PART_TYPES.has(part.type)) {
+    return false
+  }
+  if (part.type === 'text') {
+    return typeof part.content === 'string'
+  }
+  const source = part.source
+  if (typeof source !== 'object' || source === null) return false
+  const src = source as Record<string, unknown>
+  if (typeof src.value !== 'string') return false
+  // `data` sources require a mimeType (matches ContentPartDataSource); `url`
+  // sources don't. Requiring it here keeps the runtime guard consistent with
+  // the type and avoids emitting `data:undefined;base64,...` downstream.
+  if (src.type === 'data') return typeof src.mimeType === 'string'
+  return src.type === 'url'
+}
+
+/**
+ * True iff `value` is a NON-EMPTY array whose every element is a valid
+ * `ContentPart`. Empty arrays and mixed arrays return false so they continue
+ * to be treated as ordinary (stringified) data — this keeps the auto-detection
+ * footgun narrow.
+ */
+export function isContentPartArray(
+  value: unknown,
+): value is Array<ContentPart> {
+  return Array.isArray(value) && value.length > 0 && value.every(isContentPart)
+}
+
+/**
+ * Normalize a tool's return value for transport:
+ * - string            → unchanged
+ * - ContentPart array → unchanged (multimodal, passed through to the adapter)
+ * - anything else     → `JSON.stringify`
+ */
+export function normalizeToolResult(
+  result: unknown,
+): string | Array<ContentPart> {
+  if (typeof result === 'string') return result
+  if (isContentPartArray(result)) return result
+  return JSON.stringify(result)
+}

package/src/utilities/usage.ts

@@ -0,0 +1,41 @@
+import type { ProviderUsageDetails, TokenUsage } from '../types'
+
+/**
+ * Input parameters for building base TokenUsage.
+ * Provider functions should extract these from their SDK's response.
+ */
+export interface BaseUsageInput {
+  /** Total input/prompt tokens */
+  promptTokens: number
+  /** Total output/completion tokens */
+  completionTokens: number
+  /** Total tokens (prompt + completion) */
+  totalTokens: number
+}
+
+/**
+ * Builds the base TokenUsage object with core fields.
+ * Provider-specific functions should use this and then add their own details.
+ *
+ * @param input - The base token counts
+ * @returns A TokenUsage object with promptTokens, completionTokens, totalTokens
+ *
+ * @example
+ * ```typescript
+ * const base = buildBaseUsage({
+ *   promptTokens: 100,
+ *   completionTokens: 50,
+ *   totalTokens: 150
+ * });
+ * // Returns: { promptTokens: 100, completionTokens: 50, totalTokens: 150 }
+ * ```
+ */
+export function buildBaseUsage<TProviderDetails = ProviderUsageDetails>(
+  input: BaseUsageInput,
+): TokenUsage<TProviderDetails> {
+  return {
+    promptTokens: input.promptTokens,
+    completionTokens: input.completionTokens,
+    totalTokens: input.totalTokens,
+  }
+}