💾 feat: Persist Context Breakdown & Branch/Total Usage Cost (#13734)

* 💾 feat: Persist Context Breakdown & Branch/Total Usage Cost

Persist the granular context breakdown and per-response usage/cost on the
response message metadata, and re-derive branch + total usage/cost from a
per-message index so the popover survives reloads and is branch-aware live.

- Add aggregateEmittedUsage + buildPersistedContextUsage helpers in
  packages/api; capture the latest visible snapshot and every emitted
  on_token_usage payload via contextUsageSink/usageEmitSink.
- Attach metadata.contextUsage (Part A) and metadata.usage (Part B) on the
  agents response message in sendCompletion.
- Carry per-message usage on the token index; add sumTotalUsage/setEntryUsage
  and branch-scoped usage on sumBranch.
- Repurpose the session accumulator into a single in-flight pending holder;
  flush it into the index at finalize; hydrate breakdowns on load.
- Render branch cost with a conditional all-branches total in the breakdown.

* 🧹 chore: Remove orphaned com_ui_session_cost i18n key

* 🩹 fix: Address Codex review — normalize usage server-side, fix reload deltas

- Persist per-event-normalized display units in metadata.usage (TResponseUsage)
  so reloaded mixed-provider turns match the live session; client reads them
  directly instead of re-normalizing with a single stamped provider (P2).
- Persist completedOutputTokens (final call output) on metadata.contextUsage so
  a reloaded multi-call turn adds the post-snapshot delta, not the full
  tokenCount the snapshot already counts (P2).
- buildIndex preserves a prior entry's immutable usage when a rebuilt cache
  message lacks metadata.usage, so a mid-session rebuild (regenerate) keeps a
  sibling branch's flushed cost (fixes the e2e regenerate failure).
- Track costKnown so turns saved with contextCost off don't render $0.00 when
  cost display is later enabled (P3).
- Use an epsilon for the all-branches cost comparison to avoid a spurious total
  row from float summation order (P3).
- Update unit/integration/e2e tests for the new shapes; regenerate e2e asserts
  the all-branches total after reload (deterministic via persisted metadata).

* 🩹 fix: Address Codex round 2 — pending leak, cost coverage, reload delta

- Clear the in-flight pending usage on terminal abort/error (resetLive), so a
  stopped generation's tokens no longer merge into the next response (P2).
- costKnown now means COMPLETE coverage (ANDed): a branch mixing cost-bearing
  and cost-less turns is flagged incomplete and the cost row is hidden rather
  than rendering an under-reported total (P2).
- Drop the tokenCount fallback for completedOutputTokens on reload: only the
  persisted post-snapshot delta is used, so a multi-call turn whose provider
  emitted no usage_metadata no longer double-counts earlier output (P2).
- Update tokens.spec for AND coverage semantics + incomplete-cost case.

* 🩹 fix: Address Codex round 3 — no-usage snapshots, total coverage, provider-less cache

- Skip persisting metadata.contextUsage when the response emitted no primary
  usage event: without a known post-snapshot output the granular gauge would
  undercount the reply on reload, so fall back to the coarse per-message
  estimate instead (P2).
- Gate the all-branches cost row on totalUsage.costKnown so an incomplete total
  (a sibling saved without cost) never renders an under-reported figure (P2).
- aggregateEmittedUsage/finalCallOutputTokens now normalize per-event with the
  client's magnitude fallback (normalizeEventUnits) instead of billing
  splitUsage, so provider-less cached events match live on reload (P2).
- Add backend test for the provider-less cached case.

* 🩹 fix: Address Codex round 4 — abort attribution, complete cost coverage

- aggregateEmittedUsage persists cost only when EVERY call was priced; a partial
  pricing failure now omits cost so the client treats coverage as unknown rather
  than reading an under-reported sum as authoritative (P2).
- finalizeUsage flushes pending into the response entry only when events were
  folded this session (eventCount > 0), so a late/second resumable subscriber
  carrying persisted metadata.usage keeps it instead of being overwritten with
  an empty pending record (P2).
- On user stop, attribute the in-flight pending usage to the partial response
  (new attributePending handler) instead of discarding it in resetLive — the
  stopped reply's billed tokens are kept and still can't leak into the next
  response; resetLive's discard remains for the error path (P2).

* 🐛 fix: Persist branch cost across branch switches via sticky usage history

Branch cost vanished on switching to a sibling branch (until a new turn) — the
cost analog of the granularity bug. buildIndex rebuilds the token index from the
messages cache; a sibling generated this session whose cache message lacks
metadata.usage (and is transiently dropped from the cache during regenerate)
lost its live-flushed usage, so sumBranch found none and the cost row hid.

Fix: a sticky per-response usage map (conversationId → messageId → usage),
written by setEntryUsage and never rebuilt from the cache — the usage counterpart
of snapshotsByAnchorFamily for the breakdown. buildIndex/upsertEntries restore an
entry's usage from it when the message carries none; cleared on convo switch and
migrated with the index. Add unit coverage for the drop-then-readd regression and
an e2e assertion that branch cost survives a branch switch.

* 🐛 fix: Re-index on branch switch so branch cost survives the switch

The sticky usage history alone didn't fix the reported branch-switch cost drop:
on a branch switch no cache `updated` event fires, so the index subscriber never
re-ran, and the post-regenerate rebuild was skipped while `isSubmitting` was
still true — leaving the index stale and missing the now-viewed branch's
response entirely (sticky can only restore entries present in a rebuild).

Re-index from the messages cache on every tail change (created/finalize AND
branch switch), not just while submitting. The cache holds the full message set
at switch time, so the viewed branch's response is re-added and its usage
restored from metadata.usage or the sticky history → sumBranch finds it and the
branch cost renders. Verified locally: the branch-switch e2e now passes (the
cost section shows both the branch row and the all-branches total). Also fixed
that e2e assertion to target a single cost value (strict-mode safe).

* 🩹 fix: Handle stopped-stream usage — reset pending + persist abort metadata

Codex round (stop/abort edges):
- Resumable explicit-stop (intentional SSE close) reset UI state but never
  cleared pendingUsageFamily, so usage folded before the stop leaked into the
  next response in the conversation. Discard pending on intentional close
  (resetLive); a resume re-folds via backfillUsage, so nothing is lost.
- The abort save path (abortMiddleware) persisted the stopped response without
  metadata.usage/contextUsage, so its cost + breakdown vanished on reload.
  Rebuild both from the job's persisted tokenUsage (emitted payloads incl. cost)
  and contextUsage snapshot — parity with the normal sendCompletion path;
  breakdown gated on a primary usage event like buildResponseMetadata.

Deferred (per scope decision): mid-stream branch-switch transiently shows the
streaming branch's pending on the viewed sibling (cosmetic, until finalize).

* 🩹 fix: Persist abort metadata on the real agents route + tighten snapshot gate

Codex round (corrects last round's wrong-path fixes):
- Stopped AGENTS responses are saved by routes/agents/index.js (/chat/abort),
  not abortMiddleware — so last round's metadata fix never ran for them. Moved
  the rollup/snapshot builder into packages/api as buildAbortedResponseMetadata
  (shared, unit-tested) and applied it in BOTH abort save paths, so a stopped
  agent reply keeps its cost + breakdown on reload.
- Persist the breakdown only when the FINAL visible call emitted usage: track a
  per-response snapshot count and require primaryUsageCount >= snapshotCount.
  Previously any earlier primary usage event passed the gate, so a multi-call
  turn whose final call emitted no usage_metadata used an earlier call's output
  as completedOutputTokens (already counted by the latest snapshot) → reload
  over-reported. Now it falls back to the coarse estimate.

Resumable stop pending-reset (prior round, 3cde6fe035) already flows through
clearAllSubmissions → SSE close → the intentional-close handler's resetLive.
Deferred per scope: mid-stream branch-switch pending attribution (tracked).

* 🩹 fix: Abort breakdown over-count + resume re-fold after pending discard

Codex round (on the re-applied abort/snapshot work):
- buildAbortedResponseMetadata now persists ONLY the usage/cost rollup, not the
  context breakdown. The abort path can't tell whether the final call emitted
  usage (the job stores only the latest snapshot, not a count), so persisting
  the breakdown risked reusing an earlier call's output as completedOutputTokens
  (already in the snapshot) → reload over-count. Stopped/incomplete responses
  now fall back to the coarse gauge estimate, which is safe and apt.
- resetLive now also forgets the conversation's folded usage-event identities
  (clearUsageFolded). Discarding pending on a terminal/intentional close left
  the folded keys set, so a later resume's backfillUsage saw the persisted
  events as duplicates and never rebuilt pending — leaving the response's usage
  missing until a full reload. Clearing them lets the resume re-fold.
This commit is contained in:
Danny Avila 2026-06-14 10:48:07 -04:00 • committed by GitHub
parent 98704f28c1
commit b03b2a0a29
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 1321 additions and 90 deletions

View file

@ -1,7 +1,15 @@
import type { TContextUsageEvent, TTokenUsageEvent } from 'librechat-data-provider';
import type { RecordUsageDeps, RecordUsageParams, SubagentUsageEvent } from './usage';
import type { UsageMetadata } from '../stream/interfaces/IJobStore';
import type { BulkWriteDeps, PricingFns } from './transactions';
import { computeUsageCostUSD, createSubagentUsageSink, recordCollectedUsage } from './usage';
import {
computeUsageCostUSD,
aggregateEmittedUsage,
createSubagentUsageSink,
recordCollectedUsage,
buildPersistedContextUsage,
buildAbortedResponseMetadata,
} from './usage';
describe('recordCollectedUsage', () => {
let mockSpendTokens: jest.Mock;
@ -1503,3 +1511,190 @@ describe('computeUsageCostUSD', () => {
expect(cost).toBeCloseTo((1000 * 3 + 2000 * 3.75 + 10000 * 0.3 + 500 * 15) / 1e6);
});
});
describe('aggregateEmittedUsage', () => {
it('returns null for no emitted events', () => {
expect(aggregateEmittedUsage([])).toBeNull();
});
it('normalizes each call into display units (input excludes cache) and sums cost', () => {
const events: TTokenUsageEvent[] = [
{
input_tokens: 100,
output_tokens: 20,
total_tokens: 120,
model: 'gpt-4o-mini',
provider: 'openAI',
cost: 0.001,
},
{
input_tokens: 150,
output_tokens: 10,
total_tokens: 160,
input_token_details: { cache_creation: 30, cache_read: 50 },
model: 'gpt-4o-mini',
provider: 'openAI',
cost: 0.002,
},
];
/** openAI is cache-subset: input excludes cache (150−30−50=70) */
expect(aggregateEmittedUsage(events)).toEqual({
input: 170,
output: 30,
cacheWrite: 30,
cacheRead: 50,
cost: 0.003,
});
});
it('omits cost when no event carried it (contextCost off)', () => {
const rollup = aggregateEmittedUsage([
{ input_tokens: 100, output_tokens: 20, total_tokens: 120, provider: 'openAI' },
]);
expect(rollup).toEqual({ input: 100, output: 20, cacheWrite: 0, cacheRead: 0 });
expect(rollup?.cost).toBeUndefined();
});
it('omits cost when any call lacked it (partial pricing failure)', () => {
/** One call priced, one missing cost (e.g. computeUsageCostUSD threw) →
* the sum would under-report, so coverage is incomplete and cost omitted. */
const rollup = aggregateEmittedUsage([
{ input_tokens: 100, output_tokens: 20, provider: 'openAI', cost: 0.01 },
{ input_tokens: 50, output_tokens: 10, provider: 'openAI' },
]);
expect(rollup?.cost).toBeUndefined();
expect(rollup?.input).toBe(150);
expect(rollup?.output).toBe(30);
});
it('normalizes mixed-provider calls per their own provider before summing', () => {
/** anthropic is additive (cache separate from input), openAI is subset */
const rollup = aggregateEmittedUsage([
{
input_tokens: 100,
output_tokens: 20,
provider: 'anthropic',
input_token_details: { cache_read: 40 },
cost: 0.01,
},
{
input_tokens: 90,
output_tokens: 5,
usage_type: 'subagent',
provider: 'openAI',
input_token_details: { cache_read: 30 },
cost: 0.02,
},
]);
/** anthropic input stays 100 (additive); openAI input 90−30=60 → 160 */
expect(rollup?.input).toBe(160);
expect(rollup?.output).toBe(25);
expect(rollup?.cacheRead).toBe(70);
expect(rollup?.cost).toBeCloseTo(0.03);
});
it('uses the magnitude fallback for provider-less cached events (matches live)', () => {
/** No provider: the client's normalizeUsageUnits uses a magnitude heuristic
* (cache ≤ input ⇒ input includes cache), so the rollup must too — billing
* splitUsage would treat it as additive and leave input at 1000, diverging
* from the live display after reload. */
const rollup = aggregateEmittedUsage([
{ input_tokens: 1000, output_tokens: 100, input_token_details: { cache_read: 400 } },
]);
expect(rollup).toEqual({ input: 600, output: 100, cacheWrite: 0, cacheRead: 400 });
});
});
describe('buildPersistedContextUsage', () => {
const baseSnapshot: TContextUsageEvent = {
runId: 'run-1',
breakdown: {
maxContextTokens: 8000,
instructionTokens: 100,
systemMessageTokens: 80,
dynamicInstructionTokens: 20,
toolSchemaTokens: 30,
summaryTokens: 0,
toolCount: 2,
messageCount: 3,
messageTokens: 500,
availableForMessages: 7000,
toolTokenCounts: { add: 15, noop: 0 },
},
contextBudget: 7800,
};
it('trims zero-valued per-tool counts', () => {
const result = buildPersistedContextUsage(baseSnapshot);
expect(result.breakdown.toolTokenCounts).toEqual({ add: 15 });
expect(result.contextBudget).toBe(7800);
});
it('drops the tool counts object entirely when all are zero', () => {
const result = buildPersistedContextUsage({
...baseSnapshot,
breakdown: { ...baseSnapshot.breakdown, toolTokenCounts: { add: 0 } },
});
expect(result.breakdown.toolTokenCounts).toBeUndefined();
});
it('passes through a snapshot without tool counts', () => {
const { toolTokenCounts: _omit, ...breakdown } = baseSnapshot.breakdown;
const result = buildPersistedContextUsage({ ...baseSnapshot, breakdown });
expect(result.breakdown.toolTokenCounts).toBeUndefined();
expect(result.breakdown.messageTokens).toBe(500);
});
it('records the final primary call output as completedOutputTokens', () => {
/** The latest snapshot precedes the final call, so its post-snapshot delta
* is that call's output — not the full multi-call response tokenCount. */
const events: TTokenUsageEvent[] = [
{ input_tokens: 100, output_tokens: 40, total_tokens: 140, provider: 'openAI' },
{ input_tokens: 200, output_tokens: 25, total_tokens: 225, provider: 'openAI' },
{
input_tokens: 50,
output_tokens: 12,
usage_type: 'subagent',
provider: 'openAI',
},
];
const result = buildPersistedContextUsage(baseSnapshot, events);
/** Last PRIMARY call's completion (25), skipping the trailing subagent event */
expect(result.completedOutputTokens).toBe(25);
});
it('omits completedOutputTokens when there are no primary calls', () => {
expect(buildPersistedContextUsage(baseSnapshot, []).completedOutputTokens).toBeUndefined();
});
});
describe('buildAbortedResponseMetadata', () => {
it('returns undefined for an empty job', () => {
expect(buildAbortedResponseMetadata(undefined)).toBeUndefined();
expect(buildAbortedResponseMetadata({})).toBeUndefined();
expect(buildAbortedResponseMetadata({ tokenUsage: 'not json' })).toBeUndefined();
});
it('rebuilds the usage/cost rollup from the job’s persisted emitted usage', () => {
const events: TTokenUsageEvent[] = [
{ input_tokens: 100, output_tokens: 20, total_tokens: 120, provider: 'openAI', cost: 0.001 },
];
const result = buildAbortedResponseMetadata({ tokenUsage: JSON.stringify(events) });
expect(result?.usage).toEqual({
input: 100,
output: 20,
cacheWrite: 0,
cacheRead: 0,
cost: 0.001,
});
});
it('never persists a breakdown for a stopped response (avoids final-call over-count)', () => {
const events: TTokenUsageEvent[] = [
{ input_tokens: 100, output_tokens: 20, total_tokens: 120, provider: 'openAI' },
];
const result = buildAbortedResponseMetadata({ tokenUsage: JSON.stringify(events) });
expect(result).toEqual({ usage: { input: 100, output: 20, cacheWrite: 0, cacheRead: 0 } });
expect((result as { contextUsage?: unknown }).contextUsage).toBeUndefined();
});
});

View file

@ -1,6 +1,12 @@
import { logger } from '@librechat/data-schemas';
import { inputTokensIncludesCache } from 'librechat-data-provider';
import type { TCustomConfig, TTransactionsConfig } from 'librechat-data-provider';
import type {
TCustomConfig,
TResponseUsage,
TTokenUsageEvent,
TContextUsageEvent,
TTransactionsConfig,
} from 'librechat-data-provider';
import type {
StructuredTokenUsage,
BulkWriteDeps,
@ -157,6 +163,169 @@ export function computeUsageCostUSD(
return credits / 1e6;
}
/**
* Aggregates the per-model-call `on_token_usage` payloads emitted for one
* response into a single rollup, persisted on `responseMessage.metadata.usage`.
*
* Each event is normalized into display units with the SAME logic the live
* client uses (`splitUsage`: input excludes cache, output is repaired) BEFORE
* summing, so the rollup reproduces the live branch/total usage exactly even
* when a turn mixes providers (e.g. a summarization or subagent call on a
* different provider than the primary). `cost` is the additive sum of the
* authoritative per-event cost, included only when at least one event carried
* it (i.e. `interface.contextCost` was on).
*/
export function aggregateEmittedUsage(
events: ReadonlyArray<TTokenUsageEvent>,
): TResponseUsage | null {
if (events.length === 0) {
return null;
}
let input = 0;
let output = 0;
let cacheWrite = 0;
let cacheRead = 0;
let cost = 0;
/** Persist cost only with COMPLETE coverage — every call priced. A partial
* sum (e.g. one call's `computeUsageCostUSD` threw and emitted without cost)
* would read back as authoritative and under-report; omitting it makes the
* client treat coverage as unknown and hide the cost, matching the live fold.
* Naturally false when `contextCost` is off (no event carries cost). */
let allHaveCost = true;
for (const event of events) {
const units = normalizeEventUnits(event);
input += units.input;
output += units.output;
cacheWrite += units.cacheWrite;
cacheRead += units.cacheRead;
if (event.cost != null) {
cost += event.cost;
} else {
allHaveCost = false;
}
}
const rollup: TResponseUsage = { input, output, cacheWrite, cacheRead };
if (allHaveCost) {
rollup.cost = cost;
}
return rollup;
}
/**
* Per-event display-unit normalization, mirroring the client's
* `normalizeUsageUnits` EXACTLY — including the magnitude fallback when
* `provider` is absent — so a reloaded rollup matches what the live client
* folded. This is deliberately distinct from billing `splitUsage`, which treats
* a missing provider as additive (no magnitude fallback); the divergence only
* surfaces for provider-less cached events (e.g. some OpenAI-compatible/custom
* payloads), where the client subtracts cache from input but `splitUsage`
* would not. Keep in sync with `normalizeUsageUnits` in client/src/utils/tokens.ts.
*/
function normalizeEventUnits(event: TTokenUsageEvent): {
input: number;
output: number;
cacheWrite: number;
cacheRead: number;
} {
const rawInput = event.input_tokens ?? 0;
const rawOutput = event.output_tokens ?? 0;
const total = event.total_tokens ?? 0;
const cacheWrite = event.input_token_details?.cache_creation ?? 0;
const cacheRead = event.input_token_details?.cache_read ?? 0;
const includesCache =
event.provider != null
? inputTokensIncludesCache(event.provider)
: cacheWrite + cacheRead <= rawInput;
const cacheAdjustment = includesCache ? 0 : cacheRead + cacheWrite;
const output =
total > rawInput + rawOutput + cacheAdjustment ? total - rawInput - cacheAdjustment : rawOutput;
return {
input: includesCache ? Math.max(0, rawInput - cacheRead - cacheWrite) : rawInput,
output,
cacheWrite,
cacheRead,
};
}
/** Output tokens of the response's final primary model call — the call the
* latest pre-invoke snapshot precedes. Persisted as the snapshot's
* `completedOutputTokens` so a reloaded multi-call turn adds only this delta
* (matching the live finalizer) instead of the full response `tokenCount`,
* which the snapshot already counts for earlier steps. */
function finalCallOutputTokens(events: ReadonlyArray<TTokenUsageEvent>): number {
for (let i = events.length - 1; i >= 0; i--) {
if (events[i].usage_type == null) {
return normalizeEventUnits(events[i]).output;
}
}
return 0;
}
/**
* Projects the latest live context snapshot into the blob persisted on
* `responseMessage.metadata.contextUsage`. Trims zero-valued per-tool counts
* (privacy/size) and records the final call's output as `completedOutputTokens`
* so rehydration adds the same post-snapshot delta the live gauge did. The
* client re-anchors the blob to the response message id on load.
*/
export function buildPersistedContextUsage(
snapshot: TContextUsageEvent,
usageEvents: ReadonlyArray<TTokenUsageEvent> = [],
): TContextUsageEvent {
const { breakdown } = snapshot;
const completedOutputTokens = finalCallOutputTokens(usageEvents);
let toolTokenCounts = breakdown.toolTokenCounts;
if (toolTokenCounts != null) {
const trimmed: Record<string, number> = {};
for (const [name, count] of Object.entries(toolTokenCounts)) {
if (count > 0) {
trimmed[name] = count;
}
}
toolTokenCounts = Object.keys(trimmed).length > 0 ? trimmed : undefined;
}
return {
...snapshot,
breakdown: { ...breakdown, toolTokenCounts },
...(completedOutputTokens > 0 && { completedOutputTokens }),
};
}
function parseUsageEvents(value?: string | null): TTokenUsageEvent[] {
if (typeof value !== 'string' || value.length === 0) {
return [];
}
try {
const parsed = JSON.parse(value);
return Array.isArray(parsed) ? (parsed as TTokenUsageEvent[]) : [];
} catch {
return [];
}
}
/**
* Builds the response `metadata` for a STOPPED generation from the job's
* persisted emitted usage, so a stopped reply keeps its accurate cost rollup on
* reload (finding: stopped responses otherwise lose cost). Shared by every abort
* save path (agents abort route + legacy abort middleware).
*
* Deliberately persists ONLY `usage`, not `contextUsage`: unlike the live path,
* the abort path can't tell whether the FINAL call (the one the latest snapshot
* precedes) emitted usage — the job stores only the latest snapshot, not the
* snapshot count. If the final call emitted none, `completedOutputTokens` would
* reuse an earlier call's output the snapshot already counts → reload
* over-reports. A stopped/incomplete response therefore falls back to the coarse
* per-message gauge estimate on reload, which is both safe and apt for an
* interrupted turn that never reached a clean pre-invoke breakdown.
*/
export function buildAbortedResponseMetadata(
job: { tokenUsage?: string | null } | null | undefined,
): { usage?: TResponseUsage } | undefined {
const events = parseUsageEvents(job?.tokenUsage);
const usage = aggregateEmittedUsage(events);
return usage ? { usage } : undefined;
}
export interface RecordUsageParams {
user: string;
conversationId: string;