diff --git a/.changeset/persist-chat-run-usage.md b/.changeset/persist-chat-run-usage.md new file mode 100644 index 000000000..4fa64a17b --- /dev/null +++ b/.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. diff --git a/docs/persistence/chat-persistence.md b/docs/persistence/chat-persistence.md index e36638276..97a796575 100644 --- a/docs/persistence/chat-persistence.md +++ b/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 diff --git a/docs/persistence/internals.md b/docs/persistence/internals.md index 67daea2f3..777d61629 100644 --- a/docs/persistence/internals.md +++ b/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 diff --git a/docs/persistence/store-reference.md b/docs/persistence/store-reference.md index fd0633f42..83e99a49e 100644 --- a/docs/persistence/store-reference.md +++ b/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 diff --git a/packages/ai-persistence/skills/ai-persistence/stores/SKILL.md b/packages/ai-persistence/skills/ai-persistence/stores/SKILL.md index 17377e394..3989bff09 100644 --- a/packages/ai-persistence/skills/ai-persistence/stores/SKILL.md +++ b/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. diff --git a/packages/ai-persistence/src/middleware.ts b/packages/ai-persistence/src/middleware.ts index 91cabab81..25a0473a1 100644 --- a/packages/ai-persistence/src/middleware.ts +++ b/packages/ai-persistence/src/middleware.ts @@ -254,6 +254,8 @@ interface RunStateEntry { pending: Array resumeByInterruptId: Map } + /** 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 { - await runs?.createOrResume({ +): Promise { + 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( + 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) { + 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 { // `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 { 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 { await runs?.update(runId, { status: 'aborted', finishedAt: Date.now(), + ...(usage ? { usage } : {}), }) } @@ -1508,15 +1586,15 @@ export function withPersistence( } } - 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( }) } } - 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( 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( // 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. diff --git a/packages/ai-persistence/src/testkit/conformance.ts b/packages/ai-persistence/src/testkit/conformance.ts index eaa941ebb..018b4d8e5 100644 --- a/packages/ai-persistence/src/testkit/conformance.ts +++ b/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. diff --git a/packages/ai-persistence/tests/abort-status.test.ts b/packages/ai-persistence/tests/abort-status.test.ts index f23beb1ba..85f06d4d8 100644 --- a/packages/ai-persistence/tests/abort-status.test.ts +++ b/packages/ai-persistence/tests/abort-status.test.ts @@ -163,6 +163,84 @@ describe('chat onAbort status', () => { expect(run?.status).toBe('aborted') expect(run?.finishedAt).toBeTypeOf('number') }) + + it('preserves known usage when the run is cancelled during a tool call', async () => { + const persistence = memoryPersistence() + const controller = new AbortController() + const usage = { + promptTokens: 9, + completionTokens: 2, + totalTokens: 11, + } + const adapter = { + kind: 'text', + name: 'mock', + model: 'test-model', + '~types': {}, + chatStream: () => + (async function* () { + yield { + type: EventType.RUN_STARTED, + runId: 'usage-abort', + threadId: 't1', + timestamp: 1, + } satisfies StreamChunk + yield { + type: EventType.TOOL_CALL_START, + toolCallId: 'tool-1', + toolCallName: 'cancel', + toolName: 'cancel', + timestamp: 1, + } satisfies StreamChunk + yield { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'tool-1', + delta: '{}', + timestamp: 1, + } satisfies StreamChunk + yield { + type: EventType.RUN_FINISHED, + runId: 'usage-abort', + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, + usage, + } satisfies StreamChunk + })(), + structuredOutput: async () => ({ data: {}, rawText: '{}' }), + } as unknown as AnyTextAdapter + + const stream = chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + tools: [ + { + name: 'cancel', + description: 'Cancel', + execute: () => { + controller.abort(RUN_CANCEL_REASON) + return { cancelled: true } + }, + }, + ], + runId: 'usage-abort', + threadId: 't1', + abortController: controller, + middleware: [withPersistence(persistence)], + }) as AsyncIterable + try { + for await (const _ of stream) { + // drain + } + } catch { + // cancellation may reject the stream + } + + expect(await persistence.stores.runs!.get('usage-abort')).toMatchObject({ + status: 'aborted', + usage, + }) + }) }) describe('interrupt status shape', () => { diff --git a/packages/ai-persistence/tests/error-abort.test.ts b/packages/ai-persistence/tests/error-abort.test.ts index d028f0746..7c7074aa4 100644 --- a/packages/ai-persistence/tests/error-abort.test.ts +++ b/packages/ai-persistence/tests/error-abort.test.ts @@ -94,6 +94,69 @@ describe('chat persistence error/abort hooks', () => { expect(run?.error).toEqual({ message: 'provider exploded' }) }) + it('preserves known usage when structured-output finalization fails', async () => { + const persistence = memoryPersistence() + const usage = { + promptTokens: 12, + completionTokens: 4, + totalTokens: 16, + } + const { adapter } = mockAdapter([ + [ + runStarted(), + { + type: EventType.TEXT_MESSAGE_CONTENT, + messageId: 'm1', + delta: 'draft', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'stop', + timestamp: 1, + usage, + }, + ], + ]) + adapter.structuredOutput = () => + Promise.reject(new Error('finalization failed')) + + const chunks = await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + tools: [ + { + name: 'search', + description: 'Search', + execute: () => ({ hits: [] }), + }, + ], + outputSchema: { + type: 'object', + properties: { answer: { type: 'string' } }, + }, + stream: true, + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + expect(chunks).toContainEqual( + expect.objectContaining({ + type: EventType.RUN_ERROR, + message: 'finalization failed', + }), + ) + + expect(await persistence.stores.runs!.get('r1')).toMatchObject({ + status: 'failed', + usage, + }) + }) + it('coerces a non-Error thrown value into the run error message', async () => { const persistence = memoryPersistence() diff --git a/packages/ai-persistence/tests/interrupts.test.ts b/packages/ai-persistence/tests/interrupts.test.ts index ac46b31d0..f0e8794ac 100644 --- a/packages/ai-persistence/tests/interrupts.test.ts +++ b/packages/ai-persistence/tests/interrupts.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it, vi } from 'vitest' import { EventType, chat, defineChatMiddleware } from '@tanstack/ai' -import type { AnyTextAdapter, StreamChunk, Tool } from '@tanstack/ai' +import type { + AnyTextAdapter, + StreamChunk, + Tool, + TokenUsage, +} from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { withPersistence } from '../src/middleware' @@ -31,12 +36,13 @@ async function collect(stream: AsyncIterable) { return out } -const interruptFinished = (runId = 'r1'): StreamChunk => ({ +const interruptFinished = (runId = 'r1', usage?: TokenUsage): StreamChunk => ({ type: EventType.RUN_FINISHED, runId, threadId: 't1', finishReason: 'tool_calls', timestamp: 1, + ...(usage ? { usage } : {}), outcome: { type: 'interrupt', interrupts: [ @@ -80,27 +86,29 @@ const text = (delta: string): StreamChunk => ({ timestamp: 1, }) -const runFinished = (runId = 'r1'): StreamChunk => ({ +const runFinished = (runId = 'r1', usage?: TokenUsage): StreamChunk => ({ type: EventType.RUN_FINISHED, runId, threadId: 't1', finishReason: 'stop', timestamp: 1, + ...(usage ? { usage } : {}), }) -const toolCallFinished = (runId = 'r1'): StreamChunk => ({ +const toolCallFinished = (runId = 'r1', usage?: TokenUsage): StreamChunk => ({ type: EventType.RUN_FINISHED, runId, threadId: 't1', finishReason: 'tool_calls', timestamp: 1, + ...(usage ? { usage } : {}), }) -const toolCallChunks = () => [ +const toolCallChunks = (usage?: TokenUsage) => [ runStarted(), toolStart(), toolArgs(), - toolCallFinished(), + toolCallFinished('r1', usage), ] async function persistClientToolTurn( @@ -227,6 +235,69 @@ describe('interrupt persistence', () => { ) }) + it('persists client-tool interrupt usage once', async () => { + const persistence = memoryPersistence() + const usage = { + promptTokens: 8, + completionTokens: 3, + totalTokens: 11, + } + + await persistClientToolTurn( + persistence, + [approvalClientTool('clientSearch')], + toolCallChunks(usage), + ) + + expect(await persistence.stores.runs!.get('r1')).toMatchObject({ + status: 'interrupted', + usage, + }) + }) + + it('includes current usage from a direct interrupt after an earlier iteration', async () => { + const persistence = memoryPersistence() + const firstUsage = { + promptTokens: 8, + completionTokens: 3, + totalTokens: 11, + } + const interruptUsage = { + promptTokens: 13, + completionTokens: 5, + totalTokens: 18, + } + const { adapter } = mockAdapter([ + toolCallChunks(firstUsage), + [runStarted(), interruptFinished('r1', interruptUsage)], + ]) + + await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'hi' }], + tools: [ + { + ...clientTool('clientSearch'), + execute: () => ({ hits: [] }), + }, + ], + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect(await persistence.stores.runs!.get('r1')).toMatchObject({ + status: 'interrupted', + usage: { + promptTokens: 21, + completionTokens: 8, + totalTokens: 29, + }, + }) + }) + it('blocks normal new input while a thread has pending interrupts', async () => { const persistence = memoryPersistence() await persistence.stores.interrupts!.create({ @@ -255,7 +326,16 @@ describe('interrupt persistence', () => { it('treats resume entries as interrupt continuation on the same run', async () => { const persistence = memoryPersistence() - const first = mockAdapter([[runStarted(), interruptFinished()]]) + const first = mockAdapter([ + [ + runStarted(), + interruptFinished('r1', { + promptTokens: 8, + completionTokens: 3, + totalTokens: 11, + }), + ], + ]) await collect( chat({ adapter: first.adapter, @@ -270,7 +350,15 @@ describe('interrupt persistence', () => { ) const continuation = mockAdapter([ - [runStarted(), text('continued'), runFinished('r1')], + [ + runStarted(), + text('continued'), + runFinished('r1', { + promptTokens: 16, + completionTokens: 6, + totalTokens: 22, + }), + ], ]) const chunks = await collect( chat({ @@ -297,6 +385,11 @@ describe('interrupt persistence', () => { expect( (await persistence.stores.interrupts!.get('interrupt-1'))?.status, ).toBe('resolved') + expect((await persistence.stores.runs!.get('r1'))?.usage).toEqual({ + promptTokens: 24, + completionTokens: 9, + totalTokens: 33, + }) }) // The full two-phase chain for an approval-required client tool, driven diff --git a/packages/ai-persistence/tests/with-persistence.test.ts b/packages/ai-persistence/tests/with-persistence.test.ts index b70a6d751..852b38669 100644 --- a/packages/ai-persistence/tests/with-persistence.test.ts +++ b/packages/ai-persistence/tests/with-persistence.test.ts @@ -5,6 +5,7 @@ import type { ModelMessage, StreamChunk, Tool, + TokenUsage, } from '@tanstack/ai' import { memoryPersistence } from '../src/memory' import { withPersistence } from '../src/middleware' @@ -46,12 +47,17 @@ const ev = { delta, timestamp: 1, }), - runFinished: (runId = 'r1', threadId = 't1'): StreamChunk => ({ + runFinished: ( + runId = 'r1', + threadId = 't1', + usage?: TokenUsage, + ): StreamChunk => ({ type: EventType.RUN_FINISHED, runId, threadId, finishReason: 'stop', timestamp: 1, + ...(usage ? { usage } : {}), }), interrupted: (interruptId = 'interrupt-1'): StreamChunk => ({ type: EventType.RUN_FINISHED, @@ -136,6 +142,88 @@ describe('withPersistence (state-only)', () => { ]) }) + it('persists cumulative usage across model calls', async () => { + const persistence = memoryPersistence() + const { adapter } = mockAdapter([ + [ + ev.runStarted(), + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'agent-tool', + role: 'assistant', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_START, + toolCallId: 'call_1', + toolCallName: 'search', + toolName: 'search', + parentMessageId: 'agent-tool', + timestamp: 1, + }, + { + type: EventType.TOOL_CALL_ARGS, + toolCallId: 'call_1', + delta: '{}', + timestamp: 1, + }, + { + type: EventType.RUN_FINISHED, + runId: 'r1', + threadId: 't1', + finishReason: 'tool_calls', + timestamp: 1, + usage: { + promptTokens: 10, + completionTokens: 2, + totalTokens: 12, + promptTokensDetails: { cachedTokens: 3 }, + cost: 1, + }, + }, + ], + [ + ev.runStarted(), + { + type: EventType.TEXT_MESSAGE_START, + messageId: 'agent-final', + role: 'assistant', + timestamp: 1, + }, + ev.text('hello'), + ev.runFinished('r1', 't1', { + promptTokens: 20, + completionTokens: 4, + totalTokens: 24, + completionTokensDetails: { reasoningTokens: 2 }, + providerUsageDetails: { requestId: 'final' }, + cost: 2, + }), + ], + ]) + + await collect( + chat({ + adapter, + messages: [{ role: 'user', content: 'search' }], + tools: [serverSearchTool()], + runId: 'r1', + threadId: 't1', + middleware: [withPersistence(persistence)], + }) as AsyncIterable, + ) + + expect((await persistence.stores.runs!.get('r1'))?.usage).toEqual({ + promptTokens: 30, + completionTokens: 6, + totalTokens: 36, + promptTokensDetails: { cachedTokens: 3 }, + completionTokensDetails: { reasoningTokens: 2 }, + providerUsageDetails: { requestId: 'final' }, + cost: 3, + }) + }) + it('persists the pending user turn at start, so it survives a failed run', async () => { const persistence = memoryPersistence() const adapter = { @@ -372,7 +460,11 @@ describe('withPersistence (state-only)', () => { timestamp: 1, }, ev.text('hello'), - ev.runFinished(), + ev.runFinished('r1', 't1', { + promptTokens: 20, + completionTokens: 4, + totalTokens: 24, + }), ], ]) @@ -430,6 +522,11 @@ describe('withPersistence (state-only)', () => { threadId: 't1', finishReason: 'tool_calls', timestamp: 1, + usage: { + promptTokens: 10, + completionTokens: 2, + totalTokens: 12, + }, }, ], [ @@ -553,7 +650,7 @@ describe('withPersistence (state-only)', () => { } }) - it('does not let structured-output TEXT_MESSAGE_START replace the agent-loop id', async () => { + it('preserves the agent-loop message id and cumulative usage through structured output', async () => { const persistence = memoryPersistence() const { adapter } = mockAdapter([ [ @@ -584,6 +681,11 @@ describe('withPersistence (state-only)', () => { threadId: 't1', finishReason: 'tool_calls', timestamp: 1, + usage: { + promptTokens: 10, + completionTokens: 2, + totalTokens: 12, + }, }, ], [ @@ -595,12 +697,21 @@ describe('withPersistence (state-only)', () => { timestamp: 1, }, ev.text('hello'), - ev.runFinished(), + ev.runFinished('r1', 't1', { + promptTokens: 20, + completionTokens: 4, + totalTokens: 24, + }), ], ]) adapter.structuredOutput = async () => ({ data: { name: 'Ada' }, rawText: '{"name":"Ada"}', + usage: { + promptTokens: 5, + completionTokens: 1, + totalTokens: 6, + }, }) await collect( @@ -624,6 +735,11 @@ describe('withPersistence (state-only)', () => { (message) => message.role === 'assistant' && message.content === 'hello', ) expect(terminal?.id).toBe('agent-final') + expect((await persistence.stores.runs!.get('r1'))?.usage).toEqual({ + promptTokens: 35, + completionTokens: 7, + totalTokens: 42, + }) }) it('records an interrupt and marks the run interrupted', async () => { diff --git a/testing/e2e/src/routes/api.persistence-durability.ts b/testing/e2e/src/routes/api.persistence-durability.ts index 3f5cdeb55..fe5bb59a1 100644 --- a/testing/e2e/src/routes/api.persistence-durability.ts +++ b/testing/e2e/src/routes/api.persistence-durability.ts @@ -3,12 +3,19 @@ import { INTERRUPT_BINDING_METADATA_KEY, INTERRUPT_BINDING_VERSION, canonicalInterruptJson, + chat, digestInterruptJson, memoryStream, resumeServerSentEventsResponse, toServerSentEventsResponse, } from '@tanstack/ai' -import type { StreamChunk } from '@tanstack/ai' +import { memoryPersistence, withPersistence } from '@tanstack/ai-persistence' +import type { + AnyTextAdapter, + StreamChunk, + TokenUsage, + Tool, +} from '@tanstack/ai' /** * Provider-free harness route for the browser-refresh persistence story. It @@ -17,7 +24,7 @@ import type { StreamChunk } from '@tanstack/ai' * connection resumable — but streams a FIXED AG-UI sequence instead of calling * an LLM, so the e2e is deterministic with nothing to mock. * - * Three scenarios (`?scenario=`): + * Four scenarios (`?scenario=`): * * - `text` (default) — a run that streams one assistant text message and * finishes cleanly (`outcome: success`). The client persists the transcript @@ -33,6 +40,8 @@ import type { StreamChunk } from '@tanstack/ai' * GET below, which returns a `reconstructChat`-shaped JSON carrying a pending * interrupt. Proves a fresh client (empty `localStorage`) re-prompts the * approval from the server alone — the path that was previously broken. + * - `usage` — runs two provider calls through server persistence and returns + * their stored cumulative usage. * * Exempt from the aimock policy: this route streams a fixed AG-UI sequence and * never reaches an LLM provider's HTTP layer, so there is nothing to mock. @@ -46,7 +55,11 @@ const confirmSchema = { required: ['confirmed'], } -function textRun(threadId: string, runId: string): AsyncIterable { +function textRun( + threadId: string, + runId: string, + usage?: TokenUsage, +): AsyncIterable { return (async function* () { yield { type: 'RUN_STARTED', @@ -78,10 +91,58 @@ function textRun(threadId: string, runId: string): AsyncIterable { runId, timestamp: Date.now(), outcome: { type: 'success' }, + ...(usage ? { usage } : {}), } as StreamChunk })() } +const usagePersistence = memoryPersistence() +const usageTool: Tool = { + name: 'search', + description: 'Search', + execute: () => ({ hits: [] }), +} +const usageAdapter = { + kind: 'text', + name: 'fixed', + model: 'test-model', + '~types': {}, + chatStream: ({ threadId, runId }: { threadId: string; runId: string }) => + textRun(threadId, runId, { + promptTokens: 12, + completionTokens: 4, + totalTokens: 16, + }), + structuredOutput: () => + Promise.resolve({ + data: { name: 'Ada Lovelace' }, + rawText: '{"name":"Ada Lovelace"}', + usage: { + promptTokens: 7, + completionTokens: 2, + totalTokens: 9, + }, + }), +} as unknown as AnyTextAdapter + +async function cumulativeUsage(threadId: string, runId: string) { + const stream = chat({ + adapter: usageAdapter, + messages: [{ role: 'user', content: 'Name the programmer' }], + tools: [usageTool], + outputSchema: { + type: 'object', + properties: { name: { type: 'string' } }, + }, + stream: true, + threadId, + runId, + middleware: [withPersistence(usagePersistence)], + }) + for await (const _ of stream) void _ + return usagePersistence.stores.runs?.get(runId) +} + function interruptRun( threadId: string, runId: string, @@ -135,11 +196,12 @@ function stringField(body: unknown, key: string): string | undefined { function scenarioOf( request: Request, -): 'text' | 'interrupt' | 'server-interrupt' { +): 'text' | 'interrupt' | 'server-interrupt' | 'usage' { try { const value = new URL(request.url).searchParams.get('scenario') if (value === 'interrupt') return 'interrupt' if (value === 'server-interrupt') return 'server-interrupt' + if (value === 'usage') return 'usage' return 'text' } catch { return 'text' @@ -193,6 +255,10 @@ export const Route = createFileRoute('/api/persistence-durability')({ const body: unknown = await request.json() const threadId = stringField(body, 'threadId') ?? 'persistence-thread' const runId = stringField(body, 'runId') ?? crypto.randomUUID() + if (scenarioOf(request) === 'usage') { + const run = await cumulativeUsage(threadId, runId) + return Response.json({ runId, threadId, usage: run?.usage }) + } const stream = scenarioOf(request) === 'interrupt' ? interruptRun(threadId, runId) diff --git a/testing/e2e/tests/persistence-durability.spec.ts b/testing/e2e/tests/persistence-durability.spec.ts index 550debd0e..ac82d560e 100644 --- a/testing/e2e/tests/persistence-durability.spec.ts +++ b/testing/e2e/tests/persistence-durability.spec.ts @@ -140,3 +140,30 @@ test.describe('persistence durability (browser refresh)', () => { expect(stored).toBeNull() }) }) + +test.describe('server persistence', () => { + test('stores cumulative usage across model calls', async ({ request }) => { + const run = await request.post( + '/api/persistence-durability?scenario=usage', + { + data: { + threadId: `usage-${crypto.randomUUID()}`, + runId: crypto.randomUUID(), + }, + }, + ) + expect(run.ok()).toBe(true) + const body = (await run.json()) as { + usage?: { + promptTokens: number + completionTokens: number + totalTokens: number + } + } + expect(body.usage).toEqual({ + promptTokens: 19, + completionTokens: 6, + totalTokens: 25, + }) + }) +})