@tanstack/ai 0.6.2 → 0.8.0

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

@@ -10,6 +10,41 @@ import type {
   ToolCallStartEvent,
   ToolExecutionContext,
 } from '../../../types'
+import type {
+  AfterToolCallInfo,
+  BeforeToolCallDecision,
+} from '../middleware/types'
+
+function safeJsonParse(value: string): unknown {
+  try {
+    return JSON.parse(value)
+  } catch {
+    return value
+  }
+}
+
+/**
+ * Optional middleware hooks for tool execution.
+ * When provided, these callbacks are invoked before/after each tool execution.
+ */
+export interface ToolExecutionMiddlewareHooks {
+  onBeforeToolCall?: (
+    toolCall: ToolCall,
+    tool: Tool | undefined,
+    args: unknown,
+  ) => Promise<BeforeToolCallDecision>
+  onAfterToolCall?: (info: AfterToolCallInfo) => Promise<void>
+}
+
+/**
+ * Error thrown when middleware decides to abort the chat run during tool execution.
+ */
+export class MiddlewareAbortError extends Error {
+  constructor(reason: string) {
+    super(reason)
+    this.name = 'MiddlewareAbortError'
+  }
+}
 
 /**
  * Manages tool call accumulation and execution for the chat() method's automatic tool execution loop.
@@ -289,6 +324,152 @@ async function* executeWithEventPolling<T>(
   return state.result
 }
 
+/**
+ * Apply a middleware onBeforeToolCall decision.
+ * Returns the (possibly transformed) input if execution should proceed,
+ * or undefined if the tool call was skipped (result already pushed).
+ * Throws MiddlewareAbortError if the decision is 'abort'.
+ */
+async function applyBeforeToolCallDecision(
+  toolCall: ToolCall,
+  tool: Tool,
+  input: unknown,
+  toolName: string,
+  middlewareHooks: ToolExecutionMiddlewareHooks,
+  results: Array<ToolResult>,
+): Promise<{ proceed: true; input: unknown } | { proceed: false }> {
+  if (!middlewareHooks.onBeforeToolCall) {
+    return { proceed: true, input }
+  }
+
+  const decision = await middlewareHooks.onBeforeToolCall(toolCall, tool, input)
+  if (!decision) {
+    return { proceed: true, input }
+  }
+
+  if (decision.type === 'abort') {
+    throw new MiddlewareAbortError(decision.reason || 'Aborted by middleware')
+  }
+
+  if (decision.type === 'skip') {
+    const skipResult = decision.result
+    results.push({
+      toolCallId: toolCall.id,
+      toolName,
+      result:
+        typeof skipResult === 'string'
+          ? safeJsonParse(skipResult)
+          : skipResult || null,
+      duration: 0,
+    })
+    if (middlewareHooks.onAfterToolCall) {
+      await middlewareHooks.onAfterToolCall({
+        toolCall,
+        tool,
+        toolName,
+        toolCallId: toolCall.id,
+        ok: true,
+        duration: 0,
+        result: skipResult,
+      })
+    }
+    return { proceed: false }
+  }
+
+  return { proceed: true, input: decision.args }
+}
+
+/**
+ * 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(
+  toolCall: ToolCall,
+  tool: Tool,
+  toolName: string,
+  input: unknown,
+  context: ToolExecutionContext,
+  pendingEvents: Array<CustomEvent>,
+  results: Array<ToolResult>,
+  middlewareHooks?: ToolExecutionMiddlewareHooks,
+): AsyncGenerator<CustomEvent, void, void> {
+  const startTime = Date.now()
+  try {
+    const executionPromise = Promise.resolve(tool.execute!(input, context))
+    let result = yield* executeWithEventPolling(executionPromise, pendingEvents)
+    const duration = Date.now() - startTime
+
+    // Flush remaining events
+    while (pendingEvents.length > 0) {
+      yield pendingEvents.shift()!
+    }
+
+    // Validate output against outputSchema if provided
+    if (
+      tool.outputSchema &&
+      isStandardSchema(tool.outputSchema) &&
+      result !== undefined &&
+      result !== null
+    ) {
+      result = parseWithStandardSchema(tool.outputSchema, result)
+    }
+
+    const finalResult =
+      typeof result === 'string' ? safeJsonParse(result) : result || null
+
+    results.push({
+      toolCallId: toolCall.id,
+      toolName,
+      result: finalResult,
+      duration,
+    })
+
+    if (middlewareHooks?.onAfterToolCall) {
+      await middlewareHooks.onAfterToolCall({
+        toolCall,
+        tool,
+        toolName,
+        toolCallId: toolCall.id,
+        ok: true,
+        duration,
+        result: finalResult,
+      })
+    }
+  } catch (error: unknown) {
+    const duration = Date.now() - startTime
+
+    // Flush remaining events
+    while (pendingEvents.length > 0) {
+      yield pendingEvents.shift()!
+    }
+
+    if (error instanceof MiddlewareAbortError) {
+      throw error
+    }
+
+    const message = error instanceof Error ? error.message : 'Unknown error'
+    results.push({
+      toolCallId: toolCall.id,
+      toolName,
+      result: { error: message },
+      state: 'output-error',
+      duration,
+    })
+
+    if (middlewareHooks?.onAfterToolCall) {
+      await middlewareHooks.onAfterToolCall({
+        toolCall,
+        tool,
+        toolName,
+        toolCallId: toolCall.id,
+        ok: false,
+        duration,
+        error,
+      })
+    }
+  }
+}
+
 /**
  * Execute tool calls based on their configuration.
  * Yields CustomEvent chunks during tool execution for real-time progress updates.
@@ -313,6 +494,7 @@ export async function* executeToolCalls(
     eventName: string,
     value: Record<string, any>,
   ) => CustomEvent,
+  middlewareHooks?: ToolExecutionMiddlewareHooks,
 ): AsyncGenerator<CustomEvent, ExecuteToolCallsResult, void> {
   const results: Array<ToolResult> = []
   const needsApproval: Array<ApprovalRequest> = []
@@ -402,13 +584,6 @@ export async function* executeToolCalls(
       },
     }
 
-    // Helper to flush any pending events
-    function* flushEvents(): Generator<CustomEvent> {
-      while (pendingEvents.length > 0) {
-        yield pendingEvents.shift()!
-      }
-    }
-
     // CASE 1: Client-side tool (no execute function)
     if (!tool.execute) {
       // Check if tool needs approval
@@ -482,51 +657,30 @@ export async function* executeToolCalls(
         const approved = approvals.get(approvalId)
 
         if (approved) {
-          // Execute after approval
-          const startTime = Date.now()
-          try {
-            const executionPromise = Promise.resolve(
-              tool.execute(input, context),
-            )
-            let result = yield* executeWithEventPolling(
-              executionPromise,
-              pendingEvents,
-            )
-            const duration = Date.now() - startTime
-            yield* flushEvents()
-
-            // Validate output against outputSchema if provided (for Standard Schema compliant schemas)
-            if (
-              tool.outputSchema &&
-              isStandardSchema(tool.outputSchema) &&
-              result !== undefined &&
-              result !== null
-            ) {
-              result = parseWithStandardSchema(tool.outputSchema, result)
-            }
-
-            results.push({
-              toolCallId: toolCall.id,
-              toolName,
-              result:
-                typeof result === 'string'
-                  ? JSON.parse(result)
-                  : result || null,
-              duration,
-            })
-          } catch (error: unknown) {
-            const duration = Date.now() - startTime
-            yield* flushEvents()
-            const message =
-              error instanceof Error ? error.message : 'Unknown error'
-            results.push({
-              toolCallId: toolCall.id,
+          // Apply middleware before-hook for approved tools
+          if (middlewareHooks) {
+            const decision = await applyBeforeToolCallDecision(
+              toolCall,
+              tool,
+              input,
               toolName,
-              result: { error: message },
-              state: 'output-error',
-              duration,
-            })
+              middlewareHooks,
+              results,
+            )
+            if (!decision.proceed) continue
+            input = decision.input
           }
+
+          yield* executeServerTool(
+            toolCall,
+            tool,
+            toolName,
+            input,
+            context,
+            pendingEvents,
+            results,
+            middlewareHooks,
+          )
         } else {
           // User declined
           results.push({
@@ -549,45 +703,29 @@ export async function* executeToolCalls(
     }
 
     // CASE 3: Normal server tool - execute immediately
-    const startTime = Date.now()
-    try {
-      const executionPromise = Promise.resolve(tool.execute(input, context))
-      let result = yield* executeWithEventPolling(
-        executionPromise,
-        pendingEvents,
-      )
-      const duration = Date.now() - startTime
-      yield* flushEvents()
-
-      // Validate output against outputSchema if provided (for Standard Schema compliant schemas)
-      if (
-        tool.outputSchema &&
-        isStandardSchema(tool.outputSchema) &&
-        result !== undefined &&
-        result !== null
-      ) {
-        result = parseWithStandardSchema(tool.outputSchema, result)
-      }
-
-      results.push({
-        toolCallId: toolCall.id,
-        toolName,
-        result:
-          typeof result === 'string' ? JSON.parse(result) : result || null,
-        duration,
-      })
-    } catch (error: unknown) {
-      const duration = Date.now() - startTime
-      yield* flushEvents()
-      const message = error instanceof Error ? error.message : 'Unknown error'
-      results.push({
-        toolCallId: toolCall.id,
+    if (middlewareHooks) {
+      const decision = await applyBeforeToolCallDecision(
+        toolCall,
+        tool,
+        input,
         toolName,
-        result: { error: message },
-        state: 'output-error',
-        duration,
-      })
+        middlewareHooks,
+        results,
+      )
+      if (!decision.proceed) continue
+      input = decision.input
     }
+
+    yield* executeServerTool(
+      toolCall,
+      tool,
+      toolName,
+      input,
+      context,
+      pendingEvents,
+      results,
+      middlewareHooks,
+    )
   }
 
   return { results, needsApproval, needsClientExecution }

package/src/activities/generateImage/index.ts

@@ -5,9 +5,10 @@
  * This is a self-contained module with implementation, types, and JSDoc.
  */
 
-import { aiEventClient } from '../../event-client.js'
+import { aiEventClient } from '@tanstack/ai-event-client'
+import { streamGenerationResult } from '../stream-generation-result.js'
 import type { ImageAdapter } from './adapter'
-import type { ImageGenerationResult } from '../../types'
+import type { ImageGenerationResult, StreamChunk } from '../../types'
 
 // ===========================
 // Activity Kind
@@ -60,9 +61,11 @@ export type ImageSizeForModel<TAdapter, TModel extends string> =
  * The model is extracted from the adapter's model property.
  *
  * @template TAdapter - The image adapter type
+ * @template TStream - Whether to stream the output
  */
 export type ImageActivityOptions<
   TAdapter extends ImageAdapter<string, any, any, any>,
+  TStream extends boolean = false,
 > = {
   /** The image adapter to use (must be created with a model) */
   adapter: TAdapter & { kind: typeof kind }
@@ -72,6 +75,14 @@ export type ImageActivityOptions<
   numberOfImages?: number
   /** Image size in WIDTHxHEIGHT format (e.g., "1024x1024") */
   size?: ImageSizeForModel<TAdapter, TAdapter['model']>
+  /**
+   * Whether to stream the image generation result.
+   * When true, returns an AsyncIterable<StreamChunk> for streaming transport.
+   * When false or not provided, returns a Promise<ImageGenerationResult>.
+   *
+   * @default false
+   */
+  stream?: TStream
 } & ({} extends ImageProviderOptionsForModel<TAdapter, TAdapter['model']>
   ? {
       /** Provider-specific options for image generation */ modelOptions?: ImageProviderOptionsForModel<
@@ -90,8 +101,15 @@ export type ImageActivityOptions<
 // Activity Result Type
 // ===========================
 
-/** Result type for the image activity */
-export type ImageActivityResult = Promise<ImageGenerationResult>
+/**
+ * Result type for the image activity.
+ * - If stream is true: AsyncIterable<StreamChunk>
+ * - Otherwise: Promise<ImageGenerationResult>
+ */
+export type ImageActivityResult<TStream extends boolean = false> =
+  TStream extends true
+    ? AsyncIterable<StreamChunk>
+    : Promise<ImageGenerationResult>
 
 function createId(prefix: string): string {
   return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
@@ -146,10 +164,31 @@ function createId(prefix: string): string {
  * })
  * ```
  */
-export async function generateImage<
+export function generateImage<
+  TAdapter extends ImageAdapter<string, any, any, any>,
+  TStream extends boolean = false,
+>(
+  options: ImageActivityOptions<TAdapter, TStream>,
+): ImageActivityResult<TStream> {
+  if (options.stream) {
+    return streamGenerationResult(() =>
+      runGenerateImage(options),
+    ) as ImageActivityResult<TStream>
+  }
+
+  return runGenerateImage(options) as ImageActivityResult<TStream>
+}
+
+/**
+ * Internal implementation of image generation (always non-streaming).
+ * Contains all devtools event emission logic.
+ */
+async function runGenerateImage<
   TAdapter extends ImageAdapter<string, any, any, any>,
->(options: ImageActivityOptions<TAdapter>): ImageActivityResult {
-  const { adapter, ...rest } = options
+>(
+  options: ImageActivityOptions<TAdapter, boolean>,
+): Promise<ImageGenerationResult> {
+  const { adapter, stream: _stream, ...rest } = options
   const model = adapter.model
   const requestId = createId('image')
   const startTime = Date.now()
@@ -204,7 +243,10 @@ export async function generateImage<
  */
 export function createImageOptions<
   TAdapter extends ImageAdapter<string, any, any, any>,
->(options: ImageActivityOptions<TAdapter>): ImageActivityOptions<TAdapter> {
+  TStream extends boolean = false,
+>(
+  options: ImageActivityOptions<TAdapter, TStream>,
+): ImageActivityOptions<TAdapter, TStream> {
   return options
 }
 

package/src/activities/generateSpeech/index.ts

@@ -5,9 +5,10 @@
  * This is a self-contained module with implementation, types, and JSDoc.
  */
 
-import { aiEventClient } from '../../event-client.js'
+import { aiEventClient } from '@tanstack/ai-event-client'
+import { streamGenerationResult } from '../stream-generation-result.js'
 import type { TTSAdapter } from './adapter'
-import type { TTSResult } from '../../types'
+import type { StreamChunk, TTSResult } from '../../types'
 
 // ===========================
 // Activity Kind
@@ -37,9 +38,11 @@ export type TTSProviderOptions<TAdapter> =
  * The model is extracted from the adapter's model property.
  *
  * @template TAdapter - The TTS adapter type
+ * @template TStream - Whether to stream the output
  */
 export interface TTSActivityOptions<
   TAdapter extends TTSAdapter<string, object>,
+  TStream extends boolean = false,
 > {
   /** The TTS adapter to use (must be created with a model) */
   adapter: TAdapter & { kind: typeof kind }
@@ -53,14 +56,27 @@ export interface TTSActivityOptions<
   speed?: number
   /** Provider-specific options for TTS generation */
   modelOptions?: TTSProviderOptions<TAdapter>
+  /**
+   * Whether to stream the generation result.
+   * When true, returns an AsyncIterable<StreamChunk> for streaming transport.
+   * When false or not provided, returns a Promise<TTSResult>.
+   *
+   * @default false
+   */
+  stream?: TStream
 }
 
 // ===========================
 // Activity Result Type
 // ===========================
 
-/** Result type for the TTS activity */
-export type TTSActivityResult = Promise<TTSResult>
+/**
+ * Result type for the TTS activity.
+ * - If stream is true: AsyncIterable<StreamChunk>
+ * - Otherwise: Promise<TTSResult>
+ */
+export type TTSActivityResult<TStream extends boolean = false> =
+  TStream extends true ? AsyncIterable<StreamChunk> : Promise<TTSResult>
 
 function createId(prefix: string): string {
   return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
@@ -100,10 +116,25 @@ function createId(prefix: string): string {
  * })
  * ```
  */
-export async function generateSpeech<
+export function generateSpeech<
   TAdapter extends TTSAdapter<string, object>,
->(options: TTSActivityOptions<TAdapter>): TTSActivityResult {
-  const { adapter, ...rest } = options
+  TStream extends boolean = false,
+>(options: TTSActivityOptions<TAdapter, TStream>): TTSActivityResult<TStream> {
+  if (options.stream) {
+    return streamGenerationResult(() =>
+      runGenerateSpeech(options),
+    ) as TTSActivityResult<TStream>
+  }
+  return runGenerateSpeech(options) as TTSActivityResult<TStream>
+}
+
+/**
+ * Run the core TTS generation logic (non-streaming).
+ */
+async function runGenerateSpeech<TAdapter extends TTSAdapter<string, object>>(
+  options: TTSActivityOptions<TAdapter, boolean>,
+): Promise<TTSResult> {
+  const { adapter, stream: _stream, ...rest } = options
   const model = adapter.model
   const requestId = createId('speech')
   const startTime = Date.now()
@@ -149,7 +180,10 @@ export async function generateSpeech<
  */
 export function createSpeechOptions<
   TAdapter extends TTSAdapter<string, object>,
->(options: TTSActivityOptions<TAdapter>): TTSActivityOptions<TAdapter> {
+  TStream extends boolean = false,
+>(
+  options: TTSActivityOptions<TAdapter, TStream>,
+): TTSActivityOptions<TAdapter, TStream> {
   return options
 }
 

package/src/activities/generateTranscription/index.ts

@@ -5,9 +5,10 @@
  * This is a self-contained module with implementation, types, and JSDoc.
  */
 
-import { aiEventClient } from '../../event-client.js'
+import { aiEventClient } from '@tanstack/ai-event-client'
+import { streamGenerationResult } from '../stream-generation-result.js'
 import type { TranscriptionAdapter } from './adapter'
-import type { TranscriptionResult } from '../../types'
+import type { StreamChunk, TranscriptionResult } from '../../types'
 
 // ===========================
 // Activity Kind
@@ -37,9 +38,11 @@ export type TranscriptionProviderOptions<TAdapter> =
  * The model is extracted from the adapter's model property.
  *
  * @template TAdapter - The transcription adapter type
+ * @template TStream - Whether to stream the output
  */
 export interface TranscriptionActivityOptions<
   TAdapter extends TranscriptionAdapter<string, object>,
+  TStream extends boolean = false,
 > {
   /** The transcription adapter to use (must be created with a model) */
   adapter: TAdapter & { kind: typeof kind }
@@ -53,14 +56,29 @@ export interface TranscriptionActivityOptions<
   responseFormat?: 'json' | 'text' | 'srt' | 'verbose_json' | 'vtt'
   /** Provider-specific options for transcription */
   modelOptions?: TranscriptionProviderOptions<TAdapter>
+  /**
+   * Whether to stream the transcription result.
+   * When true, returns an AsyncIterable<StreamChunk> for streaming transport.
+   * When false or not provided, returns a Promise<TranscriptionResult>.
+   *
+   * @default false
+   */
+  stream?: TStream
 }
 
 // ===========================
 // Activity Result Type
 // ===========================
 
-/** Result type for the transcription activity */
-export type TranscriptionActivityResult = Promise<TranscriptionResult>
+/**
+ * Result type for the transcription activity.
+ * - If stream is true: AsyncIterable<StreamChunk>
+ * - Otherwise: Promise<TranscriptionResult>
+ */
+export type TranscriptionActivityResult<TStream extends boolean = false> =
+  TStream extends true
+    ? AsyncIterable<StreamChunk>
+    : Promise<TranscriptionResult>
 
 function createId(prefix: string): string {
   return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
@@ -101,13 +119,44 @@ function createId(prefix: string): string {
  *   console.log(`[${segment.start}s - ${segment.end}s]: ${segment.text}`)
  * })
  * ```
+ *
+ * @example Streaming transcription result
+ * ```ts
+ * for await (const chunk of generateTranscription({
+ *   adapter: openaiTranscription('whisper-1'),
+ *   audio: audioFile,
+ *   stream: true
+ * })) {
+ *   console.log(chunk)
+ * }
+ * ```
+ */
+export function generateTranscription<
+  TAdapter extends TranscriptionAdapter<string, object>,
+  TStream extends boolean = false,
+>(
+  options: TranscriptionActivityOptions<TAdapter, TStream>,
+): TranscriptionActivityResult<TStream> {
+  if (options.stream) {
+    return streamGenerationResult(() =>
+      runGenerateTranscription(options),
+    ) as TranscriptionActivityResult<TStream>
+  }
+
+  return runGenerateTranscription(
+    options,
+  ) as TranscriptionActivityResult<TStream>
+}
+
+/**
+ * Run non-streaming transcription
  */
-export async function generateTranscription<
+async function runGenerateTranscription<
   TAdapter extends TranscriptionAdapter<string, object>,
 >(
-  options: TranscriptionActivityOptions<TAdapter>,
-): TranscriptionActivityResult {
-  const { adapter, ...rest } = options
+  options: TranscriptionActivityOptions<TAdapter, boolean>,
+): Promise<TranscriptionResult> {
+  const { adapter, stream: _stream, ...rest } = options
   const model = adapter.model
   const requestId = createId('transcription')
   const startTime = Date.now()
@@ -149,9 +198,10 @@ export async function generateTranscription<
  */
 export function createTranscriptionOptions<
   TAdapter extends TranscriptionAdapter<string, object>,
+  TStream extends boolean = false,
 >(
-  options: TranscriptionActivityOptions<TAdapter>,
-): TranscriptionActivityOptions<TAdapter> {
+  options: TranscriptionActivityOptions<TAdapter, TStream>,
+): TranscriptionActivityOptions<TAdapter, TStream> {
   return options
 }