@tanstack/ai-persistence 0.6.3 → 0.6.5

package/src/merge-stored.ts

@@ -0,0 +1,67 @@
+import type { ModelMessage } from '@tanstack/ai'
+
+// Empty incoming keeps stored. Non-empty: the last incoming id that already
+// exists in stored is a cutoff (reload drops the old assistant after that
+// user). Same id is replaced in place. New ids and messages with no id are
+// appended.
+export function mergeStoredMessages(
+  stored: ReadonlyArray<ModelMessage>,
+  incoming: ReadonlyArray<ModelMessage>,
+) {
+  if (incoming.length === 0) {
+    return stored.slice()
+  }
+
+  let cutoff = stored.length
+  for (let index = incoming.length - 1; index >= 0; index--) {
+    const id = incoming[index]?.id
+    if (id === undefined) continue
+    const storedIndex = stored.findIndex((message) => message.id === id)
+    if (storedIndex >= 0) {
+      cutoff = storedIndex + 1
+      break
+    }
+  }
+  const prefix = stored.slice(0, cutoff)
+
+  const incomingById = new Map<string, ModelMessage>()
+  for (const message of incoming) {
+    const id = message.id
+    if (id) incomingById.set(id, message)
+  }
+
+  const storedIds = new Set<string>()
+  const merged: Array<ModelMessage> = []
+  for (const message of prefix) {
+    const id = message.id
+    if (id) {
+      storedIds.add(id)
+      merged.push(incomingById.get(id) ?? message)
+      continue
+    }
+    merged.push(message)
+  }
+
+  for (let index = 0; index < incoming.length; index++) {
+    const message = incoming[index]
+    if (!message) continue
+    const id = message.id
+    if (id && storedIds.has(id)) continue
+    // Attach/reload can post the stored transcript again with no ids. Keep the
+    // prefix row instead of appending a second copy of the same turn.
+    if (!id) {
+      const existing = merged[index]
+      if (
+        existing &&
+        existing.id === undefined &&
+        existing.role === message.role &&
+        existing.content === message.content
+      ) {
+        continue
+      }
+    }
+    merged.push(message)
+  }
+
+  return merged
+}

package/src/middleware.ts

@@ -29,6 +29,8 @@ import {
   providePersistence,
   providePersistenceCompletion,
 } from './capabilities'
+import { mergeStoredMessages } from './merge-stored'
+import { createSubagentRunRecorder } from './subagent-runs'
 import {
   validateChatPersistenceStores,
   validateGenerationPersistenceStores,
@@ -1915,72 +1917,6 @@ function threadMessages(
   return Array.isArray(loaded) ? loaded : loaded.messages
 }
 
-// Empty incoming keeps stored. Non-empty: the last incoming id that already
-// exists in stored is a cutoff (reload drops the old assistant after that
-// user). Same id is replaced in place. New ids and messages with no id are
-// appended.
-function mergeStoredMessages(
-  stored: ReadonlyArray<ModelMessage>,
-  incoming: ReadonlyArray<ModelMessage>,
-) {
-  if (incoming.length === 0) {
-    return stored.slice()
-  }
-
-  let cutoff = stored.length
-  for (let index = incoming.length - 1; index >= 0; index--) {
-    const id = incoming[index]?.id
-    if (id === undefined) continue
-    const storedIndex = stored.findIndex((message) => message.id === id)
-    if (storedIndex >= 0) {
-      cutoff = storedIndex + 1
-      break
-    }
-  }
-  const prefix = stored.slice(0, cutoff)
-
-  const incomingById = new Map<string, ModelMessage>()
-  for (const message of incoming) {
-    const id = message.id
-    if (id) incomingById.set(id, message)
-  }
-
-  const storedIds = new Set<string>()
-  const merged: Array<ModelMessage> = []
-  for (const message of prefix) {
-    const id = message.id
-    if (id) {
-      storedIds.add(id)
-      merged.push(incomingById.get(id) ?? message)
-      continue
-    }
-    merged.push(message)
-  }
-
-  for (let index = 0; index < incoming.length; index++) {
-    const message = incoming[index]
-    if (!message) continue
-    const id = message.id
-    if (id && storedIds.has(id)) continue
-    // Attach/reload can post the stored transcript again with no ids. Keep the
-    // prefix row instead of appending a second copy of the same turn.
-    if (!id) {
-      const existing = merged[index]
-      if (
-        existing &&
-        existing.id === undefined &&
-        existing.role === message.role &&
-        existing.content === message.content
-      ) {
-        continue
-      }
-    }
-    merged.push(message)
-  }
-
-  return merged
-}
-
 export interface WithPersistenceOptions {
   /**
    * Also persist a throttled snapshot of the in-progress assistant reply while
@@ -1993,7 +1929,8 @@ export interface WithPersistenceOptions {
   snapshotStreaming?: boolean
   /**
    * Minimum milliseconds between streaming snapshots when `snapshotStreaming`
-   * is on. Defaults to 1000.
+   * is on. Defaults to 1000. A streaming subagent child always uses the same
+   * interval for its transcript writes, with or without `snapshotStreaming`.
    */
   snapshotIntervalMs?: number
 }
@@ -2044,8 +1981,18 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
     ...(wantsInterrupts ? [InterruptsCapability] : []),
   ]
 
+  const subagentRuns = createSubagentRunRecorder({
+    messages: messageStore,
+    runs,
+    intervalMs: snapshotIntervalMs,
+    ...(wantsInterrupts && persistence.stores.interrupts
+      ? { interrupts: persistence.stores.interrupts }
+      : {}),
+  })
+
   return defineChatMiddleware({
     name: 'chat-persistence',
+    routedSubagentPersistence: subagentRuns,
     provides,
     setup(ctx: ChatMiddlewareContext) {
       providePersistence(ctx, persistence)
@@ -2183,6 +2130,11 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
     },
 
     async onChunk(ctx: ChatMiddlewareContext, chunk: StreamChunk) {
+      await subagentRuns.chunk({
+        threadId: ctx.threadId,
+        runId: ctx.runId,
+        chunk,
+      })
       // Capture the current assistant turn's identity for optional in-progress
       // snapshots. Completed messages already live in `ctx.messages`.
       if (snapshotStreaming && ctx.phase === 'modelStream') {

package/src/reconstruct.ts

@@ -1,9 +1,16 @@
 import { modelMessagesToUIMessages } from '@tanstack/ai'
-import type { ModelMessage, UIMessage } from '@tanstack/ai'
+import type {
+  ModelMessage,
+  RunRecord,
+  SubagentPart,
+  UIMessage,
+} from '@tanstack/ai'
+import { storedSubagentInfo } from './subagent-runs'
 import { validateReconstructChatStores } from './types'
 import type {
   AIPersistence,
   ChatTranscriptStores,
+  InterruptRecord,
   MessagePage,
   MessageStore,
 } from './types'
@@ -171,8 +178,15 @@ export async function reconstructChat(
           before,
         })
       : windowFromMessagePage(stored, pageSize)
+  const messages = await attachSubagentCards(
+    transcript.messages,
+    persistence.stores.runs,
+    messageStore,
+    threadId,
+    pending,
+  )
   const body: ReconstructedChat = {
-    messages: transcript.messages,
+    messages,
     activeRun: active ? { runId: active.runId } : null,
     interrupts: firstPending
       ? {
@@ -190,6 +204,145 @@ export async function reconstructChat(
   })
 }
 
+function messageRunId(message: UIMessage) {
+  const metadata = message.metadata
+  if (!metadata || typeof metadata !== 'object') return
+  const tanstack = metadata.tanstack
+  if (!tanstack || typeof tanstack !== 'object') return
+  const runId = (tanstack as { runId?: unknown }).runId
+  return typeof runId === 'string' && runId !== '' ? runId : undefined
+}
+
+type Runs = NonNullable<ChatTranscriptStores['runs']>
+
+/** Rebuild one child card from its run record and stored transcript. */
+async function childCard(
+  child: RunRecord,
+  runs: Runs,
+  messageStore: MessageStore,
+  pending: ReadonlyArray<InterruptRecord>,
+  depth: number,
+): Promise<SubagentPart> {
+  const subagentRunId = child.subagentRunId ?? child.runId
+  const stored = await messageStore.loadThread(child.threadId)
+  const info = storedSubagentInfo(stored)
+  const messages = modelMessagesToUIMessages(
+    stored.filter(
+      (message) => storedSubagentInfo([message])?.placeholder !== true,
+    ),
+  )
+  const nested =
+    runs.listByParentRun && depth < 8
+      ? await runs.listByParentRun(subagentRunId)
+      : []
+  if (nested.length > 0) {
+    const cards = await Promise.all(
+      nested.map((run) =>
+        childCard(run, runs, messageStore, pending, depth + 1),
+      ),
+    )
+    const last = messages.findLastIndex((m) => m.role === 'assistant')
+    if (last === -1) {
+      messages.push({
+        id: `child-cards:${subagentRunId}`,
+        role: 'assistant',
+        parts: cards,
+      })
+    } else {
+      const host = messages[last]
+      if (host) messages[last] = { ...host, parts: [...host.parts, ...cards] }
+    }
+  }
+  const failed = child.status === 'failed' || child.status === 'aborted'
+  const interruptIds = pending
+    .filter((record) => record.payload.subagentRunId === subagentRunId)
+    .map((record) => record.interruptId)
+  return {
+    type: 'subagent',
+    subagent: {
+      id: subagentRunId,
+      name: child.name ?? info?.name ?? 'subagent',
+      status: failed
+        ? 'error'
+        : child.status === 'running'
+          ? 'running'
+          : child.status === 'interrupted'
+            ? 'suspended'
+            : 'finished',
+      ...(child.parentRunId !== undefined && {
+        parentRunId: child.parentRunId,
+      }),
+      ...(info?.parentToolCallId !== undefined && {
+        parentToolCallId: info.parentToolCallId,
+      }),
+      ...(interruptIds.length > 0 && { interruptIds }),
+      ...(info?.metadata !== undefined && { metadata: info.metadata }),
+      messages,
+      ...(failed && child.error ? { error: child.error } : {}),
+    },
+  }
+}
+
+/**
+ * Put stored subagent cards back on the transcript. A routed child sits on
+ * the parent assistant message of its run. A child that a tool call started
+ * sits on the message that holds that tool call.
+ */
+async function attachSubagentCards(
+  messages: Array<UIMessage>,
+  runs: ChatTranscriptStores['runs'],
+  messageStore: MessageStore,
+  threadId: string,
+  pending: ReadonlyArray<InterruptRecord>,
+) {
+  if (!runs?.listByParentRun) return messages
+  const parentRunIds = new Set<string>()
+  for (const message of messages) {
+    const runId = messageRunId(message)
+    if (runId) parentRunIds.add(runId)
+  }
+  const hasToolCalls = messages.some((message) =>
+    message.parts.some((part) => part.type === 'tool-call'),
+  )
+  if (hasToolCalls && runs.listByThread && threadId !== '') {
+    for (const run of await runs.listByThread(threadId)) {
+      parentRunIds.add(run.runId)
+    }
+  }
+
+  const cardsByRun = new Map<string, Array<SubagentPart>>()
+  const cardsByToolCall = new Map<string, Array<SubagentPart>>()
+  for (const runId of parentRunIds) {
+    for (const child of await runs.listByParentRun(runId)) {
+      const card = await childCard(child, runs, messageStore, pending, 0)
+      const toolCallId = card.subagent.parentToolCallId
+      const target = toolCallId === undefined ? cardsByRun : cardsByToolCall
+      const key = toolCallId ?? runId
+      target.set(key, [...(target.get(key) ?? []), card])
+    }
+  }
+  if (cardsByRun.size === 0 && cardsByToolCall.size === 0) return messages
+
+  return messages.map((message) => {
+    if (message.role !== 'assistant') return message
+    const runId = messageRunId(message)
+    const routed = runId !== undefined ? cardsByRun.get(runId) : undefined
+    const started = message.parts.flatMap((part) =>
+      part.type === 'tool-call' ? (cardsByToolCall.get(part.id) ?? []) : [],
+    )
+    if (!routed && started.length === 0) return message
+    // A routed parent message holds only the children's text. The cards
+    // replace it.
+    const parts = routed
+      ? message.parts.filter((part) => part.type !== 'text')
+      : message.parts
+    return {
+      ...message,
+      parts: [...(routed ?? []), ...parts, ...started],
+    }
+  })
+}
+
 function parsePageSize(raw: string | null) {
   if (raw == null) return
   const pageSize = Number(raw)