@tanstack/ai-persistence 0.6.7 → 0.7.2

package/src/middleware.ts

@@ -294,6 +294,11 @@ interface RunStateEntry {
    */
   streamingMessageId?: string
   streamingMessageCreatedAt?: Date
+  /**
+   * Length of `ctx.messages` when this run started. Messages from this index
+   * on are the ones this run added, so they get its run id on save (#1061).
+   */
+  firstRunMessage?: number
   completion?: {
     promise: Promise<void>
     resolve: () => void
@@ -302,6 +307,50 @@ interface RunStateEntry {
 }
 
 const runState = new WeakMap<object, RunStateEntry>()
+/** `metadata.tanstack.run.id` of a stored message, when set. */
+function runTagOf(message: ModelMessage): string | undefined {
+  const tanstack: unknown = message.metadata?.tanstack
+  if (typeof tanstack !== 'object' || tanstack === null) return
+  const run: unknown = (tanstack as { run?: unknown }).run
+  if (typeof run !== 'object' || run === null) return
+  const id: unknown = (run as { id?: unknown }).id
+  return typeof id === 'string' && id !== '' ? id : undefined
+}
+
+function withRunTag(message: ModelMessage, runId: string): ModelMessage {
+  const metadata = message.metadata ?? {}
+  const tanstack: unknown = metadata.tanstack
+  return {
+    ...message,
+    metadata: {
+      ...metadata,
+      tanstack: {
+        ...(typeof tanstack === 'object' && tanstack !== null ? tanstack : {}),
+        run: { id: runId },
+      },
+    },
+  }
+}
+
+/**
+ * The thread to save, with this run's id on the assistant messages it added
+ * (`metadata.tanstack.run.id`). `reconstructChat` uses it to match each
+ * message to its run. Messages from earlier runs are left as they are.
+ */
+function runMessages(
+  ctx: ChatMiddlewareContext,
+  state: RunStateEntry | undefined,
+): Array<ModelMessage> {
+  const from = state?.firstRunMessage
+  if (from === undefined) return [...ctx.messages]
+  return ctx.messages.map((message, index) =>
+    index >= from &&
+    message.role === 'assistant' &&
+    runTagOf(message) === undefined
+      ? withRunTag(message, ctx.runId)
+      : message,
+  )
+}
 
 const validResumeStatuses = new Set(['resolved', 'cancelled'])
 
@@ -333,6 +382,7 @@ function mergeResumeToolState(
       left.clientToolResults,
       right.clientToolResults,
     ),
+    clientToolErrors: mergeMaps(left.clientToolErrors, right.clientToolErrors),
     genericInterrupts: mergeMaps(
       left.genericInterrupts,
       right.genericInterrupts,
@@ -856,6 +906,7 @@ function resumeToolStateFromPending(
   for (const interrupt of pending) {
     const entry = resumeByInterruptId.get(interrupt.interruptId)
     if (!entry) continue
+    if (hasReservedInterruptBinding(interrupt.payload)) continue
 
     const kind = interruptKind(interrupt)
     const reason = stringField(interrupt.payload, 'reason')
@@ -2122,6 +2173,8 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       // prior history) as soon as the run starts, so a reload mid-run rehydrates
       // it before the assistant reply exists. Best-effort: a failed eager
       // snapshot must not abort the run — the authoritative save is `onFinish`.
+      const state = runState.get(ctx)
+      if (state) state.firstRunMessage = ctx.messages.length
       try {
         await messageStore.saveThread(ctx.threadId, [...ctx.messages])
       } catch {
@@ -2235,7 +2288,7 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
           : (state.usage ?? chunkUsage)
       state.usage = usage
       await interruptRun(runs, ctx.runId, usage)
-      await messageStore.saveThread(ctx.threadId, [...ctx.messages])
+      await messageStore.saveThread(ctx.threadId, runMessages(ctx, state))
       state.interrupted = true
     },
 
@@ -2253,7 +2306,7 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       // or consuming approvals before the durable history lands leaves a
       // "finished" run whose transcript is missing the terminal turn.
       try {
-        await messageStore.saveThread(ctx.threadId, [...ctx.messages])
+        await messageStore.saveThread(ctx.threadId, runMessages(ctx, state))
         await commitPendingResumes(state, persistence.stores.interrupts)
         await completeRun(runs, ctx.runId, state?.usage ?? info.usage)
         state?.completion?.resolve()
@@ -2286,9 +2339,9 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       // and durable (`RunRecord.cancelRequested`, the only channel that reaches
       // a run being driven elsewhere).
       // A run paused at an interrupt boundary is waiting for a HUMAN, not for
-      // this socket. `chat()` skips its terminal hook at an actionable-wait
-      // boundary, so its `finally` routes the disconnect here — and
-      // terminalizing then produced a record claiming the run finished while
+      // this socket. A disconnect can land after the interrupt boundary but
+      // before the invocation ends, and `chat()` then routes it here instead
+      // of to `onFinish` — and terminalizing then produced a record claiming the run finished while
       // the interrupt rows stayed `'pending'` and `validatePendingResumes`
       // still threw on the next request. An explicit cancel is different: the
       // user gave up on the approval, so the cancel band stays authoritative.

package/src/reconstruct.ts

@@ -1,8 +1,9 @@
-import { modelMessagesToUIMessages } from '@tanstack/ai'
+import { isTerminalRunStatus, modelMessagesToUIMessages } from '@tanstack/ai'
 import type {
   ModelMessage,
   RunRecord,
   SubagentPart,
+  TerminalRunStatus,
   UIMessage,
 } from '@tanstack/ai'
 import { storedSubagentInfo } from './subagent-runs'
@@ -41,11 +42,29 @@ export interface ReconstructedChat {
     pending: Array<Record<string, unknown>>
   } | null
   page?: { truncated: false } | { truncated: true; cursor: string }
+  /**
+   * The thread's finished runs, ascending by `startedAt`. Set only when
+   * {@link ReconstructChatOptions.includeRuns} is `true` and the `runs` store
+   * implements `listByThread`. Each assistant message of a listed run also
+   * gets the timings on `message.metadata.tanstack.run`.
+   */
+  runs?: Array<{
+    runId: string
+    status: TerminalRunStatus
+    startedAt: number
+    finishedAt?: number
+  }>
 }
 
 export interface ReconstructChatOptions {
   /** Query parameter carrying the thread id. Defaults to `threadId`. */
   param?: string
+  /**
+   * Add the thread's finished runs, with `startedAt` and `finishedAt`, to the
+   * response as `runs`. Needs a `runs` store that implements `listByThread`.
+   * Default: `false`.
+   */
+  includeRuns?: boolean
   /**
    * Authorize access to the requested thread before loading history.
    *
@@ -185,8 +204,26 @@ export async function reconstructChat(
     threadId,
     pending,
   )
+  const runStore = persistence.stores.runs
+  const runs =
+    options?.includeRuns && threadId && runStore?.listByThread
+      ? (await runStore.listByThread(threadId)).flatMap((run) =>
+          isTerminalRunStatus(run.status)
+            ? [
+                {
+                  runId: run.runId,
+                  status: run.status,
+                  startedAt: run.startedAt,
+                  ...(run.finishedAt !== undefined && {
+                    finishedAt: run.finishedAt,
+                  }),
+                },
+              ]
+            : [],
+        )
+      : undefined
   const body: ReconstructedChat = {
-    messages,
+    messages: runs ? stampRunTimings(messages, runs) : messages,
     activeRun: active ? { runId: active.runId } : null,
     interrupts: firstPending
       ? {
@@ -195,6 +232,7 @@ export async function reconstructChat(
         }
       : null,
     ...('page' in transcript ? { page: transcript.page } : {}),
+    ...(runs ? { runs } : {}),
   }
   return new Response(JSON.stringify(body), {
     headers: {
@@ -213,6 +251,40 @@ function messageRunId(message: UIMessage) {
   return typeof runId === 'string' && runId !== '' ? runId : undefined
 }
 
+/**
+ * Write each finished run's timings to `metadata.tanstack.run` on the
+ * assistant messages of that run, so a client reads them from the message.
+ */
+function stampRunTimings(
+  messages: Array<UIMessage>,
+  runs: NonNullable<ReconstructedChat['runs']>,
+): Array<UIMessage> {
+  const byId = new Map(runs.map((run) => [run.runId, run]))
+  return messages.map((message) => {
+    const tanstack = message.metadata?.tanstack
+    const runId: unknown = tanstack?.run?.id
+    const run =
+      message.role === 'assistant' && typeof runId === 'string'
+        ? byId.get(runId)
+        : undefined
+    if (!run) return message
+    return {
+      ...message,
+      metadata: {
+        ...message.metadata,
+        tanstack: {
+          ...tanstack,
+          run: {
+            id: run.runId,
+            startedAt: run.startedAt,
+            ...(run.finishedAt !== undefined && { finishedAt: run.finishedAt }),
+          },
+        },
+      },
+    }
+  })
+}
+
 type Runs = NonNullable<ChatTranscriptStores['runs']>
 
 /** Rebuild one child card from its run record and stored transcript. */

package/src/testkit/conformance.ts

@@ -62,6 +62,20 @@ type OptionalRunStoreMethod =
 /** Dotted `store.method` key a backend passes to declare an omitted method. */
 export type PersistenceConformanceMethodKey = `runs.${OptionalRunStoreMethod}`
 
+/**
+ * Checks added after the suite shipped. They are off by default, so a backend
+ * that passed before still passes. Turn them on with `options.checks`.
+ *
+ * - `'messages.metadata'`: `saveThread` / `loadThread` keep message `metadata`,
+ *   including `metadata.tanstack.run.id` (run timings on reload need it).
+ * - `'runs.listByThread.state'`: `listByThread` returns each run's current
+ *   `status` and `finishedAt` after `update` (`reconstructChat`'s
+ *   `includeRuns` needs it).
+ */
+export type PersistenceConformanceCheck =
+  | 'messages.metadata'
+  | 'runs.listByThread.state'
+
 /**
  * Unwrap a value the store contract says must be present. Fails the test with a
  * readable message instead of a non-null assertion (banned in this package) or
@@ -115,6 +129,12 @@ export interface PersistenceConformanceOptions {
    * has no effect.
    */
   skipMethods?: Array<PersistenceConformanceMethodKey>
+  /**
+   * Opt-in checks, off by default so existing backends keep passing. A check
+   * that is not listed is reported as a skipped case. See
+   * {@link PersistenceConformanceCheck}.
+   */
+  checks?: Array<PersistenceConformanceCheck>
 }
 
 /**
@@ -131,6 +151,7 @@ export function runPersistenceConformance(
   const skipMethods = new Set<PersistenceConformanceMethodKey>(
     options?.skipMethods ?? [],
   )
+  const checks = new Set<PersistenceConformanceCheck>(options?.checks ?? [])
 
   describe(`AIPersistence conformance: ${name}`, () => {
     let persistence: AIPersistence
@@ -274,6 +295,32 @@ export function runPersistenceConformance(
         await store.saveThread('thread-rich', rich)
         expect(await store.loadThread('thread-rich')).toEqual(rich)
       })
+
+      it('round-trips message metadata', async (ctx) => {
+        if (!checks.has('messages.metadata')) {
+          return ctx.skip(
+            "opt-in check: pass { checks: ['messages.metadata'] }",
+          )
+        }
+        const store = resolveStore('messages')
+        if (!store) return ctx.skip('store not provided')
+
+        const withMetadata: Array<ModelMessage> = [
+          {
+            role: 'user',
+            content: 'hi',
+            metadata: { author: { id: 'user-42' } },
+          },
+          {
+            role: 'assistant',
+            content: 'hello',
+            metadata: { tanstack: { run: { id: 'run-1' } }, custom: 1 },
+          },
+        ]
+
+        await store.saveThread('thread-metadata', withMetadata)
+        expect(await store.loadThread('thread-metadata')).toEqual(withMetadata)
+      })
     })
 
     describe('runs', () => {
@@ -520,6 +567,35 @@ export function runPersistenceConformance(
         expect(listed.map((r) => r.runId)).toEqual(['lt-a', 'lt-b'])
       })
 
+      it('lists runs by thread with their current status and finishedAt', async (ctx) => {
+        if (!checks.has('runs.listByThread.state')) {
+          return ctx.skip(
+            "opt-in check: pass { checks: ['runs.listByThread.state'] }",
+          )
+        }
+        const runs = resolveStore('runs')
+        if (!runs) return ctx.skip('store not provided')
+        if (!hasRunsMethod(runs, 'listByThread')) {
+          return ctx.skip('runs.listByThread not implemented')
+        }
+
+        await runs.createOrResume({
+          runId: 'lts-a',
+          threadId: 'lts',
+          startedAt: 1,
+        })
+        await runs.update('lts-a', { status: 'completed', finishedAt: 5 })
+
+        expect(await runs.listByThread('lts')).toEqual([
+          expect.objectContaining({
+            runId: 'lts-a',
+            status: 'completed',
+            startedAt: 1,
+            finishedAt: 5,
+          }),
+        ])
+      })
+
       // `listByParentRun` is optional and is skipped when absent. A store that
       // has it returns only that parent's children, oldest `startedAt` first,
       // and [] for an unknown parent.