@tanstack/ai 0.2.2 → 0.3.1

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

@@ -1,20 +1,22 @@
 import { isStandardSchema, parseWithStandardSchema } from './schema-converter'
 import type {
-  DoneStreamChunk,
   ModelMessage,
+  RunFinishedEvent,
   Tool,
   ToolCall,
-  ToolResultStreamChunk,
+  ToolCallArgsEvent,
+  ToolCallEndEvent,
+  ToolCallStartEvent,
 } from '../../../types'
 
 /**
  * Manages tool call accumulation and execution for the chat() method's automatic tool execution loop.
  *
  * Responsibilities:
- * - Accumulates streaming tool call chunks (ID, name, arguments)
+ * - Accumulates streaming tool call events (ID, name, arguments)
  * - Validates tool calls (filters out incomplete ones)
  * - Executes tool `execute` functions with parsed arguments
- * - Emits `tool_result` chunks for client visibility
+ * - Emits `TOOL_CALL_END` events for client visibility
  * - Returns tool result messages for conversation history
  *
  * This class is used internally by the AI.chat() method to handle the automatic
@@ -26,14 +28,16 @@ import type {
  *
  * // During streaming, accumulate tool calls
  * for await (const chunk of stream) {
- *   if (chunk.type === "tool_call") {
- *     manager.addToolCallChunk(chunk);
+ *   if (chunk.type === 'TOOL_CALL_START') {
+ *     manager.addToolCallStartEvent(chunk);
+ *   } else if (chunk.type === 'TOOL_CALL_ARGS') {
+ *     manager.addToolCallArgsEvent(chunk);
  *   }
  * }
  *
  * // After stream completes, execute tools
  * if (manager.hasToolCalls()) {
- *   const toolResults = yield* manager.executeTools(doneChunk);
+ *   const toolResults = yield* manager.executeTools(finishEvent);
  *   messages = [...messages, ...toolResults];
  *   manager.clear();
  * }
@@ -48,43 +52,44 @@ export class ToolCallManager {
   }
 
   /**
-   * Add a tool call chunk to the accumulator
-   * Handles streaming tool calls by accumulating arguments
+   * Add a TOOL_CALL_START event to begin tracking a tool call (AG-UI)
    */
-  addToolCallChunk(chunk: {
-    toolCall: {
-      id: string
-      type: 'function'
+  addToolCallStartEvent(event: ToolCallStartEvent): void {
+    const index = event.index ?? this.toolCallsMap.size
+    this.toolCallsMap.set(index, {
+      id: event.toolCallId,
+      type: 'function',
       function: {
-        name: string
-        arguments: string
+        name: event.toolName,
+        arguments: '',
+      },
+    })
+  }
+
+  /**
+   * Add a TOOL_CALL_ARGS event to accumulate arguments (AG-UI)
+   */
+  addToolCallArgsEvent(event: ToolCallArgsEvent): void {
+    // Find the tool call by ID
+    for (const [, toolCall] of this.toolCallsMap.entries()) {
+      if (toolCall.id === event.toolCallId) {
+        toolCall.function.arguments += event.delta
+        break
       }
     }
-    index: number
-  }): void {
-    const index = chunk.index
-    const existing = this.toolCallsMap.get(index)
-
-    if (!existing) {
-      // Only create entry if we have a tool call ID and name
-      if (chunk.toolCall.id && chunk.toolCall.function.name) {
-        this.toolCallsMap.set(index, {
-          id: chunk.toolCall.id,
-          type: 'function',
-          function: {
-            name: chunk.toolCall.function.name,
-            arguments: chunk.toolCall.function.arguments || '',
-          },
-        })
-      }
-    } else {
-      // Update name if it wasn't set before
-      if (chunk.toolCall.function.name && !existing.function.name) {
-        existing.function.name = chunk.toolCall.function.name
-      }
-      // Accumulate arguments for streaming tool calls
-      if (chunk.toolCall.function.arguments) {
-        existing.function.arguments += chunk.toolCall.function.arguments
+  }
+
+  /**
+   * Complete a tool call with its final input
+   * Called when TOOL_CALL_END is received
+   */
+  completeToolCall(event: ToolCallEndEvent): void {
+    for (const [, toolCall] of this.toolCallsMap.entries()) {
+      if (toolCall.id === event.toolCallId) {
+        if (event.input !== undefined) {
+          toolCall.function.arguments = JSON.stringify(event.input)
+        }
+        break
       }
     }
   }
@@ -107,11 +112,12 @@ export class ToolCallManager {
 
   /**
    * Execute all tool calls and return tool result messages
-   * Also yields tool_result chunks for streaming
+   * Yields TOOL_CALL_END events for streaming
+   * @param finishEvent - RUN_FINISHED event from the stream
    */
   async *executeTools(
-    doneChunk: DoneStreamChunk,
-  ): AsyncGenerator<ToolResultStreamChunk, Array<ModelMessage>, void> {
+    finishEvent: RunFinishedEvent,
+  ): AsyncGenerator<ToolCallEndEvent, Array<ModelMessage>, void> {
     const toolCallsArray = this.getToolCalls()
     const toolResults: Array<ModelMessage> = []
 
@@ -182,14 +188,14 @@ export class ToolCallManager {
         toolResultContent = `Tool ${toolCall.function.name} does not have an execute function`
       }
 
-      // Emit tool_result chunk so callers can track tool execution
+      // Emit TOOL_CALL_END event
       yield {
-        type: 'tool_result',
-        id: doneChunk.id,
-        model: doneChunk.model,
-        timestamp: Date.now(),
+        type: 'TOOL_CALL_END',
         toolCallId: toolCall.id,
-        content: toolResultContent,
+        toolName: toolCall.function.name,
+        model: finishEvent.model,
+        timestamp: Date.now(),
+        result: toolResultContent,
       }
 
       // Add tool result message

package/src/activities/generateImage/index.ts

@@ -5,6 +5,7 @@
  * This is a self-contained module with implementation, types, and JSDoc.
  */
 
+import { aiEventClient } from '../../event-client.js'
 import type { ImageAdapter } from './adapter'
 import type { ImageGenerationResult } from '../../types'
 
@@ -82,6 +83,10 @@ export interface ImageActivityOptions<
 /** Result type for the image activity */
 export type ImageActivityResult = Promise<ImageGenerationResult>
 
+function createId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
+}
+
 // ===========================
 // Activity Implementation
 // ===========================
@@ -136,8 +141,48 @@ export async function generateImage<
 >(options: ImageActivityOptions<TAdapter>): ImageActivityResult {
   const { adapter, ...rest } = options
   const model = adapter.model
-
-  return adapter.generateImages({ ...rest, model })
+  const requestId = createId('image')
+  const startTime = Date.now()
+
+  aiEventClient.emit('image:request:started', {
+    requestId,
+    provider: adapter.name,
+    model,
+    prompt: rest.prompt,
+    numberOfImages: rest.numberOfImages,
+    size: rest.size as string | undefined,
+    modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+    timestamp: startTime,
+  })
+
+  return adapter.generateImages({ ...rest, model }).then((result) => {
+    const duration = Date.now() - startTime
+
+    aiEventClient.emit('image:request:completed', {
+      requestId,
+      provider: adapter.name,
+      model,
+      images: result.images.map((image) => ({
+        url: image.url,
+        b64Json: image.b64Json,
+      })),
+      duration,
+      modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+      timestamp: Date.now(),
+    })
+
+    if (result.usage) {
+      aiEventClient.emit('image:usage', {
+        requestId,
+        model,
+        usage: result.usage,
+        modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+        timestamp: Date.now(),
+      })
+    }
+
+    return result
+  })
 }
 
 // ===========================

package/src/activities/generateSpeech/index.ts

@@ -5,6 +5,7 @@
  * This is a self-contained module with implementation, types, and JSDoc.
  */
 
+import { aiEventClient } from '../../event-client.js'
 import type { TTSAdapter } from './adapter'
 import type { TTSResult } from '../../types'
 
@@ -61,6 +62,10 @@ export interface TTSActivityOptions<
 /** Result type for the TTS activity */
 export type TTSActivityResult = Promise<TTSResult>
 
+function createId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
+}
+
 // ===========================
 // Activity Implementation
 // ===========================
@@ -100,8 +105,39 @@ export async function generateSpeech<
 >(options: TTSActivityOptions<TAdapter>): TTSActivityResult {
   const { adapter, ...rest } = options
   const model = adapter.model
+  const requestId = createId('speech')
+  const startTime = Date.now()
+
+  aiEventClient.emit('speech:request:started', {
+    requestId,
+    provider: adapter.name,
+    model,
+    text: rest.text,
+    voice: rest.voice,
+    format: rest.format,
+    speed: rest.speed,
+    modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+    timestamp: startTime,
+  })
+
+  return adapter.generateSpeech({ ...rest, model }).then((result) => {
+    const duration = Date.now() - startTime
+
+    aiEventClient.emit('speech:request:completed', {
+      requestId,
+      provider: adapter.name,
+      model,
+      audio: result.audio,
+      format: result.format,
+      audioDuration: result.duration,
+      contentType: result.contentType,
+      duration,
+      modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+      timestamp: Date.now(),
+    })
 
-  return adapter.generateSpeech({ ...rest, model })
+    return result
+  })
 }
 
 // ===========================

package/src/activities/generateTranscription/index.ts

@@ -5,6 +5,7 @@
  * This is a self-contained module with implementation, types, and JSDoc.
  */
 
+import { aiEventClient } from '../../event-client.js'
 import type { TranscriptionAdapter } from './adapter'
 import type { TranscriptionResult } from '../../types'
 
@@ -61,6 +62,10 @@ export interface TranscriptionActivityOptions<
 /** Result type for the transcription activity */
 export type TranscriptionActivityResult = Promise<TranscriptionResult>
 
+function createId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
+}
+
 // ===========================
 // Activity Implementation
 // ===========================
@@ -104,8 +109,35 @@ export async function generateTranscription<
 ): TranscriptionActivityResult {
   const { adapter, ...rest } = options
   const model = adapter.model
+  const requestId = createId('transcription')
+  const startTime = Date.now()
+
+  aiEventClient.emit('transcription:request:started', {
+    requestId,
+    provider: adapter.name,
+    model,
+    language: rest.language,
+    prompt: rest.prompt,
+    responseFormat: rest.responseFormat,
+    modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+    timestamp: startTime,
+  })
+
+  const result = await adapter.transcribe({ ...rest, model })
+  const duration = Date.now() - startTime
+
+  aiEventClient.emit('transcription:request:completed', {
+    requestId,
+    provider: adapter.name,
+    model,
+    text: result.text,
+    language: result.language,
+    duration,
+    modelOptions: rest.modelOptions as Record<string, unknown> | undefined,
+    timestamp: Date.now(),
+  })
 
-  return adapter.transcribe({ ...rest, model })
+  return result
 }
 
 // ===========================

package/src/activities/generateVideo/index.ts

@@ -7,6 +7,7 @@
  * @experimental Video generation is an experimental feature and may change.
  */
 
+import { aiEventClient } from '../../event-client.js'
 import type { VideoAdapter } from './adapter'
 import type {
   VideoJobResult,
@@ -35,6 +36,10 @@ export type VideoProviderOptions<TAdapter> =
 
 // ===========================
 // Activity Options Types
+
+function createId(prefix: string): string {
+  return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
+}
 // ===========================
 
 /**
@@ -207,6 +212,17 @@ export async function getVideoJobStatus<
   error?: string
 }> {
   const { adapter, jobId } = options
+  const requestId = createId('video-status')
+  const startTime = Date.now()
+
+  aiEventClient.emit('video:request:started', {
+    requestId,
+    provider: adapter.name,
+    model: adapter.model,
+    requestType: 'status',
+    jobId,
+    timestamp: startTime,
+  })
 
   // Get status first
   const statusResult = await adapter.getVideoStatus(jobId)
@@ -215,12 +231,37 @@ export async function getVideoJobStatus<
   if (statusResult.status === 'completed') {
     try {
       const urlResult = await adapter.getVideoUrl(jobId)
+      aiEventClient.emit('video:request:completed', {
+        requestId,
+        provider: adapter.name,
+        model: adapter.model,
+        requestType: 'status',
+        jobId,
+        status: statusResult.status,
+        progress: statusResult.progress,
+        url: urlResult.url,
+        duration: Date.now() - startTime,
+        timestamp: Date.now(),
+      })
       return {
         status: statusResult.status,
         progress: statusResult.progress,
         url: urlResult.url,
       }
     } catch (error) {
+      aiEventClient.emit('video:request:completed', {
+        requestId,
+        provider: adapter.name,
+        model: adapter.model,
+        requestType: 'status',
+        jobId,
+        status: statusResult.status,
+        progress: statusResult.progress,
+        error:
+          error instanceof Error ? error.message : 'Failed to get video URL',
+        duration: Date.now() - startTime,
+        timestamp: Date.now(),
+      })
       // If URL fetch fails, still return status
       return {
         status: statusResult.status,
@@ -231,6 +272,19 @@ export async function getVideoJobStatus<
     }
   }
 
+  aiEventClient.emit('video:request:completed', {
+    requestId,
+    provider: adapter.name,
+    model: adapter.model,
+    requestType: 'status',
+    jobId,
+    status: statusResult.status,
+    progress: statusResult.progress,
+    error: statusResult.error,
+    duration: Date.now() - startTime,
+    timestamp: Date.now(),
+  })
+
   // Return status for non-completed jobs
   return {
     status: statusResult.status,

package/src/activities/summarize/index.ts

@@ -180,8 +180,9 @@ async function runSummarize(
   const inputLength = text.length
   const startTime = Date.now()
 
-  aiEventClient.emit('summarize:started', {
+  aiEventClient.emit('summarize:request:started', {
     requestId,
+    provider: adapter.name,
     model,
     inputLength,
     timestamp: startTime,
@@ -200,8 +201,9 @@ async function runSummarize(
   const duration = Date.now() - startTime
   const outputLength = result.summary.length
 
-  aiEventClient.emit('summarize:completed', {
+  aiEventClient.emit('summarize:request:completed', {
     requestId,
+    provider: adapter.name,
     model,
     inputLength,
     outputLength,
@@ -240,21 +242,20 @@ async function* runStreamingSummarize(
   // Fall back to non-streaming and yield as a single chunk
   const result = await adapter.summarize(summarizeOptions)
 
-  // Yield content chunk with the summary
+  // Yield TEXT_MESSAGE_CONTENT event with the summary
   yield {
-    type: 'content',
-    id: result.id,
+    type: 'TEXT_MESSAGE_CONTENT',
+    messageId: result.id,
     model: result.model,
     timestamp: Date.now(),
     delta: result.summary,
     content: result.summary,
-    role: 'assistant',
   }
 
-  // Yield done chunk
+  // Yield RUN_FINISHED event
   yield {
-    type: 'done',
-    id: result.id,
+    type: 'RUN_FINISHED',
+    runId: result.id,
     model: result.model,
     timestamp: Date.now(),
     finishReason: 'stop',