fix(ai-persistence): persist chat run usage

.changeset/persist-chat-run-usage.md

@@ -0,0 +1,5 @@
+---
+'@tanstack/ai-persistence': patch
+---
+
+Persist cumulative usage for chat runs that make multiple model calls, interrupt, fail, or abort.

docs/persistence/chat-persistence.md

@@ -72,10 +72,10 @@ schema changes through your deployment workflow instead. See
 ## Threads, runs, and turns
 
 The transcript is stored per `threadId`, and each run gets a `runs` record with its
-status, timings and usage. One thing follows from that and matters when you wire a
-client: a reconnecting client never has to present a run id it may no longer know.
-The store resolves the thread's live run with `findActiveRun(threadId)` and the client
-tails that.
+status, timings, and reported usage across provider calls. One thing follows from
+that and matters when you wire a client: a reconnecting client never has to present
+a run id it may no longer know. The store resolves the thread's live run with
+`findActiveRun(threadId)` and the client tails that.
 
 [Id map](./id-map) covers how to choose a thread id and what both ids mean on the
 generation hooks. [How persistence works](./internals) has the rest.
@@ -98,8 +98,8 @@ generation hooks. [How persistence works](./internals) has the rest.
 | Moment | What is written | Best-effort? |
 | --- | --- | --- |
 | **Start of a run** (`onStart`) | Pending turn (just-submitted user message + prior history) so a reload mid-generation still shows the question | Yes. Failure does not abort the run; finish is authoritative |
-| **Interrupt boundary** | New interrupt records, run status `interrupted`, and a thread snapshot of current messages | No. Store failures propagate |
-| **Finish** (`onFinish`) | Complete transcript (including the terminal assistant reply with its stream `messageId` for in-place reload identity), run status `completed`, and commit of consumed resumes | No. The transcript is saved **before** the run is marked completed |
+| **Interrupt boundary** | New interrupt records, run status `interrupted`, known usage, and a thread snapshot of current messages | No. Store failures propagate |
+| **Finish** (`onFinish`) | Complete transcript (including the terminal assistant reply with its stream `messageId` for in-place reload identity), run status `completed`, known usage, and commit of consumed resumes | No. The transcript is saved **before** the run is marked completed |
 | **Optionally while streaming** | Throttled partial assistant text when `snapshotStreaming: true` | Yes |
 
 ```ts group=chat-persistence
@@ -114,9 +114,10 @@ with `snapshotIntervalMs` (default `1000`).
 
 On **error**, the run is marked `failed`. On **abort**, the run is marked
 `aborted` with a `finishedAt`; `interrupted` is written only at an interrupt
-boundary, and it is not terminal. Resumes accepted in `onConfig` are **not**
-consumed until a success boundary (interrupt or finish), so a failed run leaves
-pending interrupts retryable with the same resume batch.
+boundary, and it is not terminal. Both terminal paths retain usage reported
+before the failure or abort. Resumes accepted in `onConfig` are **not** consumed
+until a success boundary (interrupt or finish), so a failed run leaves pending
+interrupts retryable with the same resume batch.
 
 One abort does **not** terminalize: a plain client disconnect on a run that some
 other middleware has declared *detachable* (a durable event log plus a run
@@ -130,8 +131,11 @@ either one makes the abort terminal again. See
 [Takeover & Detached Runs](../sandbox/takeover#detach-vs-cancel).
 
 The lifecycle a run record moves through. `completed`, `failed`, and `aborted`
-are terminal; `interrupted` is **parked**, not terminal, and a continuation
-after one is a new run with a fresh `runId`:
+are terminal; `interrupted` is **parked**, not terminal. The normal client flow
+starts a continuation with a fresh `runId`. A server integration can reuse the
+same `runId`; `createOrResume` leaves its status `interrupted` until the next
+interrupt or terminal boundary. `findActiveRun` only returns `running` records,
+so it cannot discover a same-ID continuation while that continuation executes.
 
 ```mermaid
 stateDiagram-v2
@@ -144,7 +148,11 @@ stateDiagram-v2
     completed --> [*]
     failed --> [*]
     aborted --> [*]
-    interrupted --> [*] : continuation runs under a new runId
+    interrupted --> [*] : continuation may use a new runId
+    interrupted --> interrupted : same runId pauses again
+    interrupted --> completed : same runId completes
+    interrupted --> failed : same runId fails
+    interrupted --> aborted : same runId aborts
 ```
 
 ## Interrupts survive a restart

docs/persistence/internals.md

@@ -170,17 +170,22 @@ server event state, not the client's rendered messages.
 
 1. `setup` provides persistence, interrupt, and lock capabilities when their
    stores exist.
-2. `onConfig` creates or resumes the run, loads pending interrupts, and
-   validates the request's resume batch against them, then merges stored
-   messages into the request when the request carries no history.
-3. `onChunk` reacts only to a `RUN_FINISHED` interrupt outcome by committing
-   the accepted resumes, storing the new interrupts, marking the run
-   interrupted, and saving messages.
-4. `onFinish` and `onError` terminalize the run record. So does `onAbort`, with
-   one exception: on a run another middleware has declared detachable, a plain
-   disconnect (no cancel recorded in either band) writes nothing and leaves the
-   record `'running'` for a later takeover. See
-   [Takeover & Detached Runs](../sandbox/takeover#detach-vs-cancel).
+2. `onConfig` creates or resumes the run, seeds usage from the existing run
+   record, loads pending interrupts, and validates the request's resume batch
+   against them, then merges stored messages into the request when the request
+   carries no history.
+3. `onUsage` accumulates each provider terminal.
+4. `onChunk` reacts only to a `RUN_FINISHED` interrupt outcome. A direct adapter
+   terminal arrives before `onUsage`, so the handler includes its usage and
+   ignores the following `onUsage`. A synthesized tool boundary arrives after
+   the original terminal's `onUsage`, so the handler reuses that aggregate. It
+   then commits accepted resumes, stores the new interrupts, marks the run
+   interrupted, and saves messages.
+5. `onFinish` and `onError` terminalize the run record and retain known usage.
+   So does terminal `onAbort`, with one exception: on a run another middleware
+   has declared detachable, a plain disconnect (no cancel recorded in either
+   band) writes nothing and leaves the record `'running'` for a later takeover.
+   See [Takeover & Detached Runs](../sandbox/takeover#detach-vs-cancel).
 
 Accepted resumes are committed (interrupts marked resolved/cancelled) only once
 the run reaches a successful boundary, so a provider failure or abort between

docs/persistence/store-reference.md

@@ -68,7 +68,7 @@ interface RunRecord {
   startedAt: number // epoch ms
   finishedAt?: number // epoch ms, set once the run reaches a terminal status
   error?: RunError
-  usage?: TokenUsage // token counts, from @tanstack/ai
+  usage?: TokenUsage // reported usage accumulated for this runId
   // ---------------------------------------------------------------------------
   // DURABLE SANDBOXED RUNS ONLY. A chat app never writes these four and nothing
   // in `@tanstack/ai-persistence` reads them. Leave the columns out until you
@@ -128,12 +128,18 @@ interface RunStore {
 }
 ```
 
+`withPersistence` sums reported numeric usage fields across provider calls for
+the same `runId`. The opaque `providerUsageDetails` field retains the latest
+reported bag. Known usage is persisted when the run interrupts or reaches a
+terminal status.
+
 `createOrResume`, `update`, `get`, and `findActiveRun` are the floor: a backend
 that implements those four is a valid `RunStore`. Three contracts to hold:
 
 - `createOrResume` must be idempotent. A second call for an existing `runId`
-  returns the stored record unchanged, which is what makes resuming a run safe.
-  Retries may repeat the same run id.
+  returns the complete stored record unchanged, including `usage`. This makes
+  resuming a run safe and lets usage continue accumulating. Retries may repeat
+  the same run id.
 - `update` against an unknown `runId` is a no-op.
 - `findActiveRun` must do real work. Stub it to `null` and `reconstructChat`
   always reports `activeRun: null`, so a client that reloads (or switches back

packages/ai-persistence/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.

packages/ai-persistence/src/middleware.ts

@@ -254,6 +254,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 +1289,82 @@ 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 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 } : {}),
+    ...(cost !== undefined ? { cost } : {}),
+    ...(costDetails ? { costDetails } : {}),
+    ...(providerUsageDetails ? { providerUsageDetails } : {}),
+  }
 }
 
 async function completeRun(
@@ -1311,6 +1383,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 +1393,7 @@ async function failRun(
     status: 'failed',
     finishedAt: Date.now(),
     error: { message: error instanceof Error ? error.message : String(error) },
+    ...(usage ? { usage } : {}),
   })
 }
 
@@ -1334,9 +1408,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 +1424,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 +1586,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
@@ -1629,11 +1707,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
@@ -1650,12 +1741,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) {
@@ -1678,7 +1769,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.

packages/ai-persistence/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.