@tanstack/ai 0.47.3 → 0.48.0

package/src/utilities/restore-inbound-chunk.test.ts

@@ -0,0 +1,133 @@
+import { describe, expect, it } from 'vitest'
+import { EventType } from '../types'
+import type { RunFinishedEvent, StreamChunk } from '../types'
+import {
+  restoreInboundChunk,
+  restorePublicUsage,
+} from './restore-inbound-chunk'
+
+function runFinished(
+  overrides?: Omit<Partial<RunFinishedEvent>, 'type'>,
+): RunFinishedEvent {
+  return {
+    type: EventType.RUN_FINISHED,
+    threadId: 't1',
+    runId: 'r1',
+    ...overrides,
+  }
+}
+
+describe('restorePublicUsage', () => {
+  it('rebuilds TokenUsage promptTokens from spec usage[] and leftover', () => {
+    const chunk = runFinished({
+      usage: [
+        {
+          inputTokens: 10,
+          outputTokens: 5,
+          totalTokens: 15,
+          cachedInputTokens: 3,
+        },
+      ],
+      metadata: {
+        tanstack: {
+          model: 'gpt-5.5',
+          usage: {
+            cost: 0.02,
+            promptTokensDetails: { audioTokens: 1 },
+          },
+        },
+      },
+    })
+
+    restorePublicUsage(chunk)
+
+    expect(chunk.usage).toEqual({
+      promptTokens: 10,
+      completionTokens: 5,
+      totalTokens: 15,
+      cost: 0.02,
+      promptTokensDetails: { cachedTokens: 3, audioTokens: 1 },
+    })
+    expect(chunk).not.toHaveProperty('model')
+  })
+
+  it('restores TOOL_CALL_START toolName from toolCallName', () => {
+    const chunk: StreamChunk = {
+      type: EventType.TOOL_CALL_START,
+      toolCallId: 'tc1',
+      toolCallName: 'get_weather',
+    }
+
+    restorePublicUsage(chunk)
+
+    if (chunk.type !== EventType.TOOL_CALL_START) {
+      throw new Error('expected TOOL_CALL_START')
+    }
+    expect(chunk.toolName).toBe('get_weather')
+  })
+
+  it('restores TOOL_CALL_END input from metadata.tanstack.input', () => {
+    const chunk: StreamChunk = {
+      type: EventType.TOOL_CALL_END,
+      toolCallId: 'tc1',
+      metadata: { tanstack: { input: { q: 'sf' } } },
+    }
+
+    restorePublicUsage(chunk)
+
+    if (chunk.type !== EventType.TOOL_CALL_END) {
+      throw new Error('expected TOOL_CALL_END')
+    }
+    expect(chunk.input).toEqual({ q: 'sf' })
+  })
+})
+
+describe('restoreInboundChunk', () => {
+  it('copies metadata.tanstack extras back to top-level fields', () => {
+    const restored = restoreInboundChunk(
+      runFinished({
+        usage: [{ inputTokens: 10, outputTokens: 5, totalTokens: 15 }],
+        metadata: {
+          tanstack: {
+            model: 'gpt-5.5',
+            finishReason: 'stop',
+            usage: { cost: 0.02 },
+          },
+        },
+      }),
+    )
+
+    expect(restored.finishReason).toBe('stop')
+    expect(restored.model).toBe('gpt-5.5')
+    expect(restored.usage).toEqual({
+      promptTokens: 10,
+      completionTokens: 5,
+      totalTokens: 15,
+      cost: 0.02,
+    })
+  })
+
+  it('does not overwrite extras already on the chunk', () => {
+    const restored = restoreInboundChunk(
+      runFinished({
+        finishReason: 'length',
+        model: 'kept',
+        metadata: {
+          tanstack: { finishReason: 'stop', model: 'other' },
+        },
+      }),
+    )
+
+    expect(restored.finishReason).toBe('length')
+    expect(restored.model).toBe('kept')
+  })
+
+  it('leaves spec-only chunks unchanged when there is no tanstack bag', () => {
+    const chunk = {
+      type: EventType.TEXT_MESSAGE_CONTENT,
+      messageId: 'm1',
+      delta: 'Hi',
+    } as const satisfies StreamChunk
+    expect(restoreInboundChunk(chunk)).toBe(chunk)
+  })
+})

package/src/utilities/restore-inbound-chunk.ts

@@ -0,0 +1,72 @@
+import { EventType } from '../types'
+import type { StreamChunk } from '../types'
+import type { AdapterYieldChunk } from './adapter-yield-chunk'
+import { isTanstackUsage, rebuildTokenUsage } from './ag-ui-usage'
+import { tanstackMetadata } from './merge-metadata'
+
+/**
+ * Rebuild TanStack `TokenUsage` (`promptTokens`) from spec `usage[]` plus
+ * leftover fields in `metadata.tanstack.usage`. Restore in-process aliases
+ * that the wire does not keep (`toolName`, `TOOL_CALL_END.input`). Mutates
+ * in place so WeakMap run-id stamps stay attached.
+ */
+export function restorePublicUsage(chunk: StreamChunk): StreamChunk {
+  if (
+    (chunk.type === EventType.RUN_FINISHED ||
+      chunk.type === EventType.RUN_ERROR) &&
+    (Array.isArray(chunk.usage) || isTanstackUsage(chunk.usage))
+  ) {
+    const rebuilt = rebuildTokenUsage(
+      chunk.usage,
+      tanstackMetadata(chunk)?.usage,
+    )
+    if (rebuilt !== undefined) {
+      chunk.usage = rebuilt
+    }
+  }
+
+  if (
+    chunk.type === EventType.TOOL_CALL_START &&
+    chunk.toolName === undefined &&
+    chunk.toolCallName
+  ) {
+    chunk.toolName = chunk.toolCallName
+  }
+
+  if (chunk.type === EventType.TOOL_CALL_END && chunk.input === undefined) {
+    const input = tanstackMetadata(chunk)?.input
+    if (input !== undefined) {
+      chunk.input = input
+    }
+  }
+
+  return chunk
+}
+
+/**
+ * Rebuild the pre-wire chunk shape after SSE/HTTP/WS ingest.
+ * Copies `metadata.tanstack` extras back to top-level fields and rebuilds
+ * TanStack `TokenUsage` from spec `usage[]` when present.
+ */
+export function restoreInboundChunk(chunk: StreamChunk): AdapterYieldChunk {
+  restorePublicUsage(chunk)
+  const tanstack = tanstackMetadata(chunk)
+  const next = chunk as AdapterYieldChunk & Record<string, unknown>
+
+  if (tanstack == null) {
+    return next
+  }
+
+  for (const [key, value] of Object.entries(tanstack)) {
+    if (key === 'usage' || key === 'interruptErrors') continue
+    if (next[key] === undefined && value !== undefined) {
+      next[key] = value
+    }
+  }
+
+  if (tanstack.interruptErrors !== undefined) {
+    next['tanstack:interruptErrors'] = tanstack.interruptErrors
+  }
+
+  return next
+}

package/src/utilities/spec-event-keys.test.ts

@@ -0,0 +1,34 @@
+import { describe, expect, it } from 'vitest'
+import { EventType } from '../types'
+import { specKeysFor, isSpecTopLevelKey } from './spec-event-keys'
+
+describe('specKeysFor', () => {
+  it('allows metadata on every event and usage on RUN_FINISHED / RUN_ERROR', () => {
+    expect(isSpecTopLevelKey(EventType.RUN_FINISHED, 'usage')).toBe(true)
+    expect(isSpecTopLevelKey(EventType.RUN_ERROR, 'usage')).toBe(true)
+    expect(isSpecTopLevelKey(EventType.RUN_STARTED, 'metadata')).toBe(true)
+    expect(isSpecTopLevelKey(EventType.CUSTOM, 'name')).toBe(true)
+  })
+
+  it('rejects TanStack extras', () => {
+    expect(isSpecTopLevelKey(EventType.RUN_FINISHED, 'finishReason')).toBe(
+      false,
+    )
+    expect(isSpecTopLevelKey(EventType.RUN_FINISHED, 'model')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.TOOL_CALL_START, 'toolName')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.TOOL_CALL_START, 'index')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.TEXT_MESSAGE_CONTENT, 'content')).toBe(
+      false,
+    )
+    expect(isSpecTopLevelKey(EventType.TOOL_CALL_ARGS, 'args')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.TOOL_CALL_END, 'input')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.TOOL_CALL_END, 'result')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.RUN_ERROR, 'error')).toBe(false)
+    expect(
+      isSpecTopLevelKey(EventType.RUN_ERROR, 'tanstack:interruptErrors'),
+    ).toBe(false)
+    expect(isSpecTopLevelKey(EventType.STATE_SNAPSHOT, 'state')).toBe(false)
+    expect(isSpecTopLevelKey(EventType.CUSTOM, 'threadId')).toBe(false)
+    expect(specKeysFor(EventType.RUN_ERROR).has('threadId')).toBe(false)
+  })
+})

package/src/utilities/spec-event-keys.ts

@@ -0,0 +1,74 @@
+import { EventType } from '../types'
+
+const SHARED: ReadonlySet<string> = new Set([
+  'type',
+  'timestamp',
+  'rawEvent',
+  'metadata',
+])
+
+function keys(...fields: Array<string>): ReadonlySet<string> {
+  return new Set([...SHARED, ...fields])
+}
+
+const SPEC_KEYS: ReadonlyMap<string, ReadonlySet<string>> = new Map([
+  [EventType.TEXT_MESSAGE_START, keys('messageId', 'role', 'name')],
+  [EventType.TEXT_MESSAGE_CONTENT, keys('messageId', 'delta')],
+  [EventType.TEXT_MESSAGE_END, keys('messageId')],
+  [EventType.TEXT_MESSAGE_CHUNK, keys('messageId', 'role', 'delta', 'name')],
+  [
+    EventType.TOOL_CALL_START,
+    keys('toolCallId', 'toolCallName', 'parentMessageId'),
+  ],
+  [EventType.TOOL_CALL_ARGS, keys('toolCallId', 'delta')],
+  [EventType.TOOL_CALL_END, keys('toolCallId')],
+  [
+    EventType.TOOL_CALL_CHUNK,
+    keys('toolCallId', 'toolCallName', 'parentMessageId', 'delta'),
+  ],
+  [
+    EventType.TOOL_CALL_RESULT,
+    keys('messageId', 'toolCallId', 'content', 'role'),
+  ],
+  [EventType.THINKING_START, keys('title')],
+  [EventType.THINKING_END, SHARED],
+  [EventType.THINKING_TEXT_MESSAGE_START, SHARED],
+  [EventType.THINKING_TEXT_MESSAGE_CONTENT, keys('delta')],
+  [EventType.THINKING_TEXT_MESSAGE_END, SHARED],
+  [EventType.STATE_SNAPSHOT, keys('snapshot')],
+  [EventType.STATE_DELTA, keys('delta')],
+  [EventType.MESSAGES_SNAPSHOT, keys('messages')],
+  [
+    EventType.ACTIVITY_SNAPSHOT,
+    keys('messageId', 'activityType', 'content', 'replace'),
+  ],
+  [EventType.ACTIVITY_DELTA, keys('messageId', 'activityType', 'patch')],
+  [EventType.RAW, keys('event', 'source')],
+  [EventType.CUSTOM, keys('name', 'value')],
+  [EventType.RUN_STARTED, keys('threadId', 'runId', 'parentRunId', 'input')],
+  [
+    EventType.RUN_FINISHED,
+    keys('threadId', 'runId', 'result', 'outcome', 'usage'),
+  ],
+  [EventType.RUN_ERROR, keys('message', 'code', 'usage')],
+  [EventType.STEP_STARTED, keys('stepName')],
+  [EventType.STEP_FINISHED, keys('stepName')],
+  [EventType.REASONING_START, keys('messageId')],
+  [EventType.REASONING_MESSAGE_START, keys('messageId', 'role')],
+  [EventType.REASONING_MESSAGE_CONTENT, keys('messageId', 'delta')],
+  [EventType.REASONING_MESSAGE_END, keys('messageId')],
+  [EventType.REASONING_MESSAGE_CHUNK, keys('messageId', 'delta')],
+  [EventType.REASONING_END, keys('messageId')],
+  [
+    EventType.REASONING_ENCRYPTED_VALUE,
+    keys('subtype', 'entityId', 'encryptedValue'),
+  ],
+])
+
+export function specKeysFor(type: string): ReadonlySet<string> {
+  return SPEC_KEYS.get(type) ?? SHARED
+}
+
+export function isSpecTopLevelKey(type: string, key: string): boolean {
+  return specKeysFor(type).has(key)
+}

package/src/utilities/structured-output-events.ts

@@ -1,5 +1,5 @@
 import { EventType } from '../types'
-import type { StreamChunk } from '../types'
+import type { AdapterYieldChunk } from './adapter-yield-chunk'
 
 export function structuredOutputStartChunk(args: {
   messageId: string
@@ -7,7 +7,7 @@ export function structuredOutputStartChunk(args: {
   threadId: string
   runId: string
   timestamp?: number
-}): StreamChunk {
+}): AdapterYieldChunk {
   return {
     type: EventType.CUSTOM,
     name: 'structured-output.start',
@@ -27,7 +27,7 @@ export function structuredOutputCompleteChunk(args: {
   object: unknown
   raw: string
   timestamp?: number
-}): StreamChunk {
+}): AdapterYieldChunk {
   return {
     type: EventType.CUSTOM,
     name: 'structured-output.complete',