@tanstack/ai 0.2.2 → 0.3.1

package/dist/esm/activities/chat/index.js

@@ -3,6 +3,7 @@ import { streamToText } from "../../stream-to-response.js";
 import { ToolCallManager, executeToolCalls } from "./tools/tool-calls.js";
 import { convertSchemaToJsonSchema, isStandardSchema, parseWithStandardSchema } from "./tools/schema-converter.js";
 import { maxIterations } from "./agent-loop-strategies.js";
+import { convertMessagesToModelMessages } from "./messages.js";
 const kind = "text";
 function createChatOptions(options) {
   return options;
@@ -15,7 +16,7 @@ class TextEngine {
     this.totalChunkCount = 0;
     this.currentMessageId = null;
     this.accumulatedContent = "";
-    this.doneChunk = null;
+    this.finishedEvent = null;
     this.shouldEmitStreamEnd = true;
     this.earlyTermination = false;
     this.toolPhase = "continue";
@@ -27,7 +28,14 @@ class TextEngine {
     this.loopStrategy = config.params.agentLoopStrategy || maxIterations(5);
     this.toolCallManager = new ToolCallManager(this.tools);
     this.initialMessageCount = config.params.messages.length;
-    this.messages = config.params.messages;
+    const { approvals, clientToolResults } = this.extractClientStateFromOriginalMessages(
+      config.params.messages
+    );
+    this.initialApprovals = approvals;
+    this.initialClientToolResults = clientToolResults;
+    this.messages = convertMessagesToModelMessages(
+      config.params.messages
+    );
     this.requestId = this.createId("chat");
     this.streamId = this.createId("stream");
     this.effectiveRequest = config.params.abortController ? { signal: config.params.abortController.signal } : void 0;
@@ -66,61 +74,56 @@ class TextEngine {
   }
   beforeRun() {
     this.streamStartTime = Date.now();
-    const {
-      model,
-      tools,
-      temperature,
-      topP,
-      maxTokens,
-      metadata,
-      modelOptions,
-      conversationId
-    } = this.params;
+    const { tools, temperature, topP, maxTokens, metadata } = this.params;
     const options = {};
     if (temperature !== void 0) options.temperature = temperature;
     if (topP !== void 0) options.topP = topP;
     if (maxTokens !== void 0) options.maxTokens = maxTokens;
     if (metadata !== void 0) options.metadata = metadata;
-    aiEventClient.emit("text:started", {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      model,
-      provider: this.adapter.name,
-      messageCount: this.initialMessageCount,
-      hasTools: !!tools && tools.length > 0,
-      streaming: true,
-      timestamp: Date.now(),
-      clientId: conversationId,
-      toolNames: tools?.map((t) => t.name),
-      options: Object.keys(options).length > 0 ? options : void 0,
-      modelOptions
-    });
-    aiEventClient.emit("stream:started", {
-      streamId: this.streamId,
-      model,
-      provider: this.adapter.name,
+    this.eventOptions = Object.keys(options).length > 0 ? options : void 0;
+    this.eventToolNames = tools?.map((t) => t.name);
+    aiEventClient.emit("text:request:started", {
+      ...this.buildTextEventContext(),
       timestamp: Date.now()
     });
+    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()
+        });
+      }
+    });
   }
   afterRun() {
     if (!this.shouldEmitStreamEnd) {
       return;
     }
     const now = Date.now();
-    aiEventClient.emit("text:completed", {
-      requestId: this.requestId,
-      streamId: this.streamId,
-      model: this.params.model,
+    aiEventClient.emit("text:request:completed", {
+      ...this.buildTextEventContext(),
       content: this.accumulatedContent,
       messageId: this.currentMessageId || void 0,
       finishReason: this.lastFinishReason || void 0,
-      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
     });
@@ -141,7 +144,15 @@ class TextEngine {
   beginIteration() {
     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()
+    });
   }
   async *streamModelResponse() {
     const { temperature, topP, maxTokens, metadata, modelOptions } = this.params;
@@ -176,102 +187,107 @@ class TextEngine {
   }
   handleStreamChunk(chunk) {
     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;
     }
   }
-  handleContentChunk(chunk) {
-    this.accumulatedContent = chunk.content;
-    aiEventClient.emit("stream:chunk:content", {
-      streamId: this.streamId,
+  // ===========================
+  // AG-UI Event Handlers
+  // ===========================
+  handleTextMessageContentEvent(chunk) {
+    if (chunk.content) {
+      this.accumulatedContent = chunk.content;
+    } else {
+      this.accumulatedContent += chunk.delta;
+    }
+    aiEventClient.emit("text:chunk:content", {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || void 0,
-      content: chunk.content,
+      content: this.accumulatedContent,
       delta: chunk.delta,
       timestamp: Date.now()
     });
   }
-  handleToolCallChunk(chunk) {
-    this.toolCallManager.addToolCallChunk(chunk);
-    aiEventClient.emit("stream:chunk:tool-call", {
-      streamId: this.streamId,
+  handleToolCallStartEvent(chunk) {
+    this.toolCallManager.addToolCallStartEvent(chunk);
+    aiEventClient.emit("text:chunk:tool-call", {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || void 0,
-      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()
     });
   }
-  handleToolResultChunk(chunk) {
-    aiEventClient.emit("stream:chunk:tool-result", {
-      streamId: this.streamId,
+  handleToolCallArgsEvent(chunk) {
+    this.toolCallManager.addToolCallArgsEvent(chunk);
+    aiEventClient.emit("text:chunk:tool-call", {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || void 0,
       toolCallId: chunk.toolCallId,
-      result: chunk.content,
+      toolName: "",
+      index: 0,
+      arguments: chunk.delta,
       timestamp: Date.now()
     });
   }
-  handleDoneChunk(chunk) {
-    if (this.doneChunk?.finishReason === "tool_calls" && chunk.finishReason === "stop") {
-      this.lastFinishReason = chunk.finishReason;
-      aiEventClient.emit("stream:chunk:done", {
-        streamId: this.streamId,
-        messageId: this.currentMessageId || void 0,
-        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 || void 0,
-          model: this.params.model,
-          usage: chunk.usage,
-          timestamp: Date.now()
-        });
-      }
-      return;
-    }
-    this.doneChunk = chunk;
-    this.lastFinishReason = chunk.finishReason;
-    aiEventClient.emit("stream:chunk:done", {
-      streamId: this.streamId,
+  handleToolCallEndEvent(chunk) {
+    this.toolCallManager.completeToolCall(chunk);
+    aiEventClient.emit("text:chunk:tool-result", {
+      ...this.buildTextEventContext(),
+      messageId: this.currentMessageId || void 0,
+      toolCallId: chunk.toolCallId,
+      result: chunk.result || "",
+      timestamp: Date.now()
+    });
+  }
+  handleRunFinishedEvent(chunk) {
+    aiEventClient.emit("text:chunk:done", {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || void 0,
       finishReason: chunk.finishReason,
       usage: chunk.usage,
       timestamp: Date.now()
     });
     if (chunk.usage) {
-      aiEventClient.emit("usage:tokens", {
-        requestId: this.requestId,
-        streamId: this.streamId,
+      aiEventClient.emit("text:usage", {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || void 0,
-        model: this.params.model,
         usage: chunk.usage,
         timestamp: Date.now()
       });
     }
+    if (this.finishedEvent?.finishReason === "tool_calls" && chunk.finishReason === "stop") {
+      this.lastFinishReason = chunk.finishReason;
+      return;
+    }
+    this.finishedEvent = chunk;
+    this.lastFinishReason = chunk.finishReason;
   }
-  handleErrorChunk(chunk) {
-    aiEventClient.emit("stream:chunk:error", {
-      streamId: this.streamId,
+  handleRunErrorEvent(chunk) {
+    aiEventClient.emit("text:chunk:error", {
+      ...this.buildTextEventContext(),
       messageId: this.currentMessageId || void 0,
       error: chunk.error.message,
       timestamp: Date.now()
@@ -279,29 +295,23 @@ class TextEngine {
     this.earlyTermination = true;
     this.shouldEmitStreamEnd = false;
   }
-  handleThinkingChunk(chunk) {
-    aiEventClient.emit("stream:chunk:thinking", {
-      streamId: this.streamId,
-      messageId: this.currentMessageId || void 0,
-      content: chunk.content,
-      delta: chunk.delta,
-      timestamp: Date.now()
-    });
+  handleStepFinishedEvent(chunk) {
+    if (chunk.content || chunk.delta) {
+      aiEventClient.emit("text:chunk:thinking", {
+        ...this.buildTextEventContext(),
+        messageId: this.currentMessageId || void 0,
+        content: chunk.content || "",
+        delta: chunk.delta,
+        timestamp: Date.now()
+      });
+    }
   }
   async *checkForPendingToolCalls() {
     const pendingToolCalls = this.getPendingToolCallsFromMessages();
     if (pendingToolCalls.length === 0) {
       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();
     const executionResult = await executeToolCalls(
       pendingToolCalls,
@@ -312,13 +322,13 @@ class TextEngine {
     if (executionResult.needsApproval.length > 0 || executionResult.needsClientExecution.length > 0) {
       for (const chunk of this.emitApprovalRequests(
         executionResult.needsApproval,
-        doneChunk
+        finishEvent
       )) {
         yield chunk;
       }
       for (const chunk of this.emitClientToolInputs(
         executionResult.needsClientExecution,
-        doneChunk
+        finishEvent
       )) {
         yield chunk;
       }
@@ -327,7 +337,7 @@ class TextEngine {
     }
     const toolResultChunks = this.emitToolResults(
       executionResult.results,
-      doneChunk
+      finishEvent
     );
     for (const chunk of toolResultChunks) {
       yield chunk;
@@ -340,19 +350,11 @@ class TextEngine {
       return;
     }
     const toolCalls = this.toolCallManager.getToolCalls();
-    const doneChunk = this.doneChunk;
-    if (!doneChunk || toolCalls.length === 0) {
+    const finishEvent = this.finishedEvent;
+    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();
     const executionResult = await executeToolCalls(
@@ -364,13 +366,13 @@ class TextEngine {
     if (executionResult.needsApproval.length > 0 || executionResult.needsClientExecution.length > 0) {
       for (const chunk of this.emitApprovalRequests(
         executionResult.needsApproval,
-        doneChunk
+        finishEvent
       )) {
         yield chunk;
       }
       for (const chunk of this.emitClientToolInputs(
         executionResult.needsClientExecution,
-        doneChunk
+        finishEvent
       )) {
         yield chunk;
       }
@@ -379,7 +381,7 @@ class TextEngine {
     }
     const toolResultChunks = this.emitToolResults(
       executionResult.results,
-      doneChunk
+      finishEvent
     );
     for (const chunk of toolResultChunks) {
       yield chunk;
@@ -388,9 +390,10 @@ class TextEngine {
     this.setToolPhase("continue");
   }
   shouldExecuteToolPhase() {
-    return this.doneChunk?.finishReason === "tool_calls" && this.tools.length > 0 && this.toolCallManager.hasToolCalls();
+    return this.finishedEvent?.finishReason === "tool_calls" && this.tools.length > 0 && this.toolCallManager.hasToolCalls();
   }
   addAssistantToolCallMessage(toolCalls) {
+    const messageId = this.currentMessageId ?? this.createId("msg");
     this.messages = [
       ...this.messages,
       {
@@ -399,30 +402,60 @@ class TextEngine {
         toolCalls
       }
     ];
+    aiEventClient.emit("text:message:created", {
+      ...this.buildTextEventContext(),
+      messageId,
+      role: "assistant",
+      content: this.accumulatedContent || "",
+      toolCalls,
+      timestamp: Date.now()
+    });
   }
-  collectClientState() {
+  /**
+   * Extract client state (approvals and client tool results) from original messages.
+   * This is called in the constructor BEFORE converting to ModelMessage format,
+   * because the parts array (which contains approval state) is lost during conversion.
+   */
+  extractClientStateFromOriginalMessages(originalMessages) {
     const approvals = /* @__PURE__ */ new Map();
     const clientToolResults = /* @__PURE__ */ new Map();
-    for (const message of this.messages) {
+    for (const message of originalMessages) {
       if (message.role === "assistant" && message.parts) {
-        const parts = message.parts;
-        for (const part of parts) {
-          if (part.type === "tool-call" && part.state === "approval-responded" && part.approval) {
-            approvals.set(part.approval.id, part.approval.approved);
-          }
-          if (part.type === "tool-call" && part.output !== void 0 && !part.approval) {
-            clientToolResults.set(part.id, part.output);
+        for (const part of message.parts) {
+          if (part.type === "tool-call") {
+            if (part.output !== void 0 && !part.approval) {
+              clientToolResults.set(part.id, part.output);
+            }
+            if (part.approval?.id && part.approval?.approved !== void 0 && part.state === "approval-responded") {
+              approvals.set(part.approval.id, part.approval.approved);
+            }
           }
         }
       }
     }
     return { approvals, clientToolResults };
   }
-  emitApprovalRequests(approvals, doneChunk) {
+  collectClientState() {
+    const approvals = new Map(this.initialApprovals);
+    const clientToolResults = new Map(this.initialClientToolResults);
+    for (const message of this.messages) {
+      if (message.role === "tool" && message.toolCallId) {
+        let output;
+        try {
+          output = JSON.parse(message.content);
+        } catch {
+          output = message.content;
+        }
+        clientToolResults.set(message.toolCallId, output);
+      }
+    }
+    return { approvals, clientToolResults };
+  }
+  emitApprovalRequests(approvals, finishEvent) {
     const chunks = [];
     for (const approval of approvals) {
-      aiEventClient.emit("stream:approval-requested", {
-        streamId: this.streamId,
+      aiEventClient.emit("tools:approval:requested", {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || void 0,
         toolCallId: approval.toolCallId,
         toolName: approval.toolName,
@@ -431,26 +464,28 @@ class TextEngine {
         timestamp: Date.now()
       });
       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
+          }
         }
       });
     }
     return chunks;
   }
-  emitClientToolInputs(clientRequests, doneChunk) {
+  emitClientToolInputs(clientRequests, finishEvent) {
     const chunks = [];
     for (const clientTool of clientRequests) {
-      aiEventClient.emit("stream:tool-input-available", {
-        streamId: this.streamId,
+      aiEventClient.emit("tools:input:available", {
+        ...this.buildTextEventContext(),
         messageId: this.currentMessageId || void 0,
         toolCallId: clientTool.toolCallId,
         toolName: clientTool.toolName,
@@ -458,23 +493,24 @@ class TextEngine {
         timestamp: Date.now()
       });
       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
+        }
       });
     }
     return chunks;
   }
-  emitToolResults(results, doneChunk) {
+  emitToolResults(results, finishEvent) {
     const chunks = [];
     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 || void 0,
         toolCallId: result.toolCallId,
         toolName: result.toolName,
@@ -483,15 +519,14 @@ class TextEngine {
         timestamp: Date.now()
       });
       const content = JSON.stringify(result.result);
-      const chunk = {
-        type: "tool_result",
-        id: doneChunk.id,
-        model: doneChunk.model,
+      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,
         {
@@ -500,14 +535,35 @@ class TextEngine {
           toolCallId: result.toolCallId
         }
       ];
+      aiEventClient.emit("text:message:created", {
+        ...this.buildTextEventContext(),
+        messageId: this.createId("msg"),
+        role: "tool",
+        content,
+        timestamp: Date.now()
+      });
     }
     return chunks;
   }
   getPendingToolCallsFromMessages() {
-    const completedToolIds = new Set(
-      this.messages.filter((message) => message.role === "tool" && message.toolCallId).map((message) => message.toolCallId)
-      // toolCallId exists due to filter
-    );
+    const completedToolIds = /* @__PURE__ */ new Set();
+    for (const message of this.messages) {
+      if (message.role === "tool" && message.toolCallId) {
+        let hasPendingExecution = false;
+        if (typeof message.content === "string") {
+          try {
+            const parsed = JSON.parse(message.content);
+            if (parsed.pendingExecution === true) {
+              hasPendingExecution = true;
+            }
+          } catch {
+          }
+        }
+        if (!hasPendingExecution) {
+          completedToolIds.add(message.toolCallId);
+        }
+      }
+    }
     const pending = [];
     for (const message of this.messages) {
       if (message.role === "assistant" && message.toolCalls) {
@@ -520,10 +576,10 @@ class TextEngine {
     }
     return pending;
   }
-  createSyntheticDoneChunk() {
+  createSyntheticFinishedEvent() {
     return {
-      type: "done",
-      id: this.createId("pending"),
+      type: "RUN_FINISHED",
+      runId: this.createId("pending"),
       model: this.params.model,
       timestamp: Date.now(),
       finishReason: "tool_calls"
@@ -542,6 +598,28 @@ class TextEngine {
   isAborted() {
     return !!this.effectiveSignal?.aborted;
   }
+  buildTextEventContext() {
+    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 : void 0,
+      toolNames: this.eventToolNames,
+      options: this.eventOptions,
+      modelOptions: this.params.modelOptions,
+      messageCount: this.initialMessageCount,
+      hasTools: this.tools.length > 0,
+      streaming: true
+    };
+  }
+  getContentString(content) {
+    if (typeof content === "string") return content;
+    const text = content?.map((part) => part.type === "text" ? part.content : "").join("") || "";
+    return text;
+  }
   setToolPhase(phase) {
     this.toolPhase = phase;
     if (phase === "wait") {