@tanstack/ai-persistence 0.6.5 → 0.7.1

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.