@tanstack/ai-persistence 0.1.4 → 0.2.0

package/skills/ai-persistence/stores/SKILL.md

@@ -172,6 +172,11 @@ predicate: `(status: RunStatus) => status is TerminalRunStatus`, so calling it
 inside a guard narrows `status` to `TerminalRunStatus` for the rest of that
 branch, with no cast needed.
 
+`RunRecord.usage` is optional. `withPersistence` sums reported numeric fields
+across provider calls for that `runId`, while opaque `providerUsageDetails`
+retains the latest reported bag. Known usage is persisted on interruption and
+every terminal status.
+
 `RunRecord.error` is a structured `RunError`, not a bare string:
 
 ```ts
@@ -253,10 +258,10 @@ store through `update`/`get` — but `cancelRequested` must round-trip
 faithfully (previous section) for the durable path to work at all.
 
 - **`createOrResume`** (required): if `runId` exists, return it **unchanged**,
-  ignoring the passed `threadId` / `startedAt` / `status`. Resuming a run does
-  not reset `startedAt` or overwrite its current status. Idempotent retries and
-  double-submit depend on this. `status` defaults to `'running'` on first
-  creation.
+  including its stored `usage`, and ignore the passed `threadId` / `startedAt` /
+  `status`. Resuming a run does not reset `startedAt` or overwrite its current
+  status. Idempotent retries and double-submit depend on this. `status` defaults
+  to `'running'` on first creation.
 - **`update`** (required): missing `runId` is a **no-op** (do not throw, do not
   insert).
 - **`get`** (required): current record, or `null` when unknown.

package/src/middleware.ts

@@ -35,6 +35,7 @@ import type {
   RunAgentResumeItem,
   StreamChunk,
   ToolApprovalResolution,
+  BilledUsage,
   TokenUsage,
 } from '@tanstack/ai'
 import type {
@@ -254,6 +255,8 @@ interface RunStateEntry {
     pending: Array<InterruptRecord>
     resumeByInterruptId: Map<string, RunAgentResumeItem>
   }
+  /** Usage accumulated across every model call in this chat invocation. */
+  usage?: TokenUsage
   /** Accumulated terminal-turn text, for throttled streaming snapshots (B). */
   streamingText?: string
   /** Epoch ms of the last streaming snapshot, to throttle writes (B). */
@@ -1287,12 +1290,98 @@ async function createOrResumeRun(
   runs: RunStore | undefined,
   runId: string,
   threadId: string,
-): Promise<void> {
-  await runs?.createOrResume({
+): Promise<TokenUsage | undefined> {
+  const run = await runs?.createOrResume({
     runId,
     threadId,
     startedAt: Date.now(),
   })
+  return run?.usage
+}
+
+function sumOptionalNumber(
+  current: number | undefined,
+  next: number | undefined,
+): number | undefined {
+  if (current === undefined) return next
+  if (next === undefined) return current
+  return current + next
+}
+
+function sumNumberFields<T extends object>(
+  current: T | undefined,
+  next: T | undefined,
+): T | undefined {
+  if (!current) return next
+  if (!next) return current
+
+  const result = { ...current }
+  for (const key of Object.keys(next) as Array<keyof T>) {
+    const currentValue = current[key]
+    const nextValue = next[key]
+    if (typeof nextValue === 'number') {
+      result[key] = ((typeof currentValue === 'number' ? currentValue : 0) +
+        nextValue) as T[keyof T]
+    }
+  }
+  return result
+}
+
+function accumulateTokenUsage(
+  current: TokenUsage | undefined,
+  next: TokenUsage,
+): TokenUsage {
+  if (!current) return { ...next }
+
+  const promptTokensDetails = sumNumberFields(
+    current.promptTokensDetails,
+    next.promptTokensDetails,
+  )
+  const completionTokensDetails = sumNumberFields(
+    current.completionTokensDetails,
+    next.completionTokensDetails,
+  )
+  const costDetails = sumNumberFields(current.costDetails, next.costDetails)
+  // Provider-specific details are opaque, so retain the latest reported bag.
+  const providerUsageDetails =
+    next.providerUsageDetails ?? current.providerUsageDetails
+  const durationSeconds = sumOptionalNumber(
+    current.durationSeconds,
+    next.durationSeconds,
+  )
+  const unitsBilled = sumOptionalNumber(current.unitsBilled, next.unitsBilled)
+  const billed = accumulateBilled(current.billed, next.billed)
+  const cost = sumOptionalNumber(current.cost, next.cost)
+
+  return {
+    ...current,
+    ...next,
+    promptTokens: current.promptTokens + next.promptTokens,
+    completionTokens: current.completionTokens + next.completionTokens,
+    totalTokens: current.totalTokens + next.totalTokens,
+    ...(promptTokensDetails ? { promptTokensDetails } : {}),
+    ...(completionTokensDetails ? { completionTokensDetails } : {}),
+    ...(durationSeconds !== undefined ? { durationSeconds } : {}),
+    ...(unitsBilled !== undefined ? { unitsBilled } : {}),
+    ...(billed !== undefined ? { billed } : {}),
+    ...(cost !== undefined ? { cost } : {}),
+    ...(costDetails ? { costDetails } : {}),
+    ...(providerUsageDetails ? { providerUsageDetails } : {}),
+  }
+}
+
+/**
+ * Sum billed quantities when both reports use the same unit. Different units
+ * cannot be added, so the later report wins.
+ */
+function accumulateBilled(
+  current: BilledUsage | undefined,
+  next: BilledUsage | undefined,
+): BilledUsage | undefined {
+  if (!current) return next
+  if (!next) return current
+  if (current.unit !== next.unit) return next
+  return { quantity: current.quantity + next.quantity, unit: current.unit }
 }
 
 async function completeRun(
@@ -1311,6 +1400,7 @@ async function failRun(
   runs: RunStore | undefined,
   runId: string,
   error: unknown,
+  usage?: TokenUsage,
 ): Promise<void> {
   // `RunRecord.error` is a structured `RunError`. Only `message` is filled in
   // here: the middleware sees an opaque thrown value, and inventing a `code`
@@ -1320,6 +1410,7 @@ async function failRun(
     status: 'failed',
     finishedAt: Date.now(),
     error: { message: error instanceof Error ? error.message : String(error) },
+    ...(usage ? { usage } : {}),
   })
 }
 
@@ -1334,9 +1425,11 @@ async function failRun(
 export async function interruptRun(
   runs: RunStore | undefined,
   runId: string,
+  usage?: TokenUsage,
 ): Promise<void> {
   await runs?.update(runId, {
     status: 'interrupted',
+    ...(usage ? { usage } : {}),
   })
 }
 
@@ -1348,10 +1441,12 @@ export async function interruptRun(
 export async function abortRun(
   runs: RunStore | undefined,
   runId: string,
+  usage?: TokenUsage,
 ): Promise<void> {
   await runs?.update(runId, {
     status: 'aborted',
     finishedAt: Date.now(),
+    ...(usage ? { usage } : {}),
   })
 }
 
@@ -1508,15 +1603,15 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
         }
       }
 
-      await createOrResumeRun(runs, ctx.runId, ctx.threadId)
+      const storedUsage = await createOrResumeRun(runs, ctx.runId, ctx.threadId)
 
-      {
-        const state = runState.get(ctx)
-        if (!state?.merged) {
-          if (state) state.merged = true
-          const stored = await messageStore.loadThread(ctx.threadId)
-          patch.messages = config.messages.length > 0 ? config.messages : stored
-        }
+      const state = runState.get(ctx)
+      // A continuation has a fresh middleware context but resumes the same run.
+      if (state && storedUsage) state.usage = storedUsage
+      if (!state?.merged) {
+        if (state) state.merged = true
+        const stored = await messageStore.loadThread(ctx.threadId)
+        patch.messages = config.messages.length > 0 ? config.messages : stored
       }
 
       return Object.keys(patch).length > 0 ? patch : undefined
@@ -1542,7 +1637,14 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       if (ctx.phase === 'modelStream') {
         const s = runState.get(ctx)
         if (s && chunk.type === 'TEXT_MESSAGE_START') {
-          s.streamingMessageId = chunk.messageId
+          // An empty/malformed messageId means "no identity" (matching the
+          // engine's convention), leaving room for the TOOL_CALL_START
+          // parentMessageId fallback below — but the per-turn accumulator
+          // still resets so snapshots never mix text across turns.
+          s.streamingMessageId =
+            typeof chunk.messageId === 'string' && chunk.messageId !== ''
+              ? chunk.messageId
+              : undefined
           s.streamingMessageCreatedAt = new Date()
           s.streamingText = ''
         } else if (
@@ -1622,11 +1724,24 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
           })
         }
       }
-      await interruptRun(runs, ctx.runId)
+      // Adapter terminals arrive before `onUsage`; synthesized tool boundaries
+      // arrive after it with the same usage already in state.
+      const usage =
+        ctx.phase === 'modelStream' && chunk.usage
+          ? accumulateTokenUsage(state.usage, chunk.usage)
+          : (state.usage ?? chunk.usage)
+      state.usage = usage
+      await interruptRun(runs, ctx.runId, usage)
       await messageStore.saveThread(ctx.threadId, [...ctx.messages])
       state.interrupted = true
     },
 
+    onUsage(ctx: ChatMiddlewareContext, usage: TokenUsage) {
+      const state = runState.get(ctx)
+      if (!state || state.interrupted) return
+      state.usage = accumulateTokenUsage(state.usage, usage)
+    },
+
     async onFinish(ctx: ChatMiddlewareContext, info: FinishInfo) {
       const state = runState.get(ctx)
       if (state?.interrupted) return
@@ -1643,12 +1758,12 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
           state?.streamingMessageCreatedAt,
         ),
       )
-      await completeRun(runs, ctx.runId, info.usage)
+      await completeRun(runs, ctx.runId, state?.usage ?? info.usage)
       await commitPendingResumes(state, persistence.stores.interrupts)
     },
 
     async onError(ctx: ChatMiddlewareContext, info: ErrorInfo) {
-      await failRun(runs, ctx.runId, info.error)
+      await failRun(runs, ctx.runId, info.error, runState.get(ctx)?.usage)
     },
 
     async onAbort(ctx: ChatMiddlewareContext, info: AbortInfo) {
@@ -1671,7 +1786,7 @@ export function withPersistence<TStores extends ChatTranscriptStores>(
       // user gave up on the approval, so the cancel band stays authoritative.
       const state = runState.get(ctx)
       if (cancelled || (!detachableRun(ctx) && state?.interrupted !== true)) {
-        await abortRun(runs, ctx.runId)
+        await abortRun(runs, ctx.runId, state?.usage)
         return
       }
       // A plain disconnect on a detachable or interrupted run: write NOTHING.

package/src/testkit/conformance.ts

@@ -299,6 +299,13 @@ export function runPersistenceConformance(
           usage: { promptTokens: 3, completionTokens: 4, totalTokens: 7 },
         })
 
+        const resumedAfterUpdate = await store.createOrResume({
+          runId: 'run-1',
+          threadId: 'thread-different',
+          startedAt: 9999,
+        })
+        expect(resumedAfterUpdate).toEqual(done)
+
         // `error` is a structured RunError: the prose `message` plus the
         // optional machine-branchable `code`. Both must survive the round-trip,
         // so a backend that flattens the record to a bare string fails here.