@tanstack/ai 0.2.1 → 0.3.0

package/src/activities/chat/index.ts

@@ -23,14 +23,18 @@ import type { AnyTextAdapter } from './adapter'
 import type {
   AgentLoopStrategy,
   ConstrainedModelMessage,
-  DoneStreamChunk,
   InferSchemaType,
   ModelMessage,
+  RunFinishedEvent,
   SchemaInput,
   StreamChunk,
+  TextMessageContentEvent,
   TextOptions,
   Tool,
   ToolCall,
+  ToolCallArgsEvent,
+  ToolCallEndEvent,
+  ToolCallStartEvent,
 } from '../../types'
 
 // ===========================
@@ -213,7 +217,9 @@ class TextEngine<
   private totalChunkCount = 0
   private currentMessageId: string | null = null
   private accumulatedContent = ''
-  private doneChunk: DoneStreamChunk | null = null
+  private eventOptions?: Record<string, unknown>
+  private eventToolNames?: Array<string>
+  private finishedEvent: RunFinishedEvent | null = null
   private shouldEmitStreamEnd = true
   private earlyTermination = false
   private toolPhase: ToolPhaseResult = 'continue'
@@ -278,16 +284,7 @@ class TextEngine<
 
   private beforeRun(): void {
     this.streamStartTime = Date.now()
-    const {
-      model,
-      tools,
-      temperature,
-      topP,
-      maxTokens,
-      metadata,
-      modelOptions,
-      conversationId,
-    } = this.params
+    const { tools, temperature, topP, maxTokens, metadata } = this.params
 
     // Gather flattened options into an object for event emission
     const options: Record<string, unknown> = {}
@@ -296,26 +293,49 @@ class TextEngine<
     if (maxTokens !== undefined) options.maxTokens = maxTokens
     if (metadata !== undefined) options.metadata = metadata
 
-    aiEventClient.emit('text:started', {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      model: model,
-      provider: this.adapter.name,
-      messageCount: this.initialMessageCount,
-      hasTools: !!tools && tools.length > 0,
-      streaming: true,
+    this.eventOptions = Object.keys(options).length > 0 ? options : undefined
+    this.eventToolNames = tools?.map((t) => t.name)
+
+    aiEventClient.emit('text:request:started', {
+      ...this.buildTextEventContext(),
       timestamp: Date.now(),
-      clientId: conversationId,
-      toolNames: tools?.map((t) => t.name),
-      options: Object.keys(options).length > 0 ? options : undefined,
-      modelOptions: modelOptions as Record<string, unknown> | undefined,
     })
 
-    aiEventClient.emit('stream:started', {
-      streamId: this.streamId,
-      model,
-      provider: this.adapter.name,
-      timestamp: Date.now(),
+    // Always emit messages for tracking:
+    // - For existing conversations (with conversationId): only emit the latest user message
+    // - For new conversations (without conversationId): emit all messages for reconstruction
+    const messagesToEmit = this.params.conversationId
+      ? this.messages.slice(-1).filter((m) => m.role === 'user')
+      : this.messages
+
+    messagesToEmit.forEach((message, index) => {
+      const messageIndex = this.params.conversationId
+        ? this.messages.length - 1
+        : index
+      const messageId = this.createId('msg')
+      const baseContext = this.buildTextEventContext()
+      const content = this.getContentString(message.content)
+
+      aiEventClient.emit('text:message:created', {
+        ...baseContext,
+        messageId,
+        role: message.role,
+        content,
+        toolCalls: message.toolCalls,
+        messageIndex,
+        timestamp: Date.now(),
+      })
+
+      if (message.role === 'user') {
+        aiEventClient.emit('text:message:user', {
+          ...baseContext,
+          messageId,
+          role: 'user',
+          content,
+          messageIndex,
+          timestamp: Date.now(),
+        })
+      }
     })
   }
 
@@ -325,23 +345,13 @@ class TextEngine<
     }
 
     const now = Date.now()
-
-    // Emit text:completed with final state
-    aiEventClient.emit('text:completed', {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      model: this.params.model,
+    // Emit text:request:completed with final state
+    aiEventClient.emit('text:request:completed', {
+      ...this.buildTextEventContext(),
       content: this.accumulatedContent,
       messageId: this.currentMessageId || undefined,
       finishReason: this.lastFinishReason || undefined,
-      usage: this.doneChunk?.usage,
-      timestamp: now,
-    })
-
-    aiEventClient.emit('stream:ended', {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      totalChunks: this.totalChunkCount,
+      usage: this.finishedEvent?.usage,
       duration: now - this.streamStartTime,
       timestamp: now,
     })
@@ -366,7 +376,16 @@ class TextEngine<
   private beginIteration(): void {
     this.currentMessageId = this.createId('msg')
     this.accumulatedContent = ''
-    this.doneChunk = null
+    this.finishedEvent = null
+
+    const baseContext = this.buildTextEventContext()
+    aiEventClient.emit('text:message:created', {
+      ...baseContext,
+      messageId: this.currentMessageId,
+      role: 'assistant',
+      content: '',
+      timestamp: Date.now(),
+    })
   }
 
   private async *streamModelResponse(): AsyncGenerator<StreamChunk> {
@@ -413,101 +432,96 @@ class TextEngine<
 
   private handleStreamChunk(chunk: StreamChunk): void {
     switch (chunk.type) {
-      case 'content':
-        this.handleContentChunk(chunk)
+      // AG-UI Events
+      case 'TEXT_MESSAGE_CONTENT':
+        this.handleTextMessageContentEvent(chunk)
+        break
+      case 'TOOL_CALL_START':
+        this.handleToolCallStartEvent(chunk)
         break
-      case 'tool_call':
-        this.handleToolCallChunk(chunk)
+      case 'TOOL_CALL_ARGS':
+        this.handleToolCallArgsEvent(chunk)
         break
-      case 'tool_result':
-        this.handleToolResultChunk(chunk)
+      case 'TOOL_CALL_END':
+        this.handleToolCallEndEvent(chunk)
         break
-      case 'done':
-        this.handleDoneChunk(chunk)
+      case 'RUN_FINISHED':
+        this.handleRunFinishedEvent(chunk)
         break
-      case 'error':
-        this.handleErrorChunk(chunk)
+      case 'RUN_ERROR':
+        this.handleRunErrorEvent(chunk)
         break
-      case 'thinking':
-        this.handleThinkingChunk(chunk)
+      case 'STEP_FINISHED':
+        this.handleStepFinishedEvent(chunk)
         break
+
       default:
+        // RUN_STARTED, TEXT_MESSAGE_START, TEXT_MESSAGE_END, STEP_STARTED,
+        // STATE_SNAPSHOT, STATE_DELTA, CUSTOM
+        // - no special handling needed in chat activity
         break
     }
   }
 
-  private handleContentChunk(chunk: Extract<StreamChunk, { type: 'content' }>) {
-    this.accumulatedContent = chunk.content
-    aiEventClient.emit('stream:chunk:content', {
-      streamId: this.streamId,
+  // ===========================
+  // AG-UI Event Handlers
+  // ===========================
+
+  private handleTextMessageContentEvent(chunk: TextMessageContentEvent): void {
+    if (chunk.content) {
+      this.accumulatedContent = chunk.content
+    } else {
+      this.accumulatedContent += chunk.delta
+    }
+    aiEventClient.emit('text:chunk:content', {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || undefined,
-      content: chunk.content,
+      content: this.accumulatedContent,
       delta: chunk.delta,
       timestamp: Date.now(),
     })
   }
 
-  private handleToolCallChunk(
-    chunk: Extract<StreamChunk, { type: 'tool_call' }>,
-  ): void {
-    this.toolCallManager.addToolCallChunk(chunk)
-    aiEventClient.emit('stream:chunk:tool-call', {
-      streamId: this.streamId,
+  private handleToolCallStartEvent(chunk: ToolCallStartEvent): void {
+    this.toolCallManager.addToolCallStartEvent(chunk)
+    aiEventClient.emit('text:chunk:tool-call', {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || undefined,
-      toolCallId: chunk.toolCall.id,
-      toolName: chunk.toolCall.function.name,
-      index: chunk.index,
-      arguments: chunk.toolCall.function.arguments,
+      toolCallId: chunk.toolCallId,
+      toolName: chunk.toolName,
+      index: chunk.index ?? 0,
+      arguments: '',
       timestamp: Date.now(),
     })
   }
 
-  private handleToolResultChunk(
-    chunk: Extract<StreamChunk, { type: 'tool_result' }>,
-  ): void {
-    aiEventClient.emit('stream:chunk:tool-result', {
-      streamId: this.streamId,
+  private handleToolCallArgsEvent(chunk: ToolCallArgsEvent): void {
+    this.toolCallManager.addToolCallArgsEvent(chunk)
+    aiEventClient.emit('text:chunk:tool-call', {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || undefined,
       toolCallId: chunk.toolCallId,
-      result: chunk.content,
+      toolName: '',
+      index: 0,
+      arguments: chunk.delta,
       timestamp: Date.now(),
     })
   }
 
-  private handleDoneChunk(chunk: DoneStreamChunk): void {
-    // Don't overwrite a tool_calls finishReason with a stop finishReason
-    // This can happen when adapters send multiple done chunks
-    if (
-      this.doneChunk?.finishReason === 'tool_calls' &&
-      chunk.finishReason === 'stop'
-    ) {
-      // Still emit the event and update lastFinishReason, but don't overwrite doneChunk
-      this.lastFinishReason = chunk.finishReason
-      aiEventClient.emit('stream:chunk:done', {
-        streamId: this.streamId,
-        messageId: this.currentMessageId || undefined,
-        finishReason: chunk.finishReason,
-        usage: chunk.usage,
-        timestamp: Date.now(),
-      })
-
-      if (chunk.usage) {
-        aiEventClient.emit('usage:tokens', {
-          requestId: this.requestId,
-          streamId: this.streamId,
-          messageId: this.currentMessageId || undefined,
-          model: this.params.model,
-          usage: chunk.usage,
-          timestamp: Date.now(),
-        })
-      }
-      return
-    }
+  private handleToolCallEndEvent(chunk: ToolCallEndEvent): void {
+    this.toolCallManager.completeToolCall(chunk)
+    aiEventClient.emit('text:chunk:tool-result', {
+      ...this.buildTextEventContext(),
+      messageId: this.currentMessageId || undefined,
+      toolCallId: chunk.toolCallId,
+      result: chunk.result || '',
+      timestamp: Date.now(),
+    })
+  }
 
-    this.doneChunk = chunk
-    this.lastFinishReason = chunk.finishReason
-    aiEventClient.emit('stream:chunk:done', {
-      streamId: this.streamId,
+  private handleRunFinishedEvent(chunk: RunFinishedEvent): void {
+    aiEventClient.emit('text:chunk:done', {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || undefined,
       finishReason: chunk.finishReason,
       usage: chunk.usage,
@@ -515,22 +529,32 @@ class TextEngine<
     })
 
     if (chunk.usage) {
-      aiEventClient.emit('usage:tokens', {
-        requestId: this.requestId,
-        streamId: this.streamId,
+      aiEventClient.emit('text:usage', {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || undefined,
-        model: this.params.model,
         usage: chunk.usage,
         timestamp: Date.now(),
       })
     }
+
+    // Don't overwrite a tool_calls finishReason with a stop finishReason
+    if (
+      this.finishedEvent?.finishReason === 'tool_calls' &&
+      chunk.finishReason === 'stop'
+    ) {
+      this.lastFinishReason = chunk.finishReason
+      return
+    }
+
+    this.finishedEvent = chunk
+    this.lastFinishReason = chunk.finishReason
   }
 
-  private handleErrorChunk(
-    chunk: Extract<StreamChunk, { type: 'error' }>,
+  private handleRunErrorEvent(
+    chunk: Extract<StreamChunk, { type: 'RUN_ERROR' }>,
   ): void {
-    aiEventClient.emit('stream:chunk:error', {
-      streamId: this.streamId,
+    aiEventClient.emit('text:chunk:error', {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || undefined,
       error: chunk.error.message,
       timestamp: Date.now(),
@@ -539,16 +563,19 @@ class TextEngine<
     this.shouldEmitStreamEnd = false
   }
 
-  private handleThinkingChunk(
-    chunk: Extract<StreamChunk, { type: 'thinking' }>,
+  private handleStepFinishedEvent(
+    chunk: Extract<StreamChunk, { type: 'STEP_FINISHED' }>,
   ): void {
-    aiEventClient.emit('stream:chunk:thinking', {
-      streamId: this.streamId,
-      messageId: this.currentMessageId || undefined,
-      content: chunk.content,
-      delta: chunk.delta,
-      timestamp: Date.now(),
-    })
+    // Handle thinking/reasoning content from STEP_FINISHED events
+    if (chunk.content || chunk.delta) {
+      aiEventClient.emit('text:chunk:thinking', {
+        ...this.buildTextEventContext(),
+        messageId: this.currentMessageId || undefined,
+        content: chunk.content || '',
+        delta: chunk.delta,
+        timestamp: Date.now(),
+      })
+    }
   }
 
   private async *checkForPendingToolCalls(): AsyncGenerator<
@@ -561,16 +588,7 @@ class TextEngine<
       return 'continue'
     }
 
-    const doneChunk = this.createSyntheticDoneChunk()
-
-    aiEventClient.emit('text:iteration', {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      iterationNumber: this.iterationCount + 1,
-      messageCount: this.messages.length,
-      toolCallCount: pendingToolCalls.length,
-      timestamp: Date.now(),
-    })
+    const finishEvent = this.createSyntheticFinishedEvent()
 
     const { approvals, clientToolResults } = this.collectClientState()
 
@@ -587,14 +605,14 @@ class TextEngine<
     ) {
       for (const chunk of this.emitApprovalRequests(
         executionResult.needsApproval,
-        doneChunk,
+        finishEvent,
       )) {
         yield chunk
       }
 
       for (const chunk of this.emitClientToolInputs(
         executionResult.needsClientExecution,
-        doneChunk,
+        finishEvent,
       )) {
         yield chunk
       }
@@ -605,7 +623,7 @@ class TextEngine<
 
     const toolResultChunks = this.emitToolResults(
       executionResult.results,
-      doneChunk,
+      finishEvent,
     )
 
     for (const chunk of toolResultChunks) {
@@ -622,22 +640,13 @@ class TextEngine<
     }
 
     const toolCalls = this.toolCallManager.getToolCalls()
-    const doneChunk = this.doneChunk
+    const finishEvent = this.finishedEvent
 
-    if (!doneChunk || toolCalls.length === 0) {
+    if (!finishEvent || toolCalls.length === 0) {
       this.setToolPhase('stop')
       return
     }
 
-    aiEventClient.emit('text:iteration', {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      iterationNumber: this.iterationCount + 1,
-      messageCount: this.messages.length,
-      toolCallCount: toolCalls.length,
-      timestamp: Date.now(),
-    })
-
     this.addAssistantToolCallMessage(toolCalls)
 
     const { approvals, clientToolResults } = this.collectClientState()
@@ -655,14 +664,14 @@ class TextEngine<
     ) {
       for (const chunk of this.emitApprovalRequests(
         executionResult.needsApproval,
-        doneChunk,
+        finishEvent,
       )) {
         yield chunk
       }
 
       for (const chunk of this.emitClientToolInputs(
         executionResult.needsClientExecution,
-        doneChunk,
+        finishEvent,
       )) {
         yield chunk
       }
@@ -673,7 +682,7 @@ class TextEngine<
 
     const toolResultChunks = this.emitToolResults(
       executionResult.results,
-      doneChunk,
+      finishEvent,
     )
 
     for (const chunk of toolResultChunks) {
@@ -687,13 +696,14 @@ class TextEngine<
 
   private shouldExecuteToolPhase(): boolean {
     return (
-      this.doneChunk?.finishReason === 'tool_calls' &&
+      this.finishedEvent?.finishReason === 'tool_calls' &&
       this.tools.length > 0 &&
       this.toolCallManager.hasToolCalls()
     )
   }
 
   private addAssistantToolCallMessage(toolCalls: Array<ToolCall>): void {
+    const messageId = this.currentMessageId ?? this.createId('msg')
     this.messages = [
       ...this.messages,
       {
@@ -702,6 +712,15 @@ class TextEngine<
         toolCalls,
       },
     ]
+
+    aiEventClient.emit('text:message:created', {
+      ...this.buildTextEventContext(),
+      messageId,
+      role: 'assistant',
+      content: this.accumulatedContent || '',
+      toolCalls,
+      timestamp: Date.now(),
+    })
   }
 
   private collectClientState(): {
@@ -740,13 +759,13 @@ class TextEngine<
 
   private emitApprovalRequests(
     approvals: Array<ApprovalRequest>,
-    doneChunk: DoneStreamChunk,
+    finishEvent: RunFinishedEvent,
   ): Array<StreamChunk> {
     const chunks: Array<StreamChunk> = []
 
     for (const approval of approvals) {
-      aiEventClient.emit('stream:approval-requested', {
-        streamId: this.streamId,
+      aiEventClient.emit('tools:approval:requested', {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || undefined,
         toolCallId: approval.toolCallId,
         toolName: approval.toolName,
@@ -755,17 +774,20 @@ class TextEngine<
         timestamp: Date.now(),
       })
 
+      // Emit a CUSTOM event for approval requests
       chunks.push({
-        type: 'approval-requested',
-        id: doneChunk.id,
-        model: doneChunk.model,
+        type: 'CUSTOM',
         timestamp: Date.now(),
-        toolCallId: approval.toolCallId,
-        toolName: approval.toolName,
-        input: approval.input,
-        approval: {
-          id: approval.approvalId,
-          needsApproval: true,
+        model: finishEvent.model,
+        name: 'approval-requested',
+        data: {
+          toolCallId: approval.toolCallId,
+          toolName: approval.toolName,
+          input: approval.input,
+          approval: {
+            id: approval.approvalId,
+            needsApproval: true,
+          },
         },
       })
     }
@@ -775,13 +797,13 @@ class TextEngine<
 
   private emitClientToolInputs(
     clientRequests: Array<ClientToolRequest>,
-    doneChunk: DoneStreamChunk,
+    finishEvent: RunFinishedEvent,
   ): Array<StreamChunk> {
     const chunks: Array<StreamChunk> = []
 
     for (const clientTool of clientRequests) {
-      aiEventClient.emit('stream:tool-input-available', {
-        streamId: this.streamId,
+      aiEventClient.emit('tools:input:available', {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || undefined,
         toolCallId: clientTool.toolCallId,
         toolName: clientTool.toolName,
@@ -789,14 +811,17 @@ class TextEngine<
         timestamp: Date.now(),
       })
 
+      // Emit a CUSTOM event for client tool inputs
       chunks.push({
-        type: 'tool-input-available',
-        id: doneChunk.id,
-        model: doneChunk.model,
+        type: 'CUSTOM',
         timestamp: Date.now(),
-        toolCallId: clientTool.toolCallId,
-        toolName: clientTool.toolName,
-        input: clientTool.input,
+        model: finishEvent.model,
+        name: 'tool-input-available',
+        data: {
+          toolCallId: clientTool.toolCallId,
+          toolName: clientTool.toolName,
+          input: clientTool.input,
+        },
       })
     }
 
@@ -805,14 +830,13 @@ class TextEngine<
 
   private emitToolResults(
     results: Array<ToolResult>,
-    doneChunk: DoneStreamChunk,
+    finishEvent: RunFinishedEvent,
   ): Array<StreamChunk> {
     const chunks: Array<StreamChunk> = []
 
     for (const result of results) {
-      aiEventClient.emit('tool:call-completed', {
-        requestId: this.requestId,
-        streamId: this.streamId,
+      aiEventClient.emit('tools:call:completed', {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || undefined,
         toolCallId: result.toolCallId,
         toolName: result.toolName,
@@ -822,16 +846,16 @@ class TextEngine<
       })
 
       const content = JSON.stringify(result.result)
-      const chunk: Extract<StreamChunk, { type: 'tool_result' }> = {
-        type: 'tool_result',
-        id: doneChunk.id,
-        model: doneChunk.model,
+
+      // Emit TOOL_CALL_END event
+      chunks.push({
+        type: 'TOOL_CALL_END',
         timestamp: Date.now(),
+        model: finishEvent.model,
         toolCallId: result.toolCallId,
-        content,
-      }
-
-      chunks.push(chunk)
+        toolName: result.toolName,
+        result: content,
+      })
 
       this.messages = [
         ...this.messages,
@@ -841,6 +865,14 @@ class TextEngine<
           toolCallId: result.toolCallId,
         },
       ]
+
+      aiEventClient.emit('text:message:created', {
+        ...this.buildTextEventContext(),
+        messageId: this.createId('msg'),
+        role: 'tool',
+        content,
+        timestamp: Date.now(),
+      })
     }
 
     return chunks
@@ -868,10 +900,10 @@ class TextEngine<
     return pending
   }
 
-  private createSyntheticDoneChunk(): DoneStreamChunk {
+  private createSyntheticFinishedEvent(): RunFinishedEvent {
     return {
-      type: 'done',
-      id: this.createId('pending'),
+      type: 'RUN_FINISHED',
+      runId: this.createId('pending'),
       model: this.params.model,
       timestamp: Date.now(),
       finishReason: 'tool_calls',
@@ -896,6 +928,50 @@ class TextEngine<
     return !!this.effectiveSignal?.aborted
   }
 
+  private buildTextEventContext(): {
+    requestId: string
+    streamId: string
+    provider: string
+    model: string
+    clientId?: string
+    source?: 'client' | 'server'
+    systemPrompts?: Array<string>
+    toolNames?: Array<string>
+    options?: Record<string, unknown>
+    modelOptions?: Record<string, unknown>
+    messageCount: number
+    hasTools: boolean
+    streaming: boolean
+  } {
+    return {
+      requestId: this.requestId,
+      streamId: this.streamId,
+      provider: this.adapter.name,
+      model: this.params.model,
+      clientId: this.params.conversationId,
+      source: 'server',
+      systemPrompts:
+        this.systemPrompts.length > 0 ? this.systemPrompts : undefined,
+      toolNames: this.eventToolNames,
+      options: this.eventOptions,
+      modelOptions: this.params.modelOptions as
+        | Record<string, unknown>
+        | undefined,
+      messageCount: this.initialMessageCount,
+      hasTools: this.tools.length > 0,
+      streaming: true,
+    }
+  }
+
+  private getContentString(content: ModelMessage['content']): string {
+    if (typeof content === 'string') return content
+    const text =
+      content
+        ?.map((part) => (part.type === 'text' ? part.content : ''))
+        .join('') || ''
+    return text
+  }
+
   private setToolPhase(phase: ToolPhaseResult): void {
     this.toolPhase = phase
     if (phase === 'wait') {