@tanstack/ai 0.5.1 → 0.6.1

package/src/activities/chat/index.ts

@@ -24,6 +24,7 @@ import type { AnyTextAdapter } from './adapter'
 import type {
   AgentLoopStrategy,
   ConstrainedModelMessage,
+  CustomEvent,
   InferSchemaType,
   ModelMessage,
   RunFinishedEvent,
@@ -601,17 +602,30 @@ class TextEngine<
 
     const { approvals, clientToolResults } = this.collectClientState()
 
-    const executionResult = await executeToolCalls(
+    const generator = executeToolCalls(
       pendingToolCalls,
       this.tools,
       approvals,
       clientToolResults,
+      (eventName, data) => this.createCustomEventChunk(eventName, data),
     )
 
+    // Consume the async generator, yielding custom events and collecting the return value
+    const executionResult = yield* this.drainToolCallGenerator(generator)
+
     if (
       executionResult.needsApproval.length > 0 ||
       executionResult.needsClientExecution.length > 0
     ) {
+      if (executionResult.results.length > 0) {
+        for (const chunk of this.emitToolResults(
+          executionResult.results,
+          finishEvent,
+        )) {
+          yield chunk
+        }
+      }
+
       for (const chunk of this.emitApprovalRequests(
         executionResult.needsApproval,
         finishEvent,
@@ -660,17 +674,30 @@ class TextEngine<
 
     const { approvals, clientToolResults } = this.collectClientState()
 
-    const executionResult = await executeToolCalls(
+    const generator = executeToolCalls(
       toolCalls,
       this.tools,
       approvals,
       clientToolResults,
+      (eventName, data) => this.createCustomEventChunk(eventName, data),
     )
 
+    // Consume the async generator, yielding custom events and collecting the return value
+    const executionResult = yield* this.drainToolCallGenerator(generator)
+
     if (
       executionResult.needsApproval.length > 0 ||
       executionResult.needsClientExecution.length > 0
     ) {
+      if (executionResult.results.length > 0) {
+        for (const chunk of this.emitToolResults(
+          executionResult.results,
+          finishEvent,
+        )) {
+          yield chunk
+        }
+      }
+
       for (const chunk of this.emitApprovalRequests(
         executionResult.needsApproval,
         finishEvent,
@@ -833,7 +860,7 @@ class TextEngine<
         timestamp: Date.now(),
         model: finishEvent.model,
         name: 'approval-requested',
-        data: {
+        value: {
           toolCallId: approval.toolCallId,
           toolName: approval.toolName,
           input: approval.input,
@@ -870,7 +897,7 @@ class TextEngine<
         timestamp: Date.now(),
         model: finishEvent.model,
         name: 'tool-input-available',
-        data: {
+        value: {
           toolCallId: clientTool.toolCallId,
           toolName: clientTool.toolName,
           input: clientTool.input,
@@ -1052,6 +1079,50 @@ class TextEngine<
     }
   }
 
+  /**
+   * Drain an executeToolCalls async generator, yielding any CustomEvent chunks
+   * and returning the final ExecuteToolCallsResult.
+   */
+  private async *drainToolCallGenerator(
+    generator: AsyncGenerator<
+      CustomEvent,
+      {
+        results: Array<ToolResult>
+        needsApproval: Array<ApprovalRequest>
+        needsClientExecution: Array<ClientToolRequest>
+      },
+      void
+    >,
+  ): AsyncGenerator<
+    StreamChunk,
+    {
+      results: Array<ToolResult>
+      needsApproval: Array<ApprovalRequest>
+      needsClientExecution: Array<ClientToolRequest>
+    },
+    void
+  > {
+    let next = await generator.next()
+    while (!next.done) {
+      yield next.value
+      next = await generator.next()
+    }
+    return next.value
+  }
+
+  private createCustomEventChunk(
+    eventName: string,
+    value: Record<string, any>,
+  ): CustomEvent {
+    return {
+      type: 'CUSTOM',
+      timestamp: Date.now(),
+      model: this.params.model,
+      name: eventName,
+      value,
+    }
+  }
+
   private createId(prefix: string): string {
     return `${prefix}-${Date.now()}-${Math.random().toString(36).slice(2, 9)}`
   }

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

@@ -73,6 +73,13 @@ export interface StreamProcessorEvents {
     approvalId: string
   }) => void
 
+  // Custom events from server-side tools
+  onCustomEvent?: (
+    eventType: string,
+    data: unknown,
+    context: { toolCallId?: string },
+  ) => void
+
   // Granular events for UI optimization (character-by-character, state tracking)
   onTextUpdate?: (messageId: string, content: string) => void
   onToolCallStateChange?: (
@@ -807,8 +814,8 @@ export class StreamProcessor {
     chunk: Extract<StreamChunk, { type: 'CUSTOM' }>,
   ): void {
     // Handle client tool input availability - trigger client-side execution
-    if (chunk.name === 'tool-input-available' && chunk.data) {
-      const { toolCallId, toolName, input } = chunk.data as {
+    if (chunk.name === 'tool-input-available' && chunk.value) {
+      const { toolCallId, toolName, input } = chunk.value as {
         toolCallId: string
         toolName: string
         input: any
@@ -820,11 +827,12 @@ export class StreamProcessor {
         toolName,
         input,
       })
+      return
     }
 
     // Handle approval requests
-    if (chunk.name === 'approval-requested' && chunk.data) {
-      const { toolCallId, toolName, input, approval } = chunk.data as {
+    if (chunk.name === 'approval-requested' && chunk.value) {
+      const { toolCallId, toolName, input, approval } = chunk.value as {
         toolCallId: string
         toolName: string
         input: any
@@ -849,6 +857,16 @@ export class StreamProcessor {
         input,
         approvalId: approval.id,
       })
+      return
+    }
+
+    // Forward non-system custom events to onCustomEvent callback
+    if (this.events.onCustomEvent) {
+      const toolCallId =
+        chunk.value && typeof chunk.value === 'object'
+          ? (chunk.value as any).toolCallId
+          : undefined
+      this.events.onCustomEvent(chunk.name, chunk.value, { toolCallId })
     }
   }
 

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

@@ -1,5 +1,6 @@
 import { isStandardSchema, parseWithStandardSchema } from './schema-converter'
 import type {
+  CustomEvent,
   ModelMessage,
   RunFinishedEvent,
   Tool,
@@ -7,6 +8,7 @@ import type {
   ToolCallArgsEvent,
   ToolCallEndEvent,
   ToolCallStartEvent,
+  ToolExecutionContext,
 } from '../../../types'
 
 /**
@@ -250,7 +252,46 @@ interface ExecuteToolCallsResult {
 }
 
 /**
- * Execute tool calls based on their configuration
+ * Helper that runs a tool execution promise while polling for pending custom events.
+ * Yields any custom events that are emitted during execution, then returns the
+ * execution result.
+ */
+async function* executeWithEventPolling<T>(
+  executionPromise: Promise<T>,
+  pendingEvents: Array<CustomEvent>,
+): AsyncGenerator<CustomEvent, T, void> {
+  // Use an object to track mutable state across the async boundary
+  const state = { done: false, result: undefined as T }
+  const executionWithFlag = executionPromise.then((r) => {
+    state.done = true
+    state.result = r
+    return r
+  })
+
+  while (!state.done) {
+    // Wait for either the execution to complete or a short timeout
+    await Promise.race([
+      executionWithFlag,
+      new Promise((resolve) => setTimeout(resolve, 10)),
+    ])
+
+    // Flush any pending events
+    while (pendingEvents.length > 0) {
+      yield pendingEvents.shift()!
+    }
+  }
+
+  // Final flush in case events were emitted right at completion
+  while (pendingEvents.length > 0) {
+    yield pendingEvents.shift()!
+  }
+
+  return state.result
+}
+
+/**
+ * Execute tool calls based on their configuration.
+ * Yields CustomEvent chunks during tool execution for real-time progress updates.
  *
  * Handles three cases:
  * 1. Client tools (no execute) - request client to execute
@@ -261,13 +302,18 @@ interface ExecuteToolCallsResult {
  * @param tools - Available tools with their configurations
  * @param approvals - Map of approval decisions (approval.id -> approved boolean)
  * @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(
   toolCalls: Array<ToolCall>,
   tools: ReadonlyArray<Tool>,
   approvals: Map<string, boolean> = new Map(),
   clientResults: Map<string, any> = new Map(),
-): Promise<ExecuteToolCallsResult> {
+  createCustomEventChunk?: (
+    eventName: string,
+    value: Record<string, any>,
+  ) => CustomEvent,
+): AsyncGenerator<CustomEvent, ExecuteToolCallsResult, void> {
   const results: Array<ToolResult> = []
   const needsApproval: Array<ApprovalRequest> = []
   const needsClientExecution: Array<ClientToolRequest> = []
@@ -326,6 +372,29 @@ export async function executeToolCalls(
       }
     }
 
+    // Create a ToolExecutionContext for this tool call with event emission
+    const pendingEvents: Array<CustomEvent> = []
+    const context: ToolExecutionContext = {
+      toolCallId: toolCall.id,
+      emitCustomEvent: (eventName: string, value: Record<string, any>) => {
+        if (createCustomEventChunk) {
+          pendingEvents.push(
+            createCustomEventChunk(eventName, {
+              ...value,
+              toolCallId: toolCall.id,
+            }),
+          )
+        }
+      },
+    }
+
+    // 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
@@ -402,8 +471,15 @@ export async function executeToolCalls(
           // Execute after approval
           const startTime = Date.now()
           try {
-            let result = await tool.execute(input)
+            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 (
@@ -426,6 +502,7 @@ export async function executeToolCalls(
             })
           } catch (error: unknown) {
             const duration = Date.now() - startTime
+            yield* flushEvents()
             const message =
               error instanceof Error ? error.message : 'Unknown error'
             results.push({
@@ -460,8 +537,13 @@ export async function executeToolCalls(
     // CASE 3: Normal server tool - execute immediately
     const startTime = Date.now()
     try {
-      let result = await tool.execute(input)
+      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 (
@@ -482,6 +564,7 @@ export async function executeToolCalls(
       })
     } catch (error: unknown) {
       const duration = Date.now() - startTime
+      yield* flushEvents()
       const message = error instanceof Error ? error.message : 'Unknown error'
       results.push({
         toolCallId: toolCall.id,

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

@@ -4,6 +4,7 @@ import type {
   JSONSchema,
   SchemaInput,
   Tool,
+  ToolExecutionContext,
 } from '../../../types'
 
 /**
@@ -112,6 +113,7 @@ export interface ToolDefinition<
   server: (
     execute: (
       args: InferSchemaType<TInput>,
+      context?: ToolExecutionContext,
     ) => Promise<InferSchemaType<TOutput>> | InferSchemaType<TOutput>,
   ) => ServerTool<TInput, TOutput, TName>
 
@@ -193,6 +195,7 @@ export function toolDefinition<
     server(
       execute: (
         args: InferSchemaType<TInput>,
+        context?: ToolExecutionContext,
       ) => Promise<InferSchemaType<TOutput>> | InferSchemaType<TOutput>,
     ): ServerTool<TInput, TOutput, TName> {
       return {

package/src/types.ts

@@ -344,6 +344,34 @@ export type ConstrainedModelMessage<
   content: ConstrainedContent<TInputModalitiesTypes>
 }
 
+/**
+ * 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
+}
+
 /**
  * Tool/Function definition for function calling.
  *
@@ -460,7 +488,7 @@ export interface Tool<
    *   return weather; // Can return object or string
    * }
    */
-  execute?: (args: any) => Promise<any> | any
+  execute?: (args: any, context?: ToolExecutionContext) => Promise<any> | any
 
   /** If true, tool execution requires user approval before running. Works with both server and client tools. */
   needsApproval?: boolean
@@ -894,8 +922,8 @@ export interface CustomEvent extends BaseAGUIEvent {
   type: 'CUSTOM'
   /** Custom event name */
   name: string
-  /** Custom event data */
-  data?: unknown
+  /** Custom event value */
+  value?: unknown
 }
 
 /**