@tanstack/ai 0.22.0 → 0.23.0

package/dist/esm/activities/chat/middleware/types.d.ts

@@ -1,4 +1,4 @@
-import { JSONSchema, ModelMessage, StreamChunk, Tool, ToolCall } from '../../../types.js';
+import { JSONSchema, ModelMessage, StreamChunk, Tool, ToolCall, UsageTotals } from '../../../types.js';
 import { SystemPrompt } from '../../../system-prompts.js';
 /**
  * Phase of the chat middleware lifecycle.
@@ -20,6 +20,8 @@ export interface ChatMiddlewareContext {
     requestId: string;
     /** Unique identifier for this stream */
     streamId: string;
+    /** AG-UI run identifier for correlating client and server events */
+    runId: string;
     /**
      * AG-UI thread identifier — a stable per-conversation ID used to
      * correlate client and server devtools events. Resolves to the
@@ -201,11 +203,12 @@ export interface ToolPhaseCompleteInfo {
 /**
  * Token usage statistics passed to the onUsage hook.
  * Extracted from the RUN_FINISHED chunk when usage data is present.
+ *
+ * Includes optional provider-reported `cost`/`costDetails` (see {@link UsageTotals}).
+ * Kept as an interface extending `UsageTotals` to preserve declaration merging for
+ * this publicly exported type.
  */
-export interface UsageInfo {
-    promptTokens: number;
-    completionTokens: number;
-    totalTokens: number;
+export interface UsageInfo extends UsageTotals {
 }
 /**
  * Information passed to onFinish.
@@ -217,12 +220,8 @@ export interface FinishInfo {
     duration: number;
     /** Final accumulated text content */
     content: string;
-    /** Final usage totals, if available */
-    usage?: {
-        promptTokens: number;
-        completionTokens: number;
-        totalTokens: number;
-    } | undefined;
+    /** Final usage totals, if available (optionally including provider-reported cost) */
+    usage?: UsageTotals | undefined;
 }
 /**
  * Information passed to onAbort.

package/dist/esm/activities/chat/stream/processor.d.ts

@@ -25,6 +25,17 @@ export interface StreamProcessorEvents {
     onTextUpdate?: (messageId: string, content: string) => void;
     onToolCallStateChange?: (messageId: string, toolCallId: string, state: ToolCallState, args: string) => void;
     onThinkingUpdate?: (messageId: string, stepId: string, content: string) => void;
+    onStructuredOutputChange?: (args: {
+        phase: 'start' | 'update' | 'complete' | 'error';
+        messageId: string;
+        status: 'streaming' | 'complete' | 'error';
+        raw: string;
+        partial?: unknown;
+        data?: unknown;
+        reasoning?: string;
+        errorMessage?: string;
+        delta?: string;
+    }) => void;
 }
 /**
  * Options for StreamProcessor
@@ -69,6 +80,7 @@ export declare class StreamProcessor {
     private pendingManualMessageId;
     private pendingThinkingStepId;
     private readonly structuredMessageIds;
+    private readonly structuredOutputUpdateBatches;
     private readonly activeRuns;
     private finishReason;
     private hasError;
@@ -359,6 +371,9 @@ export declare class StreamProcessor {
      * @see docs/chat-architecture.md#uimessage-part-ordering-invariants — Replace vs. push logic
      */
     private emitTextUpdateForMessage;
+    private queueStructuredOutputUpdate;
+    private flushStructuredOutputUpdate;
+    private emitStructuredOutputChange;
     /**
      * Emit messages change event
      */

package/dist/esm/activities/chat/stream/processor.js

@@ -2,6 +2,7 @@ import { generateMessageId, uiMessageToModelMessages } from "../messages.js";
 import { defaultJSONParser } from "./json-parser.js";
 import { updateToolCallWithOutput, updateToolResultPart, updateToolCallApprovalResponse, appendStructuredOutputDelta, updateToolCallPart, errorStructuredOutputPart, updateThinkingPart, completeStructuredOutputPart, updateToolCallApproval, updateTextPart } from "./message-updaters.js";
 import { ImmediateStrategy } from "./strategies.js";
+const STRUCTURED_OUTPUT_UPDATE_BATCH_SIZE = 12;
 class StreamProcessor {
   chunkStrategy;
   events;
@@ -16,6 +17,7 @@ class StreamProcessor {
   pendingManualMessageId = null;
   pendingThinkingStepId = null;
   structuredMessageIds = /* @__PURE__ */ new Set();
+  structuredOutputUpdateBatches = /* @__PURE__ */ new Map();
   // Run tracking (for concurrent run safety)
   activeRuns = /* @__PURE__ */ new Set();
   // Shared stream state
@@ -205,6 +207,9 @@ class StreamProcessor {
     for (const id of this.structuredMessageIds) {
       if (!keptIds.has(id)) this.structuredMessageIds.delete(id);
     }
+    for (const id of this.structuredOutputUpdateBatches.keys()) {
+      if (!keptIds.has(id)) this.structuredOutputUpdateBatches.delete(id);
+    }
     for (const id of this.messageStates.keys()) {
       if (!keptIds.has(id)) this.messageStates.delete(id);
     }
@@ -226,6 +231,7 @@ class StreamProcessor {
     this.activeMessageIds.clear();
     this.toolCallToMessage.clear();
     this.structuredMessageIds.clear();
+    this.structuredOutputUpdateBatches.clear();
     this.pendingManualMessageId = null;
     this.emitMessagesChange();
   }
@@ -579,6 +585,7 @@ class StreamProcessor {
           delta2
         );
         state.totalTextContent += delta2;
+        this.queueStructuredOutputUpdate(messageId, delta2);
         this.emitMessagesChange();
       }
       return;
@@ -846,12 +853,14 @@ class StreamProcessor {
       );
     }
     if (this.structuredMessageIds.has(messageId)) {
+      this.flushStructuredOutputUpdate(messageId);
       this.messages = errorStructuredOutputPart(
         this.messages,
         messageId,
         errorMessage
       );
       this.structuredMessageIds.delete(messageId);
+      this.emitStructuredOutputChange(messageId, "error");
       this.emitMessagesChange();
     }
     this.events.onError?.(new Error(errorMessage));
@@ -994,6 +1003,13 @@ class StreamProcessor {
       if (targetId) {
         this.ensureAssistantMessage(targetId);
         this.structuredMessageIds.add(targetId);
+        this.structuredOutputUpdateBatches.delete(targetId);
+        this.events.onStructuredOutputChange?.({
+          phase: "start",
+          messageId: targetId,
+          status: "streaming",
+          raw: ""
+        });
       }
       return;
     }
@@ -1001,6 +1017,7 @@ class StreamProcessor {
       const v = chunk.value;
       const targetId = v.messageId ?? messageId;
       if (targetId) {
+        this.flushStructuredOutputUpdate(targetId);
         this.messages = completeStructuredOutputPart(
           this.messages,
           targetId,
@@ -1009,6 +1026,7 @@ class StreamProcessor {
           v.reasoning
         );
         this.structuredMessageIds.delete(targetId);
+        this.emitStructuredOutputChange(targetId, "complete");
         this.emitMessagesChange();
       }
     }
@@ -1132,6 +1150,40 @@ class StreamProcessor {
     this.emitMessagesChange();
     this.events.onTextUpdate?.(messageId, state.currentSegmentText);
   }
+  queueStructuredOutputUpdate(messageId, delta) {
+    const existing = this.structuredOutputUpdateBatches.get(messageId);
+    const next = {
+      delta: `${existing?.delta ?? ""}${delta}`,
+      chunkCount: (existing?.chunkCount ?? 0) + 1
+    };
+    this.structuredOutputUpdateBatches.set(messageId, next);
+    if (next.chunkCount >= STRUCTURED_OUTPUT_UPDATE_BATCH_SIZE) {
+      this.flushStructuredOutputUpdate(messageId);
+    }
+  }
+  flushStructuredOutputUpdate(messageId) {
+    const batch = this.structuredOutputUpdateBatches.get(messageId);
+    if (!batch || batch.chunkCount === 0) return;
+    this.structuredOutputUpdateBatches.delete(messageId);
+    this.emitStructuredOutputChange(messageId, "update", batch.delta);
+  }
+  emitStructuredOutputChange(messageId, phase, delta) {
+    const part = this.messages.find((message) => message.id === messageId)?.parts.find(
+      (messagePart) => messagePart.type === "structured-output"
+    );
+    if (!part) return;
+    this.events.onStructuredOutputChange?.({
+      phase,
+      messageId,
+      status: part.status,
+      raw: part.raw,
+      ...part.partial !== void 0 ? { partial: part.partial } : {},
+      ...part.data !== void 0 ? { data: part.data } : {},
+      ...part.reasoning !== void 0 ? { reasoning: part.reasoning } : {},
+      ...part.errorMessage !== void 0 ? { errorMessage: part.errorMessage } : {},
+      ...delta !== void 0 ? { delta } : {}
+    });
+  }
   /**
    * Emit messages change event
    */
@@ -1163,13 +1215,16 @@ class StreamProcessor {
       }
     }
     for (const messageId of this.structuredMessageIds) {
+      this.flushStructuredOutputUpdate(messageId);
       this.messages = errorStructuredOutputPart(
         this.messages,
         messageId,
         "Stream ended without structured-output.complete"
       );
+      this.emitStructuredOutputChange(messageId, "error");
     }
     this.structuredMessageIds.clear();
+    this.structuredOutputUpdateBatches.clear();
     this.activeMessageIds.clear();
     if (lastAssistantMessage && !this.hasError) {
       if (this.isWhitespaceOnlyMessage(lastAssistantMessage)) {
@@ -1283,6 +1338,7 @@ class StreamProcessor {
     this.activeRuns.clear();
     this.toolCallToMessage.clear();
     this.structuredMessageIds.clear();
+    this.structuredOutputUpdateBatches.clear();
     this.pendingManualMessageId = null;
     this.pendingThinkingStepId = null;
     this.finishReason = null;