@tanstack/ai 0.18.0 → 0.19.0

package/src/activities/chat/stream/message-updaters.ts

@@ -5,7 +5,9 @@
  * These are used by StreamProcessor to manage the message array.
  */
 
+import { parsePartialJSON } from './json-parser'
 import type {
+  StructuredOutputPart,
   ThinkingPart,
   ToolCallPart,
   ToolResultPart,
@@ -251,6 +253,175 @@ export function updateToolCallApprovalResponse(
   })
 }
 
+/**
+ * Append a delta to the structured-output part on `messageId`, or create one
+ * if absent. Progressive parse of the accumulated buffer fills `partial`.
+ *
+ * Callers must only invoke this while the part is still in flight — the
+ * helper unconditionally writes `status: 'streaming'`, so feeding it a delta
+ * after a `complete`/`error` terminal would regress the part. In practice the
+ * processor gates calls via `structuredMessageIds`, which is dropped on
+ * terminal events.
+ *
+ * If the progressive parse returns null/undefined (the buffer is not yet a
+ * parseable JSON prefix), the previously-good `partial` is preserved so the
+ * UI doesn't flicker back to empty for a single render.
+ */
+export function appendStructuredOutputDelta(
+  messages: Array<UIMessage>,
+  messageId: string,
+  delta: string,
+): Array<UIMessage> {
+  return messages.map((msg) => {
+    if (msg.id !== messageId) {
+      return msg
+    }
+
+    const parts = [...msg.parts]
+    const existingIndex = parts.findIndex(
+      (p): p is StructuredOutputPart => p.type === 'structured-output',
+    )
+    const existing =
+      existingIndex >= 0 ? (parts[existingIndex] as StructuredOutputPart) : null
+
+    const nextRaw = (existing?.raw ?? '') + delta
+    const progressive = parsePartialJSON(nextRaw)
+    const nextPartial =
+      progressive !== undefined && progressive !== null
+        ? progressive
+        : existing?.partial
+
+    const nextPart: StructuredOutputPart = {
+      type: 'structured-output',
+      status: 'streaming',
+      raw: nextRaw,
+      ...(nextPartial !== undefined ? { partial: nextPartial } : {}),
+      ...(existing?.reasoning !== undefined
+        ? { reasoning: existing.reasoning }
+        : {}),
+    }
+
+    if (existingIndex >= 0) {
+      parts[existingIndex] = nextPart
+    } else {
+      parts.push(nextPart)
+    }
+
+    return { ...msg, parts }
+  })
+}
+
+/**
+ * Snap the structured-output part on `messageId` to `complete` with the
+ * validated `data`. Picks the freshest available `raw` so the wire
+ * round-trip stays internally consistent:
+ *
+ *   1. Caller-supplied `raw` (the original streamed bytes from the model).
+ *   2. The existing part's `raw` (deltas accumulated before this terminal).
+ *   3. `JSON.stringify(data)` as a defensive fallback for terminal-only
+ *      completes that never shipped raw — keeps the part self-consistent
+ *      so downstream consumers never see a complete part with empty raw.
+ */
+export function completeStructuredOutputPart(
+  messages: Array<UIMessage>,
+  messageId: string,
+  data: unknown,
+  raw: string,
+  reasoning?: string,
+): Array<UIMessage> {
+  return messages.map((msg) => {
+    if (msg.id !== messageId) {
+      return msg
+    }
+
+    const parts = [...msg.parts]
+    const existingIndex = parts.findIndex(
+      (p): p is StructuredOutputPart => p.type === 'structured-output',
+    )
+
+    const existingRaw =
+      existingIndex >= 0
+        ? (parts[existingIndex] as StructuredOutputPart).raw
+        : ''
+    let resolvedRaw = raw || existingRaw
+    if (resolvedRaw === '' && data !== undefined) {
+      try {
+        resolvedRaw = JSON.stringify(data)
+      } catch {
+        // Unserializable (circular, BigInt, throwing toJSON). Leave raw
+        // empty. Both downstream paths handle this: `ag-ui-wire.ts`
+        // `collectText` skips complete parts with empty raw entirely, and
+        // `uiMessageToModelMessages` falls back to a defensive
+        // `safeJsonStringify(data)` which itself returns `''` for the same
+        // unserializable inputs — so the turn is silently dropped from the
+        // next request rather than shipping garbage or crashing the stream.
+      }
+    }
+
+    const nextPart: StructuredOutputPart = {
+      type: 'structured-output',
+      status: 'complete',
+      data,
+      partial: data,
+      raw: resolvedRaw,
+      ...(reasoning !== undefined ? { reasoning } : {}),
+    }
+
+    if (existingIndex >= 0) {
+      parts[existingIndex] = nextPart
+    } else {
+      parts.push(nextPart)
+    }
+
+    return { ...msg, parts }
+  })
+}
+
+/**
+ * Mark the structured-output part on `messageId` as errored. If no part
+ * exists yet — RUN_ERROR fired after `structured-output.start` but before
+ * any delta — create an empty errored placeholder so consumers have
+ * something renderable. Existing complete parts are left alone (an error
+ * after a successful complete should not retroactively un-complete it).
+ */
+export function errorStructuredOutputPart(
+  messages: Array<UIMessage>,
+  messageId: string,
+  errorMessage: string,
+): Array<UIMessage> {
+  return messages.map((msg) => {
+    if (msg.id !== messageId) {
+      return msg
+    }
+
+    const parts = [...msg.parts]
+    const existingIndex = parts.findIndex(
+      (p): p is StructuredOutputPart => p.type === 'structured-output',
+    )
+
+    if (existingIndex < 0) {
+      parts.push({
+        type: 'structured-output',
+        status: 'error',
+        raw: '',
+        errorMessage,
+      })
+      return { ...msg, parts }
+    }
+
+    const existing = parts[existingIndex] as StructuredOutputPart
+    if (existing.status === 'complete') {
+      return msg
+    }
+    parts[existingIndex] = {
+      ...existing,
+      status: 'error',
+      errorMessage,
+    }
+    return { ...msg, parts }
+  })
+}
+
 /**
  * Update or add a thinking part to a message, keyed by stepId.
  * Each distinct stepId produces its own ThinkingPart.

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

@@ -20,6 +20,9 @@
 import { generateMessageId, uiMessageToModelMessages } from '../messages.js'
 import { defaultJSONParser } from './json-parser'
 import {
+  appendStructuredOutputDelta,
+  completeStructuredOutputPart,
+  errorStructuredOutputPart,
   updateTextPart,
   updateThinkingPart,
   updateToolCallApproval,
@@ -145,6 +148,8 @@ export class StreamProcessor {
   private pendingManualMessageId: string | null = null
   private pendingThinkingStepId: string | null = null
 
+  private structuredMessageIds: Set<string> = new Set()
+
   // Run tracking (for concurrent run safety)
   private activeRuns = new Set<string>()
 
@@ -383,6 +388,27 @@ export class StreamProcessor {
    * Remove messages after a certain index (for reload/retry)
    */
   removeMessagesAfter(index: number): void {
+    const keptIds = new Set(this.messages.slice(0, index + 1).map((m) => m.id))
+    // Drop routing state for messages that no longer exist; otherwise a
+    // resumed stream (`reload()` or a server that reuses messageIds across
+    // runs) could land deltas / tool args on stale map entries and corrupt
+    // the new assistant message's parts. Mirror the four routing maps that
+    // key on messageId: structuredMessageIds (custom-event routing),
+    // messageStates (per-message stream state), toolCallToMessage (tool
+    // args → message), and activeMessageIds (finalize / completeAllToolCalls
+    // iteration targets — must not include phantoms).
+    for (const id of this.structuredMessageIds) {
+      if (!keptIds.has(id)) this.structuredMessageIds.delete(id)
+    }
+    for (const id of this.messageStates.keys()) {
+      if (!keptIds.has(id)) this.messageStates.delete(id)
+    }
+    for (const [toolCallId, msgId] of this.toolCallToMessage) {
+      if (!keptIds.has(msgId)) this.toolCallToMessage.delete(toolCallId)
+    }
+    for (const id of this.activeMessageIds) {
+      if (!keptIds.has(id)) this.activeMessageIds.delete(id)
+    }
     this.messages = this.messages.slice(0, index + 1)
     this.emitMessagesChange()
   }
@@ -395,6 +421,7 @@ export class StreamProcessor {
     this.messageStates.clear()
     this.activeMessageIds.clear()
     this.toolCallToMessage.clear()
+    this.structuredMessageIds.clear()
     this.pendingManualMessageId = null
     this.emitMessagesChange()
   }
@@ -841,6 +868,41 @@ export class StreamProcessor {
     // Content arriving means all current tool calls for this message are complete
     this.completeAllToolCallsForMessage(messageId)
 
+    if (this.structuredMessageIds.has(messageId)) {
+      // `chunk.delta` is incremental; `chunk.content` is sometimes cumulative
+      // (mirrors what the plain-text branch handles below). Reconcile against
+      // the existing raw buffer so adapters that emit cumulative content
+      // don't duplicate the JSON.
+      let delta = chunk.delta || ''
+      if (delta === '' && chunk.content !== undefined && chunk.content !== '') {
+        const existingRaw = (
+          this.messages
+            .find((m) => m.id === messageId)
+            ?.parts.find(
+              (p): p is Extract<MessagePart, { type: 'structured-output' }> =>
+                p.type === 'structured-output',
+            ) ?? { raw: '' }
+        ).raw
+        if (chunk.content.startsWith(existingRaw)) {
+          delta = chunk.content.slice(existingRaw.length)
+        } else if (existingRaw.startsWith(chunk.content)) {
+          delta = ''
+        } else {
+          delta = chunk.content
+        }
+      }
+      if (delta !== '') {
+        this.messages = appendStructuredOutputDelta(
+          this.messages,
+          messageId,
+          delta,
+        )
+        state.totalTextContent += delta
+        this.emitMessagesChange()
+      }
+      return
+    }
+
     const previousSegment = state.currentSegmentText
 
     // Detect if this is a NEW text segment (after tool calls) vs continuation
@@ -1192,10 +1254,29 @@ export class StreamProcessor {
     } else {
       this.activeRuns.clear()
     }
-    this.ensureAssistantMessage()
-    // Prefer spec field `message`; fall back to deprecated `error.message`
+    const { messageId } = this.ensureAssistantMessage()
+    // Prefer spec field `message`; fall back to deprecated `error.message`.
+    // If neither is set, the chunk still carries debug context (provider
+    // error codes, request ids, etc.) — log it so the failure isn't silent.
     const errorMessage =
       chunk.message || chunk.error?.message || 'An error occurred'
+    if (!chunk.message && !chunk.error?.message) {
+      console.error(
+        '[StreamProcessor] RUN_ERROR with no message; original chunk:',
+        chunk,
+      )
+    }
+
+    if (this.structuredMessageIds.has(messageId)) {
+      this.messages = errorStructuredOutputPart(
+        this.messages,
+        messageId,
+        errorMessage,
+      )
+      this.structuredMessageIds.delete(messageId)
+      this.emitMessagesChange()
+    }
+
     this.events.onError?.(new Error(errorMessage))
   }
 
@@ -1375,6 +1456,38 @@ export class StreamProcessor {
   ): void {
     const messageId = this.getActiveAssistantMessageId()
 
+    if (chunk.name === 'structured-output.start' && chunk.value) {
+      const v = chunk.value as { messageId?: string }
+      const targetId = v.messageId ?? messageId
+      if (targetId) {
+        this.ensureAssistantMessage(targetId)
+        this.structuredMessageIds.add(targetId)
+      }
+      return
+    }
+
+    if (chunk.name === 'structured-output.complete' && chunk.value) {
+      const v = chunk.value as {
+        object: unknown
+        raw?: string
+        reasoning?: string
+        messageId?: string
+      }
+      const targetId = v.messageId ?? messageId
+      if (targetId) {
+        this.messages = completeStructuredOutputPart(
+          this.messages,
+          targetId,
+          v.object,
+          v.raw ?? '',
+          v.reasoning,
+        )
+        this.structuredMessageIds.delete(targetId)
+        this.emitMessagesChange()
+      }
+      // Fall through so user `onCustomEvent` callbacks still observe the event.
+    }
+
     // Handle client tool input availability - trigger client-side execution
     if (chunk.name === 'tool-input-available' && chunk.value) {
       const { toolCallId, toolName, input } = chunk.value as {
@@ -1591,6 +1704,27 @@ export class StreamProcessor {
       }
     }
 
+    // The stream closed but one or more structured-output runs never sent
+    // their terminal `structured-output.complete`. Snap each lingering
+    // streaming part to error so the UI doesn't appear to stream forever,
+    // and drop the routing entries so a subsequent run on the same
+    // processor instance (long-lived `subscribe()` mode) doesn't reuse
+    // the stale ids.
+    //
+    // The iteration is unconditional w.r.t. `this.hasError` — RUN_ERROR
+    // already removed its target messageId from `structuredMessageIds`
+    // before reaching finalize, so anything still in the set is by
+    // definition a non-errored, never-completed run (the multi-run case:
+    // run-A errors, run-B is still streaming when finalize fires).
+    for (const messageId of this.structuredMessageIds) {
+      this.messages = errorStructuredOutputPart(
+        this.messages,
+        messageId,
+        'Stream ended without structured-output.complete',
+      )
+    }
+    this.structuredMessageIds.clear()
+
     this.activeMessageIds.clear()
 
     // Remove whitespace-only assistant messages (handles models like Gemini
@@ -1719,6 +1853,7 @@ export class StreamProcessor {
     this.activeMessageIds.clear()
     this.activeRuns.clear()
     this.toolCallToMessage.clear()
+    this.structuredMessageIds.clear()
     this.pendingManualMessageId = null
     this.pendingThinkingStepId = null
     this.finishReason = null

package/src/adapter-internals.ts

@@ -4,5 +4,6 @@
 
 export type { ResolvedCategories } from './logger/internal-logger'
 export { InternalLogger } from './logger/internal-logger'
+export type { Logger } from './logger/types'
 export { resolveDebugOption } from './logger/resolve'
 export { toRunErrorPayload } from './activities/error-payload'

package/src/types.ts

@@ -365,7 +365,43 @@ export interface ThinkingPart {
   signature?: string
 }
 
-export type MessagePart =
+/**
+ * Recursive `Partial` — every nested field becomes optional. Used as the
+ * `partial` type on a streaming structured-output part since the progressive
+ * JSON parse hands back objects whose fields are only filled in as bytes
+ * arrive. Defaulted in `DeepPartial<unknown>` → `unknown` so untyped parts
+ * keep their existing shape.
+ */
+export type DeepPartial<T> =
+  T extends ReadonlyArray<infer U>
+    ? Array<DeepPartial<U>>
+    : T extends object
+      ? { [K in keyof T]?: DeepPartial<T[K]> }
+      : T
+
+/**
+ * StructuredOutputPart — a typed structured response attached to the assistant
+ * message that produced it. Generic over the schema-inferred data type so
+ * consumers can thread `useChat({ outputSchema })`'s schema all the way down
+ * to `messages[i].parts[j].data`. Defaults to `unknown` so untyped consumers
+ * (e.g. internal codepaths that don't know about TSchema) keep working.
+ */
+export interface StructuredOutputPart<TData = unknown> {
+  type: 'structured-output'
+  status: 'streaming' | 'complete' | 'error'
+  /** Progressive parse of `raw` via parsePartialJSON — populated while streaming and after complete. */
+  partial?: DeepPartial<TData>
+  /** Validated final object — only set when `status === 'complete'`. */
+  data?: TData
+  /** Accumulating JSON buffer. Source of truth for wire round-trip. */
+  raw: string
+  /** Optional chain-of-thought surfaced by reasoning models alongside the structured output. */
+  reasoning?: string
+  /** Populated when `status === 'error'`. */
+  errorMessage?: string
+}
+
+export type MessagePart<TData = unknown> =
   | TextPart
   | ImagePart
   | AudioPart
@@ -374,15 +410,19 @@ export type MessagePart =
   | ToolCallPart
   | ToolResultPart
   | ThinkingPart
+  | StructuredOutputPart<TData>
 
 /**
  * UIMessage - Domain-specific message format optimized for building chat UIs
- * Contains parts that can be text, tool calls, or tool results
+ * Contains parts that can be text, tool calls, or tool results. Generic over
+ * the structured-output data type so `useChat({ outputSchema })`'s schema
+ * narrows `parts.find(p => p.type === 'structured-output').data` on the
+ * consumer side without manual casts.
  */
-export interface UIMessage {
+export interface UIMessage<TData = unknown> {
   id: string
   role: 'system' | 'user' | 'assistant'
-  parts: Array<MessagePart>
+  parts: Array<MessagePart<TData>>
   createdAt?: Date
 }
 
@@ -1114,11 +1154,28 @@ export interface StructuredOutputCompleteEvent<T = unknown> extends Omit<
   value: { object: T; raw: string; reasoning?: string }
 }
 
+/**
+ * Emitted at the start of a streaming structured-output run, before the JSON
+ * deltas. Tells consumers that the upcoming `TEXT_MESSAGE_CONTENT` deltas
+ * belong to a structured response so they can route those bytes into a
+ * `StructuredOutputPart` instead of building a `TextPart`. Carries the
+ * `messageId` the deltas will be tagged with so the routing decision can be
+ * made per-message rather than globally.
+ */
+export interface StructuredOutputStartEvent extends Omit<
+  CustomEvent,
+  'name' | 'value'
+> {
+  name: 'structured-output.start'
+  value: { messageId: string }
+}
+
 /**
  * Emitted when a server tool requires approval before execution. The agent
  * loop yields this and pauses — `structured-output.complete` will not fire
  * for that run. The shape is fixed by the orchestrator's tool-approval flow
- * (see `buildApprovalChunks` in `activities/chat/index.ts`).
+ * (the agent-loop branch of `runStreamingStructuredOutputImpl` in
+ * `activities/chat/index.ts` forwards CUSTOM events from `TextEngine.run()`).
  */
 export interface ApprovalRequestedEvent extends Omit<
   CustomEvent,
@@ -1136,8 +1193,8 @@ export interface ApprovalRequestedEvent extends Omit<
 /**
  * Emitted when a client tool is invoked. The agent loop yields this and
  * pauses to let the caller run the tool client-side — `structured-output.complete`
- * will not fire for that run. Shape fixed by `buildClientToolChunks` in
- * `activities/chat/index.ts`.
+ * will not fire for that run. Shape fixed by the agent-loop forwarding in
+ * `runStreamingStructuredOutputImpl` in `activities/chat/index.ts`.
  */
 export interface ToolInputAvailableEvent extends Omit<
   CustomEvent,
@@ -1183,6 +1240,7 @@ export interface ToolInputAvailableEvent extends Omit<
  */
 export type StructuredOutputStream<T = unknown> = AsyncIterable<
   | Exclude<StreamChunk, CustomEvent>
+  | StructuredOutputStartEvent
   | StructuredOutputCompleteEvent<T>
   | ApprovalRequestedEvent
   | ToolInputAvailableEvent

package/src/utilities/ag-ui-wire.ts

@@ -1,4 +1,4 @@
-import type { ContentPart, MessagePart, TextPart, UIMessage } from '../types'
+import type { ContentPart, MessagePart, UIMessage } from '../types'
 
 type AGUITextInputContent = { type: 'text'; text: string }
 type AGUIInputContent =
@@ -113,10 +113,29 @@ export function uiMessagesToWire(
 }
 
 function collectText(parts: ReadonlyArray<MessagePart>): string {
-  return parts
-    .filter((p): p is TextPart => p.type === 'text')
-    .map((p) => p.content)
-    .join('')
+  // The streamed JSON of a completed structured-output part is the source of
+  // truth for multi-turn coherence — emitting it back as assistant content
+  // lets the LLM see its own prior structured response. Streaming/errored
+  // parts are skipped: they'd ship malformed JSON fragments and confuse the
+  // model. `completeStructuredOutputPart` tries hard to populate `raw`
+  // (caller → existing buffer → `JSON.stringify(data)`), but the stringify
+  // fallback can leave it empty when `data` is unserializable (BigInt,
+  // circular). The `p.raw !== ''` guard below is what enforces "no malformed
+  // round-trip" in that case — without it we'd ship `''` and the model would
+  // see an empty assistant turn.
+  const out: Array<string> = []
+  for (const p of parts) {
+    if (p.type === 'text') {
+      out.push(p.content)
+    } else if (
+      p.type === 'structured-output' &&
+      p.status === 'complete' &&
+      p.raw !== ''
+    ) {
+      out.push(p.raw)
+    }
+  }
+  return out.join('')
 }
 
 function collectUserContent(