Skip to content

fix(ai-persistence): persist chat run usage - #1129

Merged
AlemTuzlak merged 3 commits into
TanStack:mainfrom
kolaworld:fix/1128-persist-chat-run-usage
Aug 18, 2026
Merged

fix(ai-persistence): persist chat run usage#1129
AlemTuzlak merged 3 commits into
TanStack:mainfrom
kolaworld:fix/1128-persist-chat-run-usage

Conversation

@kolaworld

@kolaworld kolaworld commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Fixes #1128

🎯 Changes

  • Accumulates usage across every provider call in a chat invocation.
  • Persists known usage when a run completes, interrupts, fails, or aborts.
  • Preserves accumulated usage across same-run continuations.
  • Sums numeric usage fields while retaining the latest opaque providerUsageDetails.
  • Continues omitting usage when the adapter reports none.
  • Updates store conformance, persistence tests, provider-free E2E coverage, and documentation.

✅ Checklist

  • I have followed the steps in the Contributing guide.
  • I have tested this code locally with pnpm run test:pr.

🚀 Release Impact

  • This change affects published code, and I have generated a changeset.
  • This change is docs/CI/dev-only (no release).

Summary by CodeRabbit

  • New Features

    • Chat runs now retain and accumulate token usage across multiple model calls, tool calls, interruptions, failures, aborts, and resumed executions.
    • Persisted usage includes prompt, completion, and total token counts, while preserving the latest provider-specific details.
    • Interrupted runs can resume with their previously recorded usage.
  • Bug Fixes

    • Usage is no longer lost when runs are cancelled or fail during structured-output processing.
  • Documentation

    • Expanded guidance on usage reporting, run lifecycle behavior, continuation, and persistence semantics.

@coderabbitai

coderabbitai Bot commented Aug 17, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: ea2be4b6-ad2f-4ffa-8586-1a5b27c7614a

📥 Commits

Reviewing files that changed from the base of the PR and between 2c6d142 and 9fa0e0b.

📒 Files selected for processing (2)
  • packages/ai-persistence/src/middleware.ts
  • packages/ai-persistence/tests/with-persistence.test.ts
🚧 Files skipped from review as they are similar to previous changes (2)
  • packages/ai-persistence/tests/with-persistence.test.ts
  • packages/ai-persistence/src/middleware.ts

Included review availability: Your plan includes up to 8 reviews per rolling hour; 6 remain after this review.


📝 Walkthrough

Walkthrough

The persistence middleware now aggregates token usage across provider calls and preserves it through interrupted, completed, failed, and aborted runs. Tests, an end-to-end scenario, documentation, and a patch changeset cover the behavior.

Changes

Usage persistence

Layer / File(s) Summary
Middleware usage lifecycle
packages/ai-persistence/src/middleware.ts
Run state restores stored usage, accumulates numeric and provider-specific usage, and persists it across lifecycle transitions.
Persistence behavior validation
packages/ai-persistence/src/testkit/conformance.ts, packages/ai-persistence/tests/*
Tests verify cumulative usage, idempotent resumption, and retention across interrupted, failed, aborted, tool-call, and structured-output runs.
End-to-end usage scenario
testing/e2e/src/routes/api.persistence-durability.ts, testing/e2e/tests/persistence-durability.spec.ts
The durability route performs two persisted calls and validates cumulative token usage through an end-to-end test.
Persistence documentation and release metadata
docs/persistence/*, packages/ai-persistence/skills/ai-persistence/stores/SKILL.md, .changeset/persist-chat-run-usage.md
Documentation describes usage aggregation, lifecycle persistence, continuation behavior, and unchanged run restoration. A patch changeset records the package release.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Merge Risk: ⚪ Minimal · up to 9fa0e

This change persists accumulated chat-run usage across completion, interruption, failure, abort, and continuation paths; no actionable merge-blocking risk remains beyond normal checks and review.

Sequence Diagram(s)

sequenceDiagram
  participant ChatEngine
  participant withPersistence
  participant RunRecord
  ChatEngine->>withPersistence: emit usage for each provider call
  withPersistence->>withPersistence: accumulate usage
  withPersistence->>RunRecord: persist aggregate at interruption or terminal state
  RunRecord-->>withPersistence: restore usage for continuation
Loading
🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly and concisely describes the primary change: persisting chat run usage in ai-persistence.
Description check ✅ Passed The description explains the changes, confirms both checklist items, and identifies the published-code release impact and changeset.
Linked Issues check ✅ Passed The implementation addresses issue #1128 by aggregating and persisting usage across calls, interruptions, continuations, failures, and aborts.
Out of Scope Changes check ✅ Passed The code, tests, end-to-end coverage, documentation, and changeset changes directly support the linked issue and pull request objectives.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@kolaworld kolaworld changed the title Fix/1128 persist chat run usage fix(ai-persistence): persist chat run usage Aug 17, 2026
@github-actions github-actions Bot added the waiting-on: maintainer The ball is in the maintainers’ court label Aug 18, 2026
@nx-cloud

nx-cloud Bot commented Aug 18, 2026

Copy link
Copy Markdown

View your CI Pipeline Execution ↗ for commit 9fa0e0b

Command Status Duration Result
nx run-many --targets=build --exclude=examples/... ✅ Succeeded 6s View ↗

☁️ Nx Cloud last updated this comment at 2026-08-18 14:19:34 UTC

@pkg-pr-new

pkg-pr-new Bot commented Aug 18, 2026

Copy link
Copy Markdown

Open in StackBlitz

@tanstack/ai

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai@1129

@tanstack/ai-acp

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-acp@1129

@tanstack/ai-angular

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-angular@1129

@tanstack/ai-anthropic

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-anthropic@1129

@tanstack/ai-bedrock

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-bedrock@1129

@tanstack/ai-byteplus

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-byteplus@1129

@tanstack/ai-claude-code

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-claude-code@1129

@tanstack/ai-client

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-client@1129

@tanstack/ai-code-mode

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-code-mode@1129

@tanstack/ai-code-mode-snippets

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-code-mode-snippets@1129

@tanstack/ai-codex

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-codex@1129

@tanstack/ai-cohere

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-cohere@1129

@tanstack/ai-devtools-core

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-devtools-core@1129

@tanstack/ai-durable-stream

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-durable-stream@1129

@tanstack/ai-elevenlabs

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-elevenlabs@1129

@tanstack/ai-event-client

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-event-client@1129

@tanstack/ai-fal

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-fal@1129

@tanstack/ai-gemini

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-gemini@1129

@tanstack/ai-grok

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-grok@1129

@tanstack/ai-grok-build

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-grok-build@1129

@tanstack/ai-groq

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-groq@1129

@tanstack/ai-isolate-cloudflare

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-cloudflare@1129

@tanstack/ai-isolate-daytona

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-daytona@1129

@tanstack/ai-isolate-node

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-node@1129

@tanstack/ai-isolate-quickjs

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-quickjs@1129

@tanstack/ai-isolate-quickjs-bun

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-isolate-quickjs-bun@1129

@tanstack/ai-mcp

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-mcp@1129

@tanstack/ai-memory

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-memory@1129

@tanstack/ai-mistral

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-mistral@1129

@tanstack/ai-ollama

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-ollama@1129

@tanstack/ai-openai

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-openai@1129

@tanstack/ai-opencode

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-opencode@1129

@tanstack/ai-openrouter

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-openrouter@1129

@tanstack/ai-perplexity

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-perplexity@1129

@tanstack/ai-persistence

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-persistence@1129

@tanstack/ai-preact

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-preact@1129

@tanstack/ai-react

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-react@1129

@tanstack/ai-react-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-react-ui@1129

@tanstack/ai-sandbox

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox@1129

@tanstack/ai-sandbox-cloudflare

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-cloudflare@1129

@tanstack/ai-sandbox-daytona

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-daytona@1129

@tanstack/ai-sandbox-docker

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-docker@1129

@tanstack/ai-sandbox-local-process

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-local-process@1129

@tanstack/ai-sandbox-sprites

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-sprites@1129

@tanstack/ai-sandbox-vercel

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-sandbox-vercel@1129

@tanstack/ai-solid

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-solid@1129

@tanstack/ai-solid-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-solid-ui@1129

@tanstack/ai-svelte

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-svelte@1129

@tanstack/ai-utils

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-utils@1129

@tanstack/ai-vercel-gateway

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-vercel-gateway@1129

@tanstack/ai-vue

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-vue@1129

@tanstack/ai-vue-ui

npm i https://pkg.pr.new/TanStack/ai/@tanstack/ai-vue-ui@1129

@tanstack/openai-base

npm i https://pkg.pr.new/TanStack/ai/@tanstack/openai-base@1129

@tanstack/preact-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/preact-ai-devtools@1129

@tanstack/react-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/react-ai-devtools@1129

@tanstack/solid-ai-devtools

npm i https://pkg.pr.new/TanStack/ai/@tanstack/solid-ai-devtools@1129

commit: 9fa0e0b

@AlemTuzlak
AlemTuzlak merged commit 302c9f6 into TanStack:main Aug 18, 2026
9 checks passed
@github-actions github-actions Bot mentioned this pull request Aug 18, 2026
AlemTuzlak added a commit that referenced this pull request Aug 18, 2026
Resolved 4 conflicts. Both sides were kept in every case:

- packages/ai/src/adapter-internals.ts
  Additive on both sides. Kept the interrupt exports and main's
  structured-output-text / structured-output-events exports.

- packages/ai-persistence/src/middleware.ts (onFinish)
  Kept this branch's try/catch and its saveThread -> commitPendingResumes
  -> completeRun order, so a failed batch write still calls failRun and
  leaves pending resumes retryable. Took main's accumulated usage from
  #1129: completeRun now gets `state?.usage ?? info.usage`, and failRun in
  the catch gets `state?.usage` so every terminal path persists usage the
  same way onError already does.

- packages/ai-persistence/tests/interrupts.test.ts
  Unioned the imports. Kept this branch's schema fixtures and the
  isInterruptFinished predicate, plus main's usage-aware
  interruptFinished(runId, usage) signature.

- docs/persistence/chat-persistence.md
  Kept this branch's atomic-commitBatch and legacy-sequential-fallback
  paragraphs and folded in main's sentence about terminal paths retaining
  usage.
@kolaworld
kolaworld deleted the fix/1128-persist-chat-run-usage branch August 22, 2026 19:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

waiting-on: maintainer The ball is in the maintainers’ court

Projects

None yet

Development

Successfully merging this pull request may close these issues.

withPersistence drops usage from interrupted and multi-call chat runs

2 participants