mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
💾 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
GitHub
parent
98704f28c1
commit
b03b2a0a29
19 changed files with 1321 additions and 90 deletions
|
|
@ -9,7 +9,11 @@ const {
|
|||
FakeChatModel,
|
||||
createContentAggregator,
|
||||
} = require('@librechat/agents');
|
||||
const { GenerationJobManager } = require('@librechat/api');
|
||||
const {
|
||||
GenerationJobManager,
|
||||
aggregateEmittedUsage,
|
||||
buildPersistedContextUsage,
|
||||
} = require('@librechat/api');
|
||||
const { getDefaultHandlers } = require('~/server/controllers/agents/callbacks');
|
||||
|
||||
jest.mock('nanoid', () => ({
|
||||
|
|
@ -112,7 +116,14 @@ const SECOND_CALL_USAGE = {
|
|||
|
||||
const MAX_CONTEXT_TOKENS = 8000;
|
||||
|
||||
async function runToolLoop({ res, streamId = null, collectedUsage }) {
|
||||
async function runToolLoop({
|
||||
res,
|
||||
streamId = null,
|
||||
collectedUsage,
|
||||
contextUsageSink = null,
|
||||
usageEmitSink = null,
|
||||
usageCost = null,
|
||||
}) {
|
||||
const { contentParts, aggregateContent } = createContentAggregator();
|
||||
const handlers = getDefaultHandlers({
|
||||
res,
|
||||
|
|
@ -120,6 +131,9 @@ async function runToolLoop({ res, streamId = null, collectedUsage }) {
|
|||
toolEndCallback: () => {},
|
||||
collectedUsage,
|
||||
streamId,
|
||||
contextUsageSink,
|
||||
usageEmitSink,
|
||||
usageCost,
|
||||
});
|
||||
|
||||
const run = await Run.create({
|
||||
|
|
@ -230,6 +244,63 @@ describe('usage events through the real agents pipeline', () => {
|
|||
expect(firstContextIndex).toBeLessThan(firstUsageIndex);
|
||||
});
|
||||
|
||||
test('captures the usage rollup + latest context snapshot for message persistence', () => {
|
||||
const res = createMockRes();
|
||||
const contextUsageSink = { latest: null };
|
||||
const usageEmitSink = [];
|
||||
return runToolLoop({ res, collectedUsage: [], contextUsageSink, usageEmitSink }).then(() => {
|
||||
/** Both model calls' emitted payloads are captured for the rollup */
|
||||
expect(usageEmitSink).toHaveLength(2);
|
||||
|
||||
const usage = aggregateEmittedUsage(usageEmitSink);
|
||||
/** Display units: openAI is cache-subset, so input excludes cache
|
||||
* (150−30−50=70); output is repaired completion */
|
||||
expect(usage).toEqual({
|
||||
input:
|
||||
FIRST_CALL_USAGE.input_tokens +
|
||||
(SECOND_CALL_USAGE.input_tokens -
|
||||
SECOND_CALL_USAGE.input_token_details.cache_creation -
|
||||
SECOND_CALL_USAGE.input_token_details.cache_read),
|
||||
output: FIRST_CALL_USAGE.output_tokens + SECOND_CALL_USAGE.output_tokens,
|
||||
cacheWrite: SECOND_CALL_USAGE.input_token_details.cache_creation,
|
||||
cacheRead: SECOND_CALL_USAGE.input_token_details.cache_read,
|
||||
});
|
||||
/** contextCost off → no cost folded into the rollup */
|
||||
expect(usage.cost).toBeUndefined();
|
||||
|
||||
if (hasContextUsageEvent) {
|
||||
expect(contextUsageSink.latest).not.toBeNull();
|
||||
const persisted = buildPersistedContextUsage(contextUsageSink.latest);
|
||||
expect(persisted.breakdown.maxContextTokens).toBe(MAX_CONTEXT_TOKENS);
|
||||
/** Zero-valued tool counts are trimmed from the persisted blob */
|
||||
for (const count of Object.values(persisted.breakdown.toolTokenCounts ?? {})) {
|
||||
expect(count).toBeGreaterThan(0);
|
||||
}
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
test('folds authoritative per-event cost into the rollup when contextCost is on', async () => {
|
||||
const res = createMockRes();
|
||||
const usageEmitSink = [];
|
||||
/** Stub pricing mirroring getMultiplier/getCacheMultiplier shape */
|
||||
const usageCost = {
|
||||
enabled: true,
|
||||
pricing: {
|
||||
getMultiplier: ({ tokenType }) => (tokenType === 'completion' ? 15 : 3),
|
||||
getCacheMultiplier: ({ cacheType }) => (cacheType === 'write' ? 3.75 : 0.3),
|
||||
},
|
||||
};
|
||||
await runToolLoop({ res, collectedUsage: [], usageEmitSink, usageCost });
|
||||
|
||||
for (const event of usageEmitSink) {
|
||||
expect(typeof event.cost).toBe('number');
|
||||
}
|
||||
const usage = aggregateEmittedUsage(usageEmitSink);
|
||||
expect(usage.cost).toBeGreaterThan(0);
|
||||
expect(usage.cost).toBeCloseTo(usageEmitSink.reduce((sum, e) => sum + e.cost, 0));
|
||||
});
|
||||
|
||||
test('persists usage and context snapshot for resume via GenerationJobManager', async () => {
|
||||
const streamId = `usage-e2e-stream-${Date.now()}`;
|
||||
await GenerationJobManager.createJob(streamId, 'user-1', 'convo-1');
|
||||
|
|
|
|||
|
|
@ -274,6 +274,11 @@ function feedSubagentAggregator(aggregator, event) {
|
|||
* @param {string | null} [options.streamId] - The stream ID for resumable mode, or null for standard mode.
|
||||
* @param {ToolExecuteOptions} [options.toolExecuteOptions] - Options for event-driven tool execution.
|
||||
* @param {UsageCostDeps} [options.usageCost] - Pricing context for authoritative per-event cost.
|
||||
* @param {{ latest: TContextUsageEvent | null, count: number }} [options.contextUsageSink] - Mutable
|
||||
* holder for the latest visible context snapshot + a count of visible snapshots (model calls),
|
||||
* used to persist the breakdown only when the final call emitted usage.
|
||||
* @param {Array<TTokenUsageEvent>} [options.usageEmitSink] - Array collecting each emitted
|
||||
* `on_token_usage` payload (incl. cost) so the response's usage rollup can be persisted.
|
||||
* @returns {Record<string, t.EventHandler>} The default handlers.
|
||||
* @throws {Error} If the request is not found.
|
||||
*/
|
||||
|
|
@ -288,6 +293,8 @@ function getDefaultHandlers({
|
|||
summarizationOptions = null,
|
||||
subagentAggregatorsByToolCallId = null,
|
||||
usageCost = null,
|
||||
contextUsageSink = null,
|
||||
usageEmitSink = null,
|
||||
}) {
|
||||
if (!res || !aggregateContent) {
|
||||
throw new Error(
|
||||
|
|
@ -313,6 +320,11 @@ function getDefaultHandlers({
|
|||
logger.warn('[getDefaultHandlers] Failed to compute usage cost', err);
|
||||
}
|
||||
}
|
||||
/** Collect the same payload the client folds so the response's usage rollup
|
||||
* persisted on `metadata.usage` reproduces the live branch/total + cost. */
|
||||
if (usageEmitSink) {
|
||||
usageEmitSink.push(payload);
|
||||
}
|
||||
return emitEvent(res, streamId, { event: UsageEvents.ON_TOKEN_USAGE, data: payload });
|
||||
};
|
||||
const handlers = {
|
||||
|
|
@ -521,6 +533,15 @@ function getDefaultHandlers({
|
|||
!metadata?.hide_sequential_outputs
|
||||
) {
|
||||
await emitEvent(res, streamId, { event, data });
|
||||
/** Capture the latest visible snapshot (last-wins) + count visible
|
||||
* snapshots (one per model call). The count lets the save path persist
|
||||
* the breakdown only when the FINAL call emitted usage (primary usage
|
||||
* events === snapshots), so completedOutputTokens is a real
|
||||
* post-snapshot delta and reload doesn't over-report. */
|
||||
if (contextUsageSink) {
|
||||
contextUsageSink.latest = data;
|
||||
contextUsageSink.count = (contextUsageSink.count ?? 0) + 1;
|
||||
}
|
||||
}
|
||||
},
|
||||
};
|
||||
|
|
|
|||
|
|
@ -23,6 +23,8 @@ const {
|
|||
recordCollectedUsage,
|
||||
sendEvent,
|
||||
computeUsageCostUSD,
|
||||
aggregateEmittedUsage,
|
||||
buildPersistedContextUsage,
|
||||
createSubagentUsageSink,
|
||||
isDeepSeekReasoningProvider,
|
||||
GenerationJobManager,
|
||||
|
|
@ -108,11 +110,22 @@ class AgentClient extends BaseClient {
|
|||
artifactPromises,
|
||||
maxContextTokens,
|
||||
subagentAggregatorsByToolCallId,
|
||||
contextUsageSink,
|
||||
usageEmitSink,
|
||||
...clientOptions
|
||||
} = options;
|
||||
|
||||
this.agentConfigs = agentConfigs;
|
||||
this.maxContextTokens = maxContextTokens;
|
||||
/** Latest visible context snapshot for this response, captured live by the
|
||||
* ON_CONTEXT_USAGE handler; persisted on `metadata.contextUsage`.
|
||||
* @type {{ latest: import('librechat-data-provider').TContextUsageEvent | null } | undefined} */
|
||||
this.contextUsageSink = contextUsageSink;
|
||||
/** Every emitted `on_token_usage` payload for this response (primary,
|
||||
* summarization, sequential, and subagent); aggregated into the rollup
|
||||
* persisted on `metadata.usage`.
|
||||
* @type {Array<import('librechat-data-provider').TTokenUsageEvent> | undefined} */
|
||||
this.usageEmitSink = usageEmitSink;
|
||||
/** @type {MessageContentComplex[]} */
|
||||
this.contentParts = contentParts;
|
||||
/** @type {Array<UsageMetadata>} */
|
||||
|
|
@ -828,11 +841,49 @@ class AgentClient extends BaseClient {
|
|||
});
|
||||
|
||||
const completion = filterMalformedContentParts(this.contentParts);
|
||||
const metadata = this.buildResponseMetadata();
|
||||
return metadata ? { completion, metadata } : { completion };
|
||||
}
|
||||
|
||||
/**
|
||||
* Assembles the response message `metadata`: Vertex thought signatures plus
|
||||
* the persisted context breakdown (Part A) and the usage/cost rollup (Part B),
|
||||
* which rebuild the gauge breakdown and branch/total cost across reloads.
|
||||
* Returns undefined when nothing was captured.
|
||||
* @returns {{
|
||||
* thoughtSignatures?: Record<string, string>,
|
||||
* contextUsage?: import('librechat-data-provider').TContextUsageEvent,
|
||||
* usage?: import('librechat-data-provider').TResponseUsage,
|
||||
* } | undefined}
|
||||
*/
|
||||
buildResponseMetadata() {
|
||||
/** @type {{
|
||||
* thoughtSignatures?: Record<string, string>,
|
||||
* contextUsage?: import('librechat-data-provider').TContextUsageEvent,
|
||||
* usage?: import('librechat-data-provider').TResponseUsage,
|
||||
* }} */
|
||||
const metadata = {};
|
||||
const signatures = this.collectedThoughtSignatures;
|
||||
if (!signatures || Object.keys(signatures).length === 0) {
|
||||
return { completion };
|
||||
if (signatures && Object.keys(signatures).length > 0) {
|
||||
metadata.thoughtSignatures = signatures;
|
||||
}
|
||||
return { completion, metadata: { thoughtSignatures: signatures } };
|
||||
const usageEvents = this.usageEmitSink ?? [];
|
||||
/** Persist the breakdown only when the FINAL visible call (the one the latest
|
||||
* snapshot precedes) emitted usage — i.e. as many primary usage events as
|
||||
* visible snapshots. If the final call emitted no usage_metadata (provider
|
||||
* gap, or interrupted after an earlier call did emit), `completedOutputTokens`
|
||||
* would be an earlier call's output the latest snapshot already counts, so
|
||||
* reload would over-report; fall back to the coarse per-message estimate. */
|
||||
const primaryUsageCount = usageEvents.filter((event) => event.usage_type == null).length;
|
||||
const snapshotCount = this.contextUsageSink?.count ?? 0;
|
||||
if (this.contextUsageSink?.latest && snapshotCount > 0 && primaryUsageCount >= snapshotCount) {
|
||||
metadata.contextUsage = buildPersistedContextUsage(this.contextUsageSink.latest, usageEvents);
|
||||
}
|
||||
const usage = aggregateEmittedUsage(usageEvents);
|
||||
if (usage) {
|
||||
metadata.usage = usage;
|
||||
}
|
||||
return Object.keys(metadata).length > 0 ? metadata : undefined;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -920,6 +971,12 @@ class AgentClient extends BaseClient {
|
|||
)
|
||||
: undefined,
|
||||
};
|
||||
/** Fold into the response's usage rollup (synchronously, regardless of
|
||||
* emit success) so the persisted total matches the live session, which
|
||||
* also folds subagent usage into its cost/totals. */
|
||||
if (this.usageEmitSink) {
|
||||
this.usageEmitSink.push(data);
|
||||
}
|
||||
/** The sink fires this without awaiting, so retain the promise and flush
|
||||
* it in chatCompletion's finally — emitChunk persists (HSET) before
|
||||
* publishing, and job cleanup must not race that persist or resumed
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ const {
|
|||
GenerationJobManager,
|
||||
recordCollectedUsage,
|
||||
sanitizeMessageForTransmit,
|
||||
buildAbortedResponseMetadata,
|
||||
} = require('@librechat/api');
|
||||
const { truncateText, smartTruncateText } = require('~/app/clients/prompts');
|
||||
const clearPendingReq = require('~/cache/clearPendingReq');
|
||||
|
|
@ -110,6 +111,14 @@ async function abortMessage(req, res) {
|
|||
tokenCount: completionTokens,
|
||||
};
|
||||
|
||||
/** Persist the usage/cost rollup + context breakdown for the stopped response
|
||||
* so its branch/total cost and granular rows survive a reload, matching the
|
||||
* normal completion path. */
|
||||
const abortMetadata = buildAbortedResponseMetadata(jobData);
|
||||
if (abortMetadata) {
|
||||
responseMessage.metadata = abortMetadata;
|
||||
}
|
||||
|
||||
// Spend tokens for ALL models from collectedUsage (handles parallel agents/addedConvo)
|
||||
if (collectedUsage && collectedUsage.length > 0) {
|
||||
await spendCollectedUsage({
|
||||
|
|
|
|||
|
|
@ -1,5 +1,10 @@
|
|||
const express = require('express');
|
||||
const { isEnabled, GenerationJobManager, hasPersistableAbortContent } = require('@librechat/api');
|
||||
const {
|
||||
isEnabled,
|
||||
GenerationJobManager,
|
||||
hasPersistableAbortContent,
|
||||
buildAbortedResponseMetadata,
|
||||
} = require('@librechat/api');
|
||||
const { createSseStreamTelemetry } = require('@librechat/api/telemetry');
|
||||
const { logger } = require('@librechat/data-schemas');
|
||||
const {
|
||||
|
|
@ -291,6 +296,15 @@ router.post('/chat/abort', async (req, res) => {
|
|||
user: userId,
|
||||
};
|
||||
|
||||
/** Persist the usage/cost rollup + context breakdown for the stopped
|
||||
* response (from the job's tracked tokenUsage/contextUsage) so its
|
||||
* branch/total cost and granular rows survive a reload — parity with the
|
||||
* normal completion path. */
|
||||
const abortMetadata = buildAbortedResponseMetadata(jobData);
|
||||
if (abortMetadata) {
|
||||
responseMessage.metadata = abortMetadata;
|
||||
}
|
||||
|
||||
try {
|
||||
await saveMessage(
|
||||
{
|
||||
|
|
|
|||
|
|
@ -252,6 +252,14 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
pricing: { getMultiplier: db.getMultiplier, getCacheMultiplier: db.getCacheMultiplier },
|
||||
};
|
||||
|
||||
/** Latest visible context snapshot + every emitted usage payload for this
|
||||
* response, captured by the handlers and persisted on the response message's
|
||||
* metadata so the breakdown and branch/total cost survive a reload.
|
||||
* @type {{ latest: import('librechat-data-provider').TContextUsageEvent | null, count: number }} */
|
||||
const contextUsageSink = { latest: null, count: 0 };
|
||||
/** @type {Array<import('librechat-data-provider').TTokenUsageEvent>} */
|
||||
const usageEmitSink = [];
|
||||
|
||||
const eventHandlers = getDefaultHandlers({
|
||||
res,
|
||||
toolExecuteOptions,
|
||||
|
|
@ -263,6 +271,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
streamId,
|
||||
subagentAggregatorsByToolCallId,
|
||||
usageCost,
|
||||
contextUsageSink,
|
||||
usageEmitSink,
|
||||
});
|
||||
|
||||
if (!endpointOption.agent) {
|
||||
|
|
@ -920,6 +930,10 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => {
|
|||
/** Resolved endpoint token/pricing config so spending and cost reflect
|
||||
* configured rates for custom-endpoint agents instead of defaults. */
|
||||
endpointTokenConfig: primaryConfig.endpointTokenConfig,
|
||||
/** Capture sinks the handlers fill during the run; `sendCompletion` reads
|
||||
* them to persist the breakdown + usage rollup on the response message. */
|
||||
contextUsageSink,
|
||||
usageEmitSink,
|
||||
});
|
||||
|
||||
if (streamId) {
|
||||
|
|
|
|||
|
|
@ -32,7 +32,12 @@ interface BreakdownProps {
|
|||
|
||||
export default function Breakdown({ view, showCost }: BreakdownProps) {
|
||||
const localize = useLocalize();
|
||||
const { usedTokens, maxTokens, percent, snapshot, snapshotActive, usageTotals } = view;
|
||||
const { usedTokens, maxTokens, percent, snapshot, snapshotActive, branchUsage, hasUsage } = view;
|
||||
/** Show the all-branches total only when it (a) exceeds the active branch —
|
||||
* epsilon guards against float summation order surfacing a spurious row in an
|
||||
* unbranched conversation — and (b) has COMPLETE cost coverage, so a sibling
|
||||
* branch saved without cost can't render an under-reported total. */
|
||||
const showTotal = view.totalUsage.costKnown && view.totalCost - view.branchCost > 1e-9;
|
||||
|
||||
const breakdown = snapshotActive ? snapshot?.breakdown : undefined;
|
||||
const instructionTokens =
|
||||
|
|
@ -139,28 +144,40 @@ export default function Breakdown({ view, showCost }: BreakdownProps) {
|
|||
)}
|
||||
</div>
|
||||
|
||||
{usageTotals.eventCount > 0 && (
|
||||
{hasUsage && (
|
||||
<>
|
||||
<div className="border-t border-border-light" role="separator" />
|
||||
<div className="space-y-1.5" data-testid="token-usage-totals">
|
||||
<Row label={localize('com_ui_input')} value={usageTotals.input} />
|
||||
<Row label={localize('com_ui_output')} value={usageTotals.output} />
|
||||
{usageTotals.cacheRead > 0 && (
|
||||
<Row label={localize('com_ui_cache_read')} value={usageTotals.cacheRead} />
|
||||
<Row label={localize('com_ui_input')} value={branchUsage.input} />
|
||||
<Row label={localize('com_ui_output')} value={branchUsage.output} />
|
||||
{branchUsage.cacheRead > 0 && (
|
||||
<Row label={localize('com_ui_cache_read')} value={branchUsage.cacheRead} />
|
||||
)}
|
||||
{usageTotals.cacheWrite > 0 && (
|
||||
<Row label={localize('com_ui_cache_write')} value={usageTotals.cacheWrite} />
|
||||
{branchUsage.cacheWrite > 0 && (
|
||||
<Row label={localize('com_ui_cache_write')} value={branchUsage.cacheWrite} />
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{showCost && view.costUSD != null && (
|
||||
{showCost && hasUsage && branchUsage.costKnown && (
|
||||
<>
|
||||
<div className="border-t border-border-light" role="separator" />
|
||||
<div className="flex items-center justify-between text-sm">
|
||||
<span className="text-text-secondary">{localize('com_ui_session_cost')}</span>
|
||||
<span className="font-medium text-text-primary">{formatCost(view.costUSD)}</span>
|
||||
<div className="space-y-1.5" data-testid="token-usage-cost">
|
||||
<div className="flex items-center justify-between text-sm">
|
||||
<span className="text-text-secondary">
|
||||
{showTotal
|
||||
? localize('com_ui_context_cost_branch')
|
||||
: localize('com_ui_context_cost')}
|
||||
</span>
|
||||
<span className="font-medium text-text-primary">{formatCost(view.branchCost)}</span>
|
||||
</div>
|
||||
{showTotal && (
|
||||
<div className="flex items-center justify-between text-xs">
|
||||
<span className="text-text-secondary">{localize('com_ui_context_cost_total')}</span>
|
||||
<span className="text-text-secondary">{formatCost(view.totalCost)}</span>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
|
|
|||
|
|
@ -3,17 +3,26 @@ import { useAtomValue, useSetAtom } from 'jotai';
|
|||
import { useQueryClient } from '@tanstack/react-query';
|
||||
import { Constants, QueryKeys } from 'librechat-data-provider';
|
||||
import type { TMessage, TConversation, TModelTokenomics } from 'librechat-data-provider';
|
||||
import type { ContextSnapshot, UsageTotals } from '~/store/usage';
|
||||
import type { BranchTotals } from '~/utils/tokens';
|
||||
import type { BranchTotals, BranchUsage } from '~/utils/tokens';
|
||||
import type { ContextSnapshot } from '~/store/usage';
|
||||
import {
|
||||
liveTokensFamily,
|
||||
totalUsageFamily,
|
||||
removeUsageAtoms,
|
||||
usageTotalsFamily,
|
||||
hydrateSnapshots,
|
||||
pendingUsageFamily,
|
||||
branchTotalsFamily,
|
||||
contextSnapshotFamily,
|
||||
snapshotsByAnchorFamily,
|
||||
} from '~/store/usage';
|
||||
import { buildIndex, sumBranch, clearIndex, findBranchSnapshotAnchor } from '~/utils';
|
||||
import {
|
||||
buildIndex,
|
||||
sumBranch,
|
||||
clearIndex,
|
||||
mergeUsage,
|
||||
sumTotalUsage,
|
||||
findBranchSnapshotAnchor,
|
||||
} from '~/utils';
|
||||
import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage';
|
||||
import useTokenLimits from './useTokenLimits';
|
||||
|
||||
|
|
@ -33,11 +42,18 @@ export interface TokenUsageView {
|
|||
snapshot: ContextSnapshot | null;
|
||||
snapshotActive: boolean;
|
||||
branchTotals: BranchTotals;
|
||||
usageTotals: UsageTotals;
|
||||
/** Provider usage along the active branch (matches the gauge), incl. in-flight */
|
||||
branchUsage: BranchUsage;
|
||||
/** Provider usage across all branches of the conversation */
|
||||
totalUsage: BranchUsage;
|
||||
/** Whether any usage is available to display (branch has token usage) */
|
||||
hasUsage: boolean;
|
||||
/** Authoritative branch cost; the cost row is gated on `interface.contextCost` at render */
|
||||
branchCost: number;
|
||||
/** Authoritative cost across all branches (shown when it differs from branch) */
|
||||
totalCost: number;
|
||||
liveTokens: number;
|
||||
rates?: TModelTokenomics;
|
||||
/** Session cost from provider-reported usage; undefined until usage events arrive */
|
||||
costUSD?: number;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -55,17 +71,39 @@ export default function useTokenUsage({
|
|||
const tailId = useLatestMessageId(index);
|
||||
const snapshot = useAtomValue(contextSnapshotFamily(conversationKey));
|
||||
const snapshotsByAnchor = useAtomValue(snapshotsByAnchorFamily(conversationKey));
|
||||
const usageTotals = useAtomValue(usageTotalsFamily(conversationKey));
|
||||
const pendingUsage = useAtomValue(pendingUsageFamily(conversationKey));
|
||||
const totalUsageBase = useAtomValue(totalUsageFamily(conversationKey));
|
||||
const branchTotals = useAtomValue(branchTotalsFamily(conversationKey));
|
||||
const liveTokens = useAtomValue(liveTokensFamily(conversationKey));
|
||||
const setBranchTotals = useSetAtom(branchTotalsFamily(conversationKey));
|
||||
const setTotalUsage = useSetAtom(totalUsageFamily(conversationKey));
|
||||
const limits = useTokenLimits(conversation);
|
||||
|
||||
/** Authoritative session cost: the backend prices each call (premium tiers,
|
||||
* cache rates) and emits it on the usage event; we just sum. Undefined
|
||||
* until usage events arrive — the cost row is additionally gated on
|
||||
* `interface.contextCost`, under which the backend actually emits cost. */
|
||||
const costUSD = usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined;
|
||||
/** Branch/total provider usage is index-derived; the in-flight response is
|
||||
* the only live add (the pending holder), counted into both — it sits on the
|
||||
* active branch tail and inside the conversation. The backend prices each
|
||||
* call (premium tiers, cache rates), so cost sums authoritatively. */
|
||||
const pendingAsUsage: BranchUsage = useMemo(
|
||||
() => ({
|
||||
input: pendingUsage.input,
|
||||
output: pendingUsage.output,
|
||||
cacheWrite: pendingUsage.cacheWrite,
|
||||
cacheRead: pendingUsage.cacheRead,
|
||||
cost: pendingUsage.costUSD,
|
||||
costKnown: pendingUsage.costKnown,
|
||||
}),
|
||||
[pendingUsage],
|
||||
);
|
||||
const branchUsage = useMemo(
|
||||
() => mergeUsage(branchTotals.usage, pendingAsUsage),
|
||||
[branchTotals.usage, pendingAsUsage],
|
||||
);
|
||||
const totalUsage = useMemo(
|
||||
() => mergeUsage(totalUsageBase, pendingAsUsage),
|
||||
[totalUsageBase, pendingAsUsage],
|
||||
);
|
||||
const hasUsage =
|
||||
branchUsage.input + branchUsage.output + branchUsage.cacheRead + branchUsage.cacheWrite > 0;
|
||||
|
||||
const isSubmittingRef = useRef(isSubmitting);
|
||||
isSubmittingRef.current = isSubmitting;
|
||||
|
|
@ -85,7 +123,11 @@ export default function useTokenUsage({
|
|||
}
|
||||
lastIndexed = messages;
|
||||
buildIndex(conversationKey, messages);
|
||||
/** Restore each branch's persisted breakdown (Part A) without clobbering
|
||||
* a live finalized snapshot for the same response id. */
|
||||
hydrateSnapshots(conversationKey, messages);
|
||||
setBranchTotals(sumBranch(conversationKey, tailIdRef.current, anchorIdRef.current));
|
||||
setTotalUsage(sumTotalUsage(conversationKey));
|
||||
};
|
||||
|
||||
rebuild(queryClient.getQueryData<TMessage[]>([QueryKeys.messages, conversationKey]));
|
||||
|
|
@ -115,24 +157,25 @@ export default function useTokenUsage({
|
|||
removeUsageAtoms(conversationKey);
|
||||
}
|
||||
};
|
||||
}, [conversationKey, queryClient, setBranchTotals]);
|
||||
}, [conversationKey, queryClient, setBranchTotals, setTotalUsage]);
|
||||
|
||||
useEffect(() => {
|
||||
/** The cache subscriber is muted during streaming to avoid per-chunk O(n)
|
||||
* rebuilds, but the `created` event still moves the tail to the new
|
||||
* response message. Without a snapshot (non-agent streams, or a lib that
|
||||
* predates on_context_usage) sumBranch would miss that tail in the stale
|
||||
* index and drop history + prompt. Re-index from the cache on tail change
|
||||
* while submitting — bounded, since tailId only shifts on
|
||||
* created/finalize/branch-switch, never per chunk. */
|
||||
if (isSubmittingRef.current) {
|
||||
buildIndex(
|
||||
conversationKey,
|
||||
queryClient.getQueryData<TMessage[]>([QueryKeys.messages, conversationKey]),
|
||||
);
|
||||
}
|
||||
/** Re-index from the cache on every tail change (created/finalize during a
|
||||
* stream AND branch switches). Branch switches don't fire a cache `updated`
|
||||
* event, so the subscriber below can't catch them; without rebuilding here
|
||||
* the index stays on whatever the last stream left it — which may have
|
||||
* dropped the now-viewed branch's response, so sumBranch would find no
|
||||
* tokens/usage and the gauge + branch cost would blank out. Bounded: tailId
|
||||
* only shifts on created/finalize/branch-switch, never per chunk. Usage for
|
||||
* responses whose cache message lacks `metadata.usage` is restored from the
|
||||
* sticky history inside buildIndex. */
|
||||
buildIndex(
|
||||
conversationKey,
|
||||
queryClient.getQueryData<TMessage[]>([QueryKeys.messages, conversationKey]),
|
||||
);
|
||||
setBranchTotals(sumBranch(conversationKey, tailId, anchorId));
|
||||
}, [conversationKey, tailId, anchorId, setBranchTotals, queryClient]);
|
||||
setTotalUsage(sumTotalUsage(conversationKey));
|
||||
}, [conversationKey, tailId, anchorId, setBranchTotals, setTotalUsage, queryClient]);
|
||||
|
||||
return useMemo(() => {
|
||||
/** The granular snapshot is for one specific generation. Show the live one
|
||||
|
|
@ -179,10 +222,13 @@ export default function useTokenUsage({
|
|||
snapshot: activeSnapshot,
|
||||
snapshotActive: true,
|
||||
branchTotals,
|
||||
usageTotals,
|
||||
branchUsage,
|
||||
totalUsage,
|
||||
hasUsage,
|
||||
branchCost: branchUsage.cost,
|
||||
totalCost: totalUsage.cost,
|
||||
liveTokens,
|
||||
rates: limits.rates,
|
||||
costUSD,
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -197,19 +243,23 @@ export default function useTokenUsage({
|
|||
snapshot: null,
|
||||
snapshotActive: false,
|
||||
branchTotals,
|
||||
usageTotals,
|
||||
branchUsage,
|
||||
totalUsage,
|
||||
hasUsage,
|
||||
branchCost: branchUsage.cost,
|
||||
totalCost: totalUsage.cost,
|
||||
liveTokens,
|
||||
rates: limits.rates,
|
||||
costUSD,
|
||||
};
|
||||
}, [
|
||||
snapshot,
|
||||
isSubmitting,
|
||||
branchTotals,
|
||||
usageTotals,
|
||||
branchUsage,
|
||||
totalUsage,
|
||||
hasUsage,
|
||||
liveTokens,
|
||||
limits,
|
||||
costUSD,
|
||||
snapshotsByAnchor,
|
||||
conversationKey,
|
||||
]);
|
||||
|
|
|
|||
|
|
@ -1031,6 +1031,11 @@ export default function useResumableSSE(
|
|||
setIsSubmitting(false);
|
||||
setShowStopButton(false);
|
||||
setStreamId(null);
|
||||
/** Intentional close without a final event (explicit stop, or navigation
|
||||
* while generating): discard the in-flight pending usage so it can't
|
||||
* merge into the next response in this conversation. On a resume the
|
||||
* collected usage is re-folded via backfillUsage, so nothing is lost. */
|
||||
resetLive({ ...currentSubmission, userMessage });
|
||||
});
|
||||
|
||||
// Start the SSE connection
|
||||
|
|
|
|||
|
|
@ -67,8 +67,15 @@ export default function useSSE(
|
|||
const balanceQuery = useGetUserBalance({
|
||||
enabled: !!isAuthenticated && startupConfig?.balance?.enabled,
|
||||
});
|
||||
const { contextHandler, usageHandler, tapStream, tapContent, finalizeUsage, resetLive } =
|
||||
useUsageHandler();
|
||||
const {
|
||||
contextHandler,
|
||||
usageHandler,
|
||||
tapStream,
|
||||
tapContent,
|
||||
finalizeUsage,
|
||||
resetLive,
|
||||
attributePending,
|
||||
} = useUsageHandler();
|
||||
|
||||
useEffect(() => {
|
||||
if (submission == null || Object.keys(submission).length === 0) {
|
||||
|
|
@ -186,9 +193,15 @@ export default function useSSE(
|
|||
}
|
||||
|
||||
setCompleted((prev) => new Set(prev.add(streamKey)));
|
||||
resetLive({ ...submission, userMessage });
|
||||
const latestMessages = getMessages();
|
||||
const conversationId = latestMessages?.[latestMessages.length - 1]?.conversationId;
|
||||
/** Attribute usage billed before the stop to the partial response (the
|
||||
* branch tail), then reset pending — so it neither drops nor leaks into
|
||||
* the next response. Falls back to a plain reset when no response exists. */
|
||||
const tail = latestMessages?.[latestMessages.length - 1];
|
||||
const partialResponseId =
|
||||
tail != null && tail.isCreatedByUser === false ? tail.messageId : null;
|
||||
attributePending(partialResponseId, { ...submission, userMessage });
|
||||
try {
|
||||
await abortConversation(
|
||||
conversationId ??
|
||||
|
|
|
|||
|
|
@ -11,18 +11,23 @@ import type { ContextSnapshot } from '~/store/usage';
|
|||
import {
|
||||
markUsageFolded,
|
||||
liveTokensFamily,
|
||||
totalUsageFamily,
|
||||
removeUsageAtoms,
|
||||
clearUsageFolded,
|
||||
calibrationFamily,
|
||||
usageTotalsFamily,
|
||||
pendingUsageFamily,
|
||||
branchTotalsFamily,
|
||||
migrateUsageFolded,
|
||||
EMPTY_USAGE_TOTALS,
|
||||
contextSnapshotFamily,
|
||||
snapshotsByAnchorFamily,
|
||||
} from '~/store/usage';
|
||||
import {
|
||||
sumBranch,
|
||||
setEntryUsage,
|
||||
upsertEntries,
|
||||
migrateIndex,
|
||||
sumTotalUsage,
|
||||
estimateTokens,
|
||||
normalizeUsageUnits,
|
||||
} from '~/utils';
|
||||
|
|
@ -49,6 +54,10 @@ export interface UsageHandlers {
|
|||
tapContent: (text: unknown, submission: UsageSubmissionLike) => void;
|
||||
finalizeUsage: (data: FinalDataLike, submission: UsageSubmissionLike) => void;
|
||||
resetLive: (submission: UsageSubmissionLike) => void;
|
||||
/** Terminal stop: attribute the in-flight pending usage to the stopped partial
|
||||
* response (so its billed tokens aren't dropped), then reset pending so it
|
||||
* can't leak into the next response. Discards when no response id is known. */
|
||||
attributePending: (responseId: string | null, submission: UsageSubmissionLike) => void;
|
||||
/** Idempotently folds the resumed run's collected usage into the totals */
|
||||
backfillUsage: (entries: TTokenUsageEvent[], submission: UsageSubmissionLike) => void;
|
||||
/** Seeds the live estimate from already-streamed output chars on resume */
|
||||
|
|
@ -111,6 +120,27 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
jotai.set(liveTokensFamily(convoKey), value);
|
||||
};
|
||||
|
||||
/** Flush the in-flight pending usage into a response's index entry, then
|
||||
* reset pending. Only flushes when events were actually folded this session
|
||||
* (eventCount > 0), so a finalize that carries persisted `metadata.usage`
|
||||
* but folded nothing — a late/second resumable subscriber — keeps the entry
|
||||
* loaded by `upsertEntries` instead of overwriting it with an empty record. */
|
||||
const flushPendingInto = (convoKey: string, responseId: string | null) => {
|
||||
const pendingAtom = pendingUsageFamily(convoKey);
|
||||
const pending = jotai.get(pendingAtom);
|
||||
if (responseId != null && pending.eventCount > 0) {
|
||||
setEntryUsage(convoKey, responseId, {
|
||||
input: pending.input,
|
||||
output: pending.output,
|
||||
cacheWrite: pending.cacheWrite,
|
||||
cacheRead: pending.cacheRead,
|
||||
cost: pending.costUSD,
|
||||
costKnown: pending.costKnown,
|
||||
});
|
||||
}
|
||||
jotai.set(pendingAtom, EMPTY_USAGE_TOTALS);
|
||||
};
|
||||
|
||||
const contextHandler: UsageHandlers['contextHandler'] = (data, submission) => {
|
||||
const convoKey = getConvoKey(submission);
|
||||
jotai.set(contextSnapshotFamily(convoKey), {
|
||||
|
|
@ -125,9 +155,11 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
setLive(convoKey, 0);
|
||||
};
|
||||
|
||||
/** Folds one usage event into the totals exactly once per conversation.
|
||||
* Returns false when the event was already counted (live then replayed
|
||||
* on resume), so callers can skip the live-estimate bump too. */
|
||||
/** Folds one usage event into the in-flight pending holder exactly once per
|
||||
* conversation. Returns false when the event was already counted (live then
|
||||
* replayed on resume), so callers can skip the live-estimate bump too.
|
||||
* `finalizeUsage` flushes the accumulated pending into the per-message index
|
||||
* and resets it, so branch/total stay index-derived (no double count). */
|
||||
const foldUsage = (data: TTokenUsageEvent, submission: UsageSubmissionLike): boolean => {
|
||||
const convoKey = getConvoKey(submission);
|
||||
/** runId+seq is unique per model call; fall back to the payload when a
|
||||
|
|
@ -142,9 +174,9 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
* the uncached portion, output includes repaired completion tokens */
|
||||
const units = normalizeUsageUnits(data);
|
||||
|
||||
const totalsAtom = usageTotalsFamily(convoKey);
|
||||
const prev = jotai.get(totalsAtom);
|
||||
jotai.set(totalsAtom, {
|
||||
const pendingAtom = pendingUsageFamily(convoKey);
|
||||
const prev = jotai.get(pendingAtom);
|
||||
jotai.set(pendingAtom, {
|
||||
input: prev.input + units.input,
|
||||
output: prev.output + units.output,
|
||||
cacheWrite: prev.cacheWrite + units.cacheWrite,
|
||||
|
|
@ -153,6 +185,8 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
/** Authoritative per-event cost from the backend (premium tiers, cache
|
||||
* rates); absent when contextCost is disabled — sums to 0 then */
|
||||
costUSD: prev.costUSD + (data.cost ?? 0),
|
||||
/** Coverage is complete only if EVERY folded event carried a cost */
|
||||
costKnown: prev.costKnown && data.cost != null,
|
||||
});
|
||||
return true;
|
||||
};
|
||||
|
|
@ -209,7 +243,31 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
const resetLive: UsageHandlers['resetLive'] = (submission) => {
|
||||
streamCharsRef.current = 0;
|
||||
confirmedRef.current = 0;
|
||||
setLive(getConvoKey(submission), 0);
|
||||
const convoKey = getConvoKey(submission);
|
||||
setLive(convoKey, 0);
|
||||
/** Terminal path with no salvageable response (stream error / intentional
|
||||
* close): discard the in-flight pending usage so it can't merge into the
|
||||
* next response. The user-stop path uses `attributePending` to keep it on
|
||||
* the partial reply. Also forget the folded-event identities so a resume's
|
||||
* `backfillUsage` can rebuild pending — otherwise it sees them as already
|
||||
* folded and the response's usage stays missing until a full reload. */
|
||||
jotai.set(pendingUsageFamily(convoKey), EMPTY_USAGE_TOTALS);
|
||||
clearUsageFolded(convoKey);
|
||||
};
|
||||
|
||||
const attributePending: UsageHandlers['attributePending'] = (responseId, submission) => {
|
||||
const convoKey = getConvoKey(submission);
|
||||
/** Flush the billed-but-uncommitted usage onto the stopped partial reply
|
||||
* (when its id is known and events were folded), then reset pending and
|
||||
* the live estimate. Index-derived branch/total then reflect it. */
|
||||
flushPendingInto(convoKey, responseId);
|
||||
if (responseId != null) {
|
||||
jotai.set(branchTotalsFamily(convoKey), sumBranch(convoKey, responseId, responseId));
|
||||
jotai.set(totalUsageFamily(convoKey), sumTotalUsage(convoKey));
|
||||
}
|
||||
streamCharsRef.current = 0;
|
||||
confirmedRef.current = 0;
|
||||
setLive(convoKey, 0);
|
||||
};
|
||||
|
||||
const backfillUsage: UsageHandlers['backfillUsage'] = (entries, submission) => {
|
||||
|
|
@ -246,7 +304,7 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
migrateUsageFolded(fromKey, realId);
|
||||
jotai.set(contextSnapshotFamily(realId), jotai.get(contextSnapshotFamily(fromKey)));
|
||||
jotai.set(snapshotsByAnchorFamily(realId), jotai.get(snapshotsByAnchorFamily(fromKey)));
|
||||
jotai.set(usageTotalsFamily(realId), jotai.get(usageTotalsFamily(fromKey)));
|
||||
jotai.set(pendingUsageFamily(realId), jotai.get(pendingUsageFamily(fromKey)));
|
||||
jotai.set(calibrationFamily(realId), jotai.get(calibrationFamily(fromKey)));
|
||||
removeUsageAtoms(fromKey);
|
||||
}
|
||||
|
|
@ -258,10 +316,18 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
|
||||
const userMsgId = submission.userMessage?.messageId ?? null;
|
||||
const responseId = data.responseMessage?.messageId ?? null;
|
||||
|
||||
/** Flush the in-flight response's pending usage into its index entry, then
|
||||
* reset pending. Branch/total are summed from the index, so this single
|
||||
* add is counted exactly once; the persisted `metadata.usage` reproduces
|
||||
* it on reload. */
|
||||
flushPendingInto(realId, responseId);
|
||||
|
||||
const tailId = responseId ?? data.requestMessage?.messageId ?? null;
|
||||
if (tailId) {
|
||||
jotai.set(branchTotalsFamily(realId), sumBranch(realId, tailId, responseId ?? userMsgId));
|
||||
}
|
||||
jotai.set(totalUsageFamily(realId), sumTotalUsage(realId));
|
||||
|
||||
const snapshotAtom = contextSnapshotFamily(realId);
|
||||
const snapshot = jotai.get(snapshotAtom);
|
||||
|
|
@ -300,6 +366,7 @@ export default function useUsageHandler(): UsageHandlers {
|
|||
tapContent,
|
||||
finalizeUsage,
|
||||
resetLive,
|
||||
attributePending,
|
||||
backfillUsage,
|
||||
seedLive,
|
||||
};
|
||||
|
|
|
|||
|
|
@ -884,6 +884,9 @@
|
|||
"com_ui_connecting": "Connecting",
|
||||
"com_ui_contact_admin_if_issue_persists": "Contact the Admin if the issue persists",
|
||||
"com_ui_context": "Context",
|
||||
"com_ui_context_cost": "Cost",
|
||||
"com_ui_context_cost_branch": "Cost (this branch)",
|
||||
"com_ui_context_cost_total": "All branches",
|
||||
"com_ui_context_filter_sort": "Filter and Sort by Context",
|
||||
"com_ui_context_free": "Free space",
|
||||
"com_ui_context_messages": "Messages",
|
||||
|
|
@ -1526,7 +1529,6 @@
|
|||
"com_ui_select_search_model": "Search model by name",
|
||||
"com_ui_select_search_provider": "Search provider by name",
|
||||
"com_ui_select_search_region": "Search region by name",
|
||||
"com_ui_session_cost": "Session cost",
|
||||
"com_ui_set": "Set",
|
||||
"com_ui_share": "Share",
|
||||
"com_ui_share_create_message": "Your name and any messages you add after sharing stay private.",
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
import { atom } from 'jotai';
|
||||
import { atomFamily } from 'jotai/utils';
|
||||
import type { TContextUsageEvent } from 'librechat-data-provider';
|
||||
import type { BranchTotals } from '~/utils/tokens';
|
||||
import { EMPTY_BRANCH } from '~/utils/tokens';
|
||||
import { atom, getDefaultStore } from 'jotai';
|
||||
import type { TMessage, TContextUsageEvent } from 'librechat-data-provider';
|
||||
import type { BranchTotals, BranchUsage } from '~/utils/tokens';
|
||||
import { EMPTY_BRANCH, EMPTY_USAGE } from '~/utils/tokens';
|
||||
|
||||
/** Latest backend context snapshot, anchored to the run's user message for staleness checks */
|
||||
export interface ContextSnapshot extends TContextUsageEvent {
|
||||
|
|
@ -11,7 +11,13 @@ export interface ContextSnapshot extends TContextUsageEvent {
|
|||
completedOutputTokens?: number;
|
||||
}
|
||||
|
||||
/** Cumulative provider-reported usage for the conversation's current session */
|
||||
/**
|
||||
* In-flight usage for the streaming response only — a single-response pending
|
||||
* holder. `foldUsage` accumulates the current response's `on_token_usage`
|
||||
* events here; `finalizeUsage` flushes it into the per-message index and resets
|
||||
* it. Branch/total figures are otherwise derived by summing the index, so this
|
||||
* is the only live add (counted exactly once at finalize).
|
||||
*/
|
||||
export interface UsageTotals {
|
||||
input: number;
|
||||
output: number;
|
||||
|
|
@ -21,6 +27,9 @@ export interface UsageTotals {
|
|||
/** Summed authoritative per-event cost from the backend (premium tiers,
|
||||
* cache rates). Populated only when `interface.contextCost` is enabled. */
|
||||
costUSD: number;
|
||||
/** Whether cost coverage is complete — every folded event carried a cost
|
||||
* (ANDed, vacuously true when empty); gates the cost row. */
|
||||
costKnown: boolean;
|
||||
}
|
||||
|
||||
export const EMPTY_USAGE_TOTALS: UsageTotals = {
|
||||
|
|
@ -30,6 +39,7 @@ export const EMPTY_USAGE_TOTALS: UsageTotals = {
|
|||
cacheRead: 0,
|
||||
eventCount: 0,
|
||||
costUSD: 0,
|
||||
costKnown: true,
|
||||
};
|
||||
|
||||
/**
|
||||
|
|
@ -55,10 +65,16 @@ export const snapshotsByAnchorFamily = atomFamily((_conversationId: string) =>
|
|||
atom<Map<string, ContextSnapshot>>(new Map()),
|
||||
);
|
||||
|
||||
export const usageTotalsFamily = atomFamily((_conversationId: string) =>
|
||||
/** In-flight usage of the streaming response; flushed into the index at finalize. */
|
||||
export const pendingUsageFamily = atomFamily((_conversationId: string) =>
|
||||
atom<UsageTotals>(EMPTY_USAGE_TOTALS),
|
||||
);
|
||||
|
||||
/** Provider usage/cost summed across all branches of the conversation. */
|
||||
export const totalUsageFamily = atomFamily((_conversationId: string) =>
|
||||
atom<BranchUsage>(EMPTY_USAGE),
|
||||
);
|
||||
|
||||
/** Throttled in-flight output token estimate for the current model call */
|
||||
export const liveTokensFamily = atomFamily((_conversationId: string) => atom<number>(0));
|
||||
|
||||
|
|
@ -101,13 +117,67 @@ export function migrateUsageFolded(fromId: string, toId: string): void {
|
|||
foldedUsageKeys.delete(fromId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Forgets a conversation's folded usage-event identities so a subsequent resume
|
||||
* can re-fold them. Used when a terminal close discards the in-flight pending
|
||||
* usage: `backfillUsage` would otherwise treat the persisted events as already
|
||||
* folded and never rebuild pending after a navigate-away-then-resume.
|
||||
*/
|
||||
export function clearUsageFolded(conversationId: string): void {
|
||||
foldedUsageKeys.delete(conversationId);
|
||||
}
|
||||
|
||||
/** Jotai atomFamily entries are never GC'd — call on conversation switch/cleanup */
|
||||
export function removeUsageAtoms(conversationId: string): void {
|
||||
branchTotalsFamily.remove(conversationId);
|
||||
contextSnapshotFamily.remove(conversationId);
|
||||
snapshotsByAnchorFamily.remove(conversationId);
|
||||
usageTotalsFamily.remove(conversationId);
|
||||
pendingUsageFamily.remove(conversationId);
|
||||
totalUsageFamily.remove(conversationId);
|
||||
liveTokensFamily.remove(conversationId);
|
||||
calibrationFamily.remove(conversationId);
|
||||
foldedUsageKeys.delete(conversationId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Rehydrates per-branch context breakdowns from persisted `metadata.contextUsage`
|
||||
* into the snapshot-history map so the granular gauge survives a reload / opening
|
||||
* an existing conversation. Merges, never clobbers — a live finalized snapshot
|
||||
* for the same response id wins. Re-anchors each blob to its (branch-unique)
|
||||
* response message id and reads the response's `tokenCount` as completed output.
|
||||
*/
|
||||
export function hydrateSnapshots(conversationId: string, messages?: TMessage[] | null): void {
|
||||
if (messages == null || messages.length === 0) {
|
||||
return;
|
||||
}
|
||||
const store = getDefaultStore();
|
||||
const historyAtom = snapshotsByAnchorFamily(conversationId);
|
||||
const current = store.get(historyAtom);
|
||||
let next: Map<string, ContextSnapshot> | null = null;
|
||||
for (const message of messages) {
|
||||
const id = message?.messageId;
|
||||
if (!id || current.has(id)) {
|
||||
continue;
|
||||
}
|
||||
const blob = message.metadata?.contextUsage;
|
||||
if (blob == null || typeof blob !== 'object') {
|
||||
continue;
|
||||
}
|
||||
const event = blob as TContextUsageEvent;
|
||||
/** Use ONLY the persisted post-snapshot delta (the final call's output). Do
|
||||
* NOT fall back to the full response `tokenCount`: the snapshot already
|
||||
* counts earlier steps' output for multi-call turns, so adding the whole
|
||||
* tokenCount would double-count after reload. Absent (rare: no usage event)
|
||||
* contributes 0, matching the snapshot's pre-final-call base. */
|
||||
const snapshot: ContextSnapshot = {
|
||||
...event,
|
||||
anchorMessageId: id,
|
||||
completedOutputTokens: event.completedOutputTokens,
|
||||
};
|
||||
next ??= new Map(current);
|
||||
next.set(id, snapshot);
|
||||
}
|
||||
if (next != null) {
|
||||
store.set(historyAtom, next);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import { Constants, Providers } from 'librechat-data-provider';
|
||||
import type { TMessage } from 'librechat-data-provider';
|
||||
import type { TMessage, TResponseUsage } from 'librechat-data-provider';
|
||||
import {
|
||||
buildIndex,
|
||||
upsertEntries,
|
||||
|
|
@ -7,6 +7,9 @@ import {
|
|||
clearIndex,
|
||||
hasIndex,
|
||||
sumBranch,
|
||||
mergeUsage,
|
||||
setEntryUsage,
|
||||
sumTotalUsage,
|
||||
findBranchSnapshotAnchor,
|
||||
estimateTokens,
|
||||
normalizeUsageUnits,
|
||||
|
|
@ -14,6 +17,7 @@ import {
|
|||
groupToolTokens,
|
||||
countTrailingOutputChars,
|
||||
EMPTY_BRANCH,
|
||||
EMPTY_USAGE,
|
||||
EMPTY_TOOL_GROUPS,
|
||||
} from './tokens';
|
||||
|
||||
|
|
@ -35,6 +39,28 @@ function msg(
|
|||
} as TMessage;
|
||||
}
|
||||
|
||||
/** Response message carrying a persisted `metadata.usage` rollup (the backend
|
||||
* persists already-normalized display units). */
|
||||
function responseMsg(
|
||||
messageId: string,
|
||||
parentMessageId: string | null,
|
||||
tokenCount: number,
|
||||
usage: TResponseUsage,
|
||||
): TMessage {
|
||||
return {
|
||||
messageId,
|
||||
parentMessageId,
|
||||
isCreatedByUser: false,
|
||||
tokenCount,
|
||||
conversationId: CONVO,
|
||||
text: '',
|
||||
metadata: { usage },
|
||||
} as TMessage;
|
||||
}
|
||||
|
||||
const USAGE_A: TResponseUsage = { input: 100, output: 50, cacheWrite: 0, cacheRead: 0, cost: 0.01 };
|
||||
const USAGE_B: TResponseUsage = { input: 200, output: 80, cacheWrite: 0, cacheRead: 0, cost: 0.02 };
|
||||
|
||||
describe('token index', () => {
|
||||
afterEach(() => {
|
||||
clearIndex(CONVO);
|
||||
|
|
@ -271,3 +297,197 @@ describe('countTrailingOutputChars', () => {
|
|||
expect(countTrailingOutputChars([tool()])).toBe(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('per-message usage index (branch + total)', () => {
|
||||
afterEach(() => {
|
||||
clearIndex(CONVO);
|
||||
});
|
||||
|
||||
it('reads metadata.usage onto entries and sums it along the branch', () => {
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
responseMsg('a1', 'u1', 50, USAGE_A),
|
||||
msg('u2', 'a1', true, 30),
|
||||
responseMsg('a2', 'u2', 80, USAGE_B),
|
||||
]);
|
||||
|
||||
const { usage } = sumBranch(CONVO, 'a2');
|
||||
expect(usage).toEqual({
|
||||
input: 300,
|
||||
output: 130,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 0,
|
||||
cost: 0.03,
|
||||
costKnown: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('reads persisted display units (incl. cache) directly', () => {
|
||||
/** The backend already normalized per-event; the client reads as-is */
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
responseMsg('a1', 'u1', 100, {
|
||||
input: 600,
|
||||
output: 100,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 400,
|
||||
cost: 0.03,
|
||||
}),
|
||||
]);
|
||||
|
||||
const { usage } = sumBranch(CONVO, 'a1');
|
||||
expect(usage).toEqual({
|
||||
input: 600,
|
||||
output: 100,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 400,
|
||||
cost: 0.03,
|
||||
costKnown: true,
|
||||
});
|
||||
});
|
||||
|
||||
it('marks cost unknown when persisted usage omits cost (contextCost off)', () => {
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
responseMsg('a1', 'u1', 50, { input: 100, output: 50, cacheWrite: 0, cacheRead: 0 }),
|
||||
]);
|
||||
const { usage } = sumBranch(CONVO, 'a1');
|
||||
expect(usage.costKnown).toBe(false);
|
||||
expect(usage.cost).toBe(0);
|
||||
expect(usage.input).toBe(100);
|
||||
});
|
||||
|
||||
it('messages without metadata.usage contribute zero (backward compat)', () => {
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10), msg('a1', 'u1', false, 50)]);
|
||||
expect(sumBranch(CONVO, 'a1').usage).toEqual(EMPTY_USAGE);
|
||||
expect(sumTotalUsage(CONVO)).toEqual(EMPTY_USAGE);
|
||||
});
|
||||
|
||||
it('preserves a prior entry usage when a rebuilt message lacks metadata.usage', () => {
|
||||
/** Live finalize flushes usage into the index before persisted metadata
|
||||
* reaches the cache; a mid-session rebuild from a cache message without
|
||||
* metadata.usage must not wipe it (regenerate keeps the sibling's cost). */
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10), msg('a1', 'u1', false, 50)]);
|
||||
setEntryUsage(CONVO, 'a1', {
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 0,
|
||||
cost: 0.01,
|
||||
costKnown: true,
|
||||
});
|
||||
/** Rebuild from the same cache (still no metadata.usage on a1) */
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10), msg('a1', 'u1', false, 50)]);
|
||||
expect(sumBranch(CONVO, 'a1').usage.cost).toBeCloseTo(0.01);
|
||||
expect(sumBranch(CONVO, 'a1').usage.input).toBe(100);
|
||||
});
|
||||
|
||||
it('restores branch cost from sticky history after regenerate drops then re-adds a sibling', () => {
|
||||
/** a1 generated live; its cache message never carries metadata.usage */
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10), msg('a1', 'u1', false, 50)]);
|
||||
setEntryUsage(CONVO, 'a1', {
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 0,
|
||||
cost: 0.01,
|
||||
costKnown: true,
|
||||
});
|
||||
|
||||
/** Regenerate streaming shows only the active branch — a1 is dropped */
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10), msg('a1-alt', 'u1', false, 60)]);
|
||||
expect(sumBranch(CONVO, 'a1').usage).toEqual(EMPTY_USAGE);
|
||||
|
||||
/** Post-regenerate full rebuild re-adds a1 (still no metadata.usage) */
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
msg('a1', 'u1', false, 50),
|
||||
msg('a1-alt', 'u1', false, 60),
|
||||
]);
|
||||
|
||||
/** Switching back to branch A must still show its cost (the reported bug) */
|
||||
expect(sumBranch(CONVO, 'a1').usage.cost).toBeCloseTo(0.01);
|
||||
expect(sumBranch(CONVO, 'a1').usage.input).toBe(100);
|
||||
});
|
||||
|
||||
it('scopes branch usage to the active thread while total spans all branches', () => {
|
||||
/** Regenerate: a1 and a1-alt are sibling responses under the same user msg */
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
responseMsg('a1', 'u1', 50, USAGE_A),
|
||||
responseMsg('a1-alt', 'u1', 80, USAGE_B),
|
||||
]);
|
||||
|
||||
/** Viewing branch B (the regenerated response) */
|
||||
expect(sumBranch(CONVO, 'a1-alt').usage.cost).toBeCloseTo(0.02);
|
||||
/** Viewing branch A (the original) */
|
||||
expect(sumBranch(CONVO, 'a1').usage.cost).toBeCloseTo(0.01);
|
||||
/** Total spans both abandoned + active branches */
|
||||
const total = sumTotalUsage(CONVO);
|
||||
expect(total.cost).toBeCloseTo(0.03);
|
||||
expect(total.input).toBe(300);
|
||||
expect(total.output).toBe(130);
|
||||
});
|
||||
|
||||
it('is idempotent across rebuilds', () => {
|
||||
const messages = [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
responseMsg('a1', 'u1', 50, USAGE_A),
|
||||
responseMsg('a1-alt', 'u1', 80, USAGE_B),
|
||||
];
|
||||
buildIndex(CONVO, messages);
|
||||
const first = sumTotalUsage(CONVO);
|
||||
buildIndex(CONVO, messages);
|
||||
const second = sumTotalUsage(CONVO);
|
||||
expect(second).toEqual(first);
|
||||
});
|
||||
|
||||
it('flushes a live response usage via setEntryUsage (no metadata yet)', () => {
|
||||
buildIndex(CONVO, [msg('u1', Constants.NO_PARENT, true, 10), msg('a1', 'u1', false, 50)]);
|
||||
/** Live response entry has no metadata.usage until persisted; finalize flushes it */
|
||||
expect(sumBranch(CONVO, 'a1').usage).toEqual(EMPTY_USAGE);
|
||||
setEntryUsage(CONVO, 'a1', {
|
||||
input: 100,
|
||||
output: 50,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 0,
|
||||
cost: 0.01,
|
||||
costKnown: true,
|
||||
});
|
||||
expect(sumBranch(CONVO, 'a1').usage.cost).toBeCloseTo(0.01);
|
||||
expect(sumBranch(CONVO, 'a1').usage.input).toBe(100);
|
||||
});
|
||||
|
||||
it('mergeUsage sums records and ANDs costKnown (incomplete coverage wins)', () => {
|
||||
const a = { input: 1, output: 2, cacheWrite: 3, cacheRead: 4, cost: 0.5, costKnown: false };
|
||||
const b = { input: 10, output: 20, cacheWrite: 30, cacheRead: 40, cost: 1.5, costKnown: true };
|
||||
expect(mergeUsage(a, b)).toEqual({
|
||||
input: 11,
|
||||
output: 22,
|
||||
cacheWrite: 33,
|
||||
cacheRead: 44,
|
||||
cost: 2,
|
||||
costKnown: false,
|
||||
});
|
||||
});
|
||||
|
||||
it('marks branch cost incomplete when any usage-bearing entry lacks cost', () => {
|
||||
/** A turn saved before cost display was on (no cost) alongside one with cost
|
||||
* → the summed cost under-reports, so coverage is incomplete. */
|
||||
buildIndex(CONVO, [
|
||||
msg('u1', Constants.NO_PARENT, true, 10),
|
||||
responseMsg('a1', 'u1', 50, { input: 100, output: 50, cacheWrite: 0, cacheRead: 0 }),
|
||||
responseMsg('a2', 'a1', 60, USAGE_B),
|
||||
]);
|
||||
const { usage } = sumBranch(CONVO, 'a2');
|
||||
expect(usage.costKnown).toBe(false);
|
||||
/** Cost still sums what it can, but coverage is flagged incomplete */
|
||||
expect(usage.cost).toBeCloseTo(0.02);
|
||||
expect(sumTotalUsage(CONVO).costKnown).toBe(false);
|
||||
});
|
||||
|
||||
it('EMPTY_BRANCH carries an empty usage record', () => {
|
||||
expect(sumBranch('missing-convo', 'x')).toBe(EMPTY_BRANCH);
|
||||
expect(EMPTY_BRANCH.usage).toEqual(EMPTY_USAGE);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,10 +1,36 @@
|
|||
import { Tools, Constants, inputTokensIncludesCache } from 'librechat-data-provider';
|
||||
import type { TMessage, TTokenUsageEvent } from 'librechat-data-provider';
|
||||
import type { TMessage, TResponseUsage, TTokenUsageEvent } from 'librechat-data-provider';
|
||||
|
||||
/** Provider-reported usage of one response, in display units (post-normalize). */
|
||||
export interface BranchUsage {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheWrite: number;
|
||||
cacheRead: number;
|
||||
/** Authoritative USD cost; 0 when `interface.contextCost` was off at save */
|
||||
cost: number;
|
||||
/** Whether cost coverage is COMPLETE — every usage-bearing response summed
|
||||
* here carried a cost. Gates the cost row so a partial sum (some turns saved
|
||||
* before cost display was on) never renders an under-reported total. Starts
|
||||
* true (vacuous) and is ANDed; pair with `hasUsage` to require ≥1 entry. */
|
||||
costKnown: boolean;
|
||||
}
|
||||
|
||||
export const EMPTY_USAGE: BranchUsage = {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheWrite: 0,
|
||||
cacheRead: 0,
|
||||
cost: 0,
|
||||
costKnown: true,
|
||||
};
|
||||
|
||||
export interface TokenEntry {
|
||||
tokenCount: number;
|
||||
isCreatedByUser: boolean;
|
||||
parentMessageId: string | null;
|
||||
/** Per-response provider usage from `metadata.usage` (response messages only) */
|
||||
usage?: BranchUsage;
|
||||
}
|
||||
|
||||
export interface BranchTotals {
|
||||
|
|
@ -19,6 +45,8 @@ export interface BranchTotals {
|
|||
tailId: string | null;
|
||||
/** Whether the latest run's anchor message is on this branch */
|
||||
containsAnchor: boolean;
|
||||
/** Provider usage/cost summed along the active branch */
|
||||
usage: BranchUsage;
|
||||
}
|
||||
|
||||
export const EMPTY_BRANCH: BranchTotals = {
|
||||
|
|
@ -28,26 +56,102 @@ export const EMPTY_BRANCH: BranchTotals = {
|
|||
total: 0,
|
||||
tailId: null,
|
||||
containsAnchor: false,
|
||||
usage: EMPTY_USAGE,
|
||||
};
|
||||
|
||||
/** Module-level token index: conversationId → messageId → entry. Not render state. */
|
||||
const registry = new Map<string, Map<string, TokenEntry>>();
|
||||
|
||||
/**
|
||||
* Sticky per-response usage: conversationId → messageId → usage. Written only by
|
||||
* the live finalize/stop flush (`setEntryUsage`) and never rebuilt from the
|
||||
* messages cache, so a response's usage survives mid-session index rebuilds even
|
||||
* when its cache message lacks `metadata.usage` AND it was transiently dropped
|
||||
* from the cache (e.g. a sibling branch during a regenerate). `buildIndex`
|
||||
* restores from here so branch cost persists across branch switches — the cost
|
||||
* analog of `snapshotsByAnchorFamily` for the breakdown. Cleared on convo switch.
|
||||
*/
|
||||
const usageHistory = new Map<string, Map<string, BranchUsage>>();
|
||||
|
||||
function stickyUsage(conversationId: string, messageId: string): BranchUsage | undefined {
|
||||
return usageHistory.get(conversationId)?.get(messageId);
|
||||
}
|
||||
|
||||
/** Reads the persisted per-response usage rollup off a message's metadata.
|
||||
* The backend already normalized per-event into display units, so this reads
|
||||
* them directly. Absent for user messages and pre-feature responses (they
|
||||
* contribute 0 to branch/total). */
|
||||
function readPersistedUsage(message: Partial<TMessage>): BranchUsage | undefined {
|
||||
const usage = message.metadata?.usage;
|
||||
if (usage == null || typeof usage !== 'object') {
|
||||
return undefined;
|
||||
}
|
||||
const persisted = usage as TResponseUsage;
|
||||
return {
|
||||
input: persisted.input ?? 0,
|
||||
output: persisted.output ?? 0,
|
||||
cacheWrite: persisted.cacheWrite ?? 0,
|
||||
cacheRead: persisted.cacheRead ?? 0,
|
||||
cost: persisted.cost ?? 0,
|
||||
/** Cost is omitted when saved with `contextCost` off — don't render $0.00 */
|
||||
costKnown: typeof persisted.cost === 'number',
|
||||
};
|
||||
}
|
||||
|
||||
/** Pure sum of two usage records — for combining branch/total with pending. */
|
||||
export function mergeUsage(a: BranchUsage, b: BranchUsage): BranchUsage {
|
||||
return {
|
||||
input: a.input + b.input,
|
||||
output: a.output + b.output,
|
||||
cacheWrite: a.cacheWrite + b.cacheWrite,
|
||||
cacheRead: a.cacheRead + b.cacheRead,
|
||||
cost: a.cost + b.cost,
|
||||
/** Coverage stays complete only if BOTH sides are complete */
|
||||
costKnown: a.costKnown && b.costKnown,
|
||||
};
|
||||
}
|
||||
|
||||
/** Accumulates one entry's usage into a running total (in place). */
|
||||
function addUsage(target: BranchUsage, usage?: BranchUsage): void {
|
||||
if (usage == null) {
|
||||
return;
|
||||
}
|
||||
target.input += usage.input;
|
||||
target.output += usage.output;
|
||||
target.cacheWrite += usage.cacheWrite;
|
||||
target.cacheRead += usage.cacheRead;
|
||||
target.cost += usage.cost;
|
||||
/** A usage-bearing entry without a cost breaks complete coverage */
|
||||
if (!usage.costKnown) {
|
||||
target.costKnown = false;
|
||||
}
|
||||
}
|
||||
|
||||
function toEntry(message: Partial<TMessage>): TokenEntry {
|
||||
return {
|
||||
tokenCount: typeof message.tokenCount === 'number' ? message.tokenCount : 0,
|
||||
isCreatedByUser: message.isCreatedByUser === true,
|
||||
parentMessageId: message.parentMessageId ?? null,
|
||||
usage: readPersistedUsage(message),
|
||||
};
|
||||
}
|
||||
|
||||
/** Full O(n) rebuild — only on discrete cache replacements (load, refetch, edits) */
|
||||
/** Full O(n) rebuild — only on discrete cache replacements (load, refetch, edits).
|
||||
* Restores a response's `usage` from the sticky history when the rebuilt message
|
||||
* carries none: per-message usage is immutable and the live finalize flushes it
|
||||
* into the index (`setEntryUsage`) before the persisted `metadata.usage` reaches
|
||||
* the cache, so a mid-session rebuild (e.g. during regenerate) must not wipe an
|
||||
* earlier branch's flushed usage — which would drop its branch cost on switch. */
|
||||
export function buildIndex(conversationId: string, messages?: TMessage[] | null): void {
|
||||
const index = new Map<string, TokenEntry>();
|
||||
if (messages != null) {
|
||||
for (const message of messages) {
|
||||
if (message?.messageId) {
|
||||
index.set(message.messageId, toEntry(message));
|
||||
const entry = toEntry(message);
|
||||
if (entry.usage == null) {
|
||||
entry.usage = stickyUsage(conversationId, message.messageId);
|
||||
}
|
||||
index.set(message.messageId, entry);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -66,7 +170,11 @@ export function upsertEntries(
|
|||
}
|
||||
for (const message of messages) {
|
||||
if (message?.messageId) {
|
||||
index.set(message.messageId, toEntry(message));
|
||||
const entry = toEntry(message);
|
||||
if (entry.usage == null) {
|
||||
entry.usage = stickyUsage(conversationId, message.messageId);
|
||||
}
|
||||
index.set(message.messageId, entry);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -77,15 +185,20 @@ export function migrateIndex(fromId: string, toId: string): void {
|
|||
return;
|
||||
}
|
||||
const index = registry.get(fromId);
|
||||
if (!index) {
|
||||
return;
|
||||
if (index) {
|
||||
registry.delete(fromId);
|
||||
registry.set(toId, index);
|
||||
}
|
||||
const usage = usageHistory.get(fromId);
|
||||
if (usage) {
|
||||
usageHistory.delete(fromId);
|
||||
usageHistory.set(toId, usage);
|
||||
}
|
||||
registry.delete(fromId);
|
||||
registry.set(toId, index);
|
||||
}
|
||||
|
||||
export function clearIndex(conversationId: string): void {
|
||||
registry.delete(conversationId);
|
||||
usageHistory.delete(conversationId);
|
||||
}
|
||||
|
||||
export function hasIndex(conversationId: string): boolean {
|
||||
|
|
@ -104,6 +217,7 @@ export function sumBranch(
|
|||
}
|
||||
|
||||
const totals = { input: 0, output: 0, counted: 0, total: 0, containsAnchor: false };
|
||||
const usage: BranchUsage = { ...EMPTY_USAGE };
|
||||
let currentId: string | null = tailId;
|
||||
let guard = index.size;
|
||||
|
||||
|
|
@ -124,10 +238,50 @@ export function sumBranch(
|
|||
totals.output += entry.tokenCount;
|
||||
}
|
||||
}
|
||||
addUsage(usage, entry.usage);
|
||||
currentId = entry.parentMessageId;
|
||||
}
|
||||
|
||||
return { ...totals, tailId };
|
||||
return { ...totals, tailId, usage };
|
||||
}
|
||||
|
||||
/**
|
||||
* Sums provider usage/cost across EVERY message in the conversation (all
|
||||
* branches, including regenerated/abandoned responses) — the conversation
|
||||
* total, shown alongside the branch figure when they differ.
|
||||
*/
|
||||
export function sumTotalUsage(conversationId: string): BranchUsage {
|
||||
const usage: BranchUsage = { ...EMPTY_USAGE };
|
||||
const index = registry.get(conversationId);
|
||||
if (!index) {
|
||||
return usage;
|
||||
}
|
||||
for (const entry of index.values()) {
|
||||
addUsage(usage, entry.usage);
|
||||
}
|
||||
return usage;
|
||||
}
|
||||
|
||||
/**
|
||||
* Attaches a response's usage to its index entry AND the sticky usage history.
|
||||
* Used by the live finalize/stop flush. The history copy lets `buildIndex`
|
||||
* restore the usage after a rebuild (the persisted `metadata.usage` reaches
|
||||
* `buildIndex` on reload; the sticky copy covers the in-session window before
|
||||
* that, including a sibling transiently dropped from the cache on regenerate).
|
||||
*/
|
||||
export function setEntryUsage(conversationId: string, messageId: string, usage: BranchUsage): void {
|
||||
/** Remember it durably first so a later rebuild — or a transient cache drop
|
||||
* during regenerate — can restore it even when the entry isn't present yet. */
|
||||
let history = usageHistory.get(conversationId);
|
||||
if (!history) {
|
||||
history = new Map<string, BranchUsage>();
|
||||
usageHistory.set(conversationId, history);
|
||||
}
|
||||
history.set(messageId, usage);
|
||||
const entry = registry.get(conversationId)?.get(messageId);
|
||||
if (entry) {
|
||||
entry.usage = usage;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
|
|||
|
|
@ -74,17 +74,26 @@ test.describe('context usage gauge', () => {
|
|||
await expect(usageSection.getByText('Output', { exact: true })).toBeVisible();
|
||||
|
||||
/** Cost row: interface.contextCost is enabled in the harness yaml, the
|
||||
* token-config endpoint prices mock models at the default rate, and
|
||||
* the fake model emits usage — so a $ value must render */
|
||||
await expect(popover.getByText('Session cost')).toBeVisible();
|
||||
await expect(popover.getByText(/\$\d|<\$0\.01/)).toBeVisible();
|
||||
* token-config endpoint prices mock models at the default rate, and the
|
||||
* fake model emits usage — so a $ value must render. A single (unbranched)
|
||||
* conversation shows only the branch cost, no all-branches total line. */
|
||||
const costSection = popover.getByTestId('token-usage-cost');
|
||||
await expect(costSection).toBeVisible();
|
||||
await expect(costSection.getByText(/\$\d|<\$0\.01/)).toBeVisible();
|
||||
await expect(costSection.getByText('All branches')).toHaveCount(0);
|
||||
await page.keyboard.press('Escape');
|
||||
|
||||
/** Fallback path: after reload there is no snapshot — the gauge rebuilds
|
||||
* from per-message tokenCount history returned by the messages query */
|
||||
/** Persistence (Parts A + B): after reload the breakdown rehydrates from
|
||||
* the response message's metadata.contextUsage + metadata.usage — the
|
||||
* granular rows AND the branch cost survive without generating a turn. */
|
||||
await page.reload({ timeout: 15000 });
|
||||
await expect(mockReply(page)).toBeVisible({ timeout: 20000 });
|
||||
await expectGaugeAboveZero(page);
|
||||
const reloaded = await openBreakdown(page);
|
||||
await expect(reloaded.getByTestId('context-breakdown')).toBeVisible({ timeout: 10000 });
|
||||
const reloadedCost = reloaded.getByTestId('token-usage-cost');
|
||||
await expect(reloadedCost).toBeVisible();
|
||||
await expect(reloadedCost.getByText(/\$\d|<\$0\.01/)).toBeVisible();
|
||||
});
|
||||
|
||||
test('renders the granular breakdown from the live context snapshot', async ({ page }) => {
|
||||
|
|
@ -99,6 +108,48 @@ test.describe('context usage gauge', () => {
|
|||
await expectGranular(page);
|
||||
});
|
||||
|
||||
test('shows branch cost with an all-branches total after regenerating', async ({ page }) => {
|
||||
test.setTimeout(150000);
|
||||
await page.goto(NEW_CHAT_PATH, { timeout: 10000 });
|
||||
await selectMockEndpoint(page, MOCK_ENDPOINTS[0]);
|
||||
|
||||
await sendAndAwaitReply(page, 'hello');
|
||||
|
||||
/** Single branch: branch cost == total, so no all-branches line renders. */
|
||||
let popover = await openBreakdown(page);
|
||||
await expect(popover.getByTestId('token-usage-cost')).toBeVisible();
|
||||
await expect(popover.getByText('All branches')).toHaveCount(0);
|
||||
await page.keyboard.press('Escape');
|
||||
|
||||
/** Regenerate to create a sibling branch (B). */
|
||||
const assistantMessage = messagesView(page).locator('.message-render').nth(1);
|
||||
await assistantMessage.hover();
|
||||
const regenerateButton = assistantMessage.locator('button[title="Regenerate"]').last();
|
||||
await expect(regenerateButton).toBeVisible();
|
||||
const [regen] = await Promise.all([
|
||||
page.waitForResponse(isAgentsStream, { timeout: 30000 }),
|
||||
regenerateButton.click(),
|
||||
]);
|
||||
expect(regen.ok()).toBeTruthy();
|
||||
await expect(page.getByText('2 / 2')).toBeVisible({ timeout: 20000 });
|
||||
|
||||
/** Branch cost is shown live for the regenerated branch. */
|
||||
popover = await openBreakdown(page);
|
||||
await expect(popover.getByTestId('token-usage-cost')).toBeVisible();
|
||||
await page.keyboard.press('Escape');
|
||||
|
||||
/** After reload both branches rehydrate from persisted metadata.usage, so
|
||||
* the cost is branch-scoped and a muted all-branches total appears (it
|
||||
* exceeds the single viewed branch). */
|
||||
await page.reload({ timeout: 15000 });
|
||||
await expect(mockReply(page)).toBeVisible({ timeout: 20000 });
|
||||
const reloaded = await openBreakdown(page);
|
||||
const costSection = reloaded.getByTestId('token-usage-cost');
|
||||
await expect(costSection).toBeVisible();
|
||||
await expect(costSection.getByText('Cost (this branch)')).toBeVisible();
|
||||
await expect(costSection.getByText('All branches')).toBeVisible();
|
||||
});
|
||||
|
||||
test('preserves the granular breakdown after switching branches', async ({ page }) => {
|
||||
test.setTimeout(150000);
|
||||
await page.goto(NEW_CHAT_PATH, { timeout: 10000 });
|
||||
|
|
@ -130,5 +181,15 @@ test.describe('context usage gauge', () => {
|
|||
await expect(page.getByText('1 / 2')).toBeVisible({ timeout: 10000 });
|
||||
|
||||
await expectGranular(page);
|
||||
|
||||
/** Branch cost must also survive the switch (live, no reload): branch A's
|
||||
* flushed usage is restored from the sticky usage history even though its
|
||||
* cache message lacks metadata.usage and B's regenerate dropped it. */
|
||||
const popover = page.getByRole('region', { name: 'Context usage' });
|
||||
const costSection = popover.getByTestId('token-usage-cost');
|
||||
await expect(costSection).toBeVisible();
|
||||
/** Branch A also has siblings, so both a branch-cost row and an all-branches
|
||||
* total render — assert at least one cost value is present. */
|
||||
await expect(costSection.getByText(/\$\d|<\$0\.01/).first()).toBeVisible();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -78,6 +78,28 @@ export type TContextUsageEvent = {
|
|||
/** Tokens still free after instructions + pruned messages */
|
||||
remainingContextTokens?: number;
|
||||
calibrationRatio?: number;
|
||||
/** Output tokens of the response's final model call (the call this pre-invoke
|
||||
* snapshot precedes). Populated only on the persisted `metadata.contextUsage`
|
||||
* blob so a reloaded multi-call turn adds the same post-snapshot delta the
|
||||
* live finalizer did — not the full response `tokenCount`, which the snapshot
|
||||
* already includes for earlier steps. */
|
||||
completedOutputTokens?: number;
|
||||
};
|
||||
|
||||
/**
|
||||
* Per-response usage rollup persisted on `responseMessage.metadata.usage`, in
|
||||
* display units (input excludes cache; output includes repaired completion).
|
||||
* Normalized per-event on the backend before summing so a reloaded conversation
|
||||
* reproduces the live branch/total usage exactly, even for mixed-provider turns
|
||||
* (summarization/subagent calls on a different provider than the primary).
|
||||
*/
|
||||
export type TResponseUsage = {
|
||||
input: number;
|
||||
output: number;
|
||||
cacheWrite: number;
|
||||
cacheRead: number;
|
||||
/** Authoritative USD cost; present only when `interface.contextCost` was on at save */
|
||||
cost?: number;
|
||||
};
|
||||
|
||||
/** Provider-reported usage for a single completed model call. */
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue