diff --git a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js index 0668d0b6c1..5764c159e9 100644 --- a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js +++ b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js @@ -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'); diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index 0b543f6528..e23ffa64bd 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -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} [options.usageEmitSink] - Array collecting each emitted + * `on_token_usage` payload (incl. cost) so the response's usage rollup can be persisted. * @returns {Record} 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; + } } }, }; diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 4ba5fafde3..6c3d3d9c61 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -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 | undefined} */ + this.usageEmitSink = usageEmitSink; /** @type {MessageContentComplex[]} */ this.contentParts = contentParts; /** @type {Array} */ @@ -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, + * contextUsage?: import('librechat-data-provider').TContextUsageEvent, + * usage?: import('librechat-data-provider').TResponseUsage, + * } | undefined} + */ + buildResponseMetadata() { + /** @type {{ + * thoughtSignatures?: Record, + * 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 diff --git a/api/server/middleware/abortMiddleware.js b/api/server/middleware/abortMiddleware.js index e0c5ae0ff0..8b339b05ba 100644 --- a/api/server/middleware/abortMiddleware.js +++ b/api/server/middleware/abortMiddleware.js @@ -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({ diff --git a/api/server/routes/agents/index.js b/api/server/routes/agents/index.js index db5976bb38..145a6c0316 100644 --- a/api/server/routes/agents/index.js +++ b/api/server/routes/agents/index.js @@ -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( { diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index 7330f22916..f72ef653e9 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -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} */ + 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) { diff --git a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx index 9b01a7f6b8..d17c01334b 100644 --- a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx +++ b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx @@ -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) { )} - {usageTotals.eventCount > 0 && ( + {hasUsage && ( <>
- - - {usageTotals.cacheRead > 0 && ( - + + + {branchUsage.cacheRead > 0 && ( + )} - {usageTotals.cacheWrite > 0 && ( - + {branchUsage.cacheWrite > 0 && ( + )}
)} - {showCost && view.costUSD != null && ( + {showCost && hasUsage && branchUsage.costKnown && ( <>
-
- {localize('com_ui_session_cost')} - {formatCost(view.costUSD)} +
+
+ + {showTotal + ? localize('com_ui_context_cost_branch') + : localize('com_ui_context_cost')} + + {formatCost(view.branchCost)} +
+ {showTotal && ( +
+ {localize('com_ui_context_cost_total')} + {formatCost(view.totalCost)} +
+ )}
)} diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts index 90bb4a25e2..44fe7d592d 100644 --- a/client/src/hooks/Chat/useTokenUsage.ts +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -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([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([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([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, ]); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 0fa9125e73..160efb81b1 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/client/src/hooks/SSE/useResumableSSE.ts @@ -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 diff --git a/client/src/hooks/SSE/useSSE.ts b/client/src/hooks/SSE/useSSE.ts index 745b268417..8f18f1d556 100644 --- a/client/src/hooks/SSE/useSSE.ts +++ b/client/src/hooks/SSE/useSSE.ts @@ -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 ?? diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index 9ae51d6920..152890252e 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -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, }; diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 3f9ea230fc..89049a4d1b 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -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.", diff --git a/client/src/store/usage.ts b/client/src/store/usage.ts index 6501648a72..a4bd064a58 100644 --- a/client/src/store/usage.ts +++ b/client/src/store/usage.ts @@ -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>(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(EMPTY_USAGE_TOTALS), ); +/** Provider usage/cost summed across all branches of the conversation. */ +export const totalUsageFamily = atomFamily((_conversationId: string) => + atom(EMPTY_USAGE), +); + /** Throttled in-flight output token estimate for the current model call */ export const liveTokensFamily = atomFamily((_conversationId: string) => atom(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 | 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); + } +} diff --git a/client/src/utils/tokens.spec.ts b/client/src/utils/tokens.spec.ts index f25cda84ec..5ce386b73c 100644 --- a/client/src/utils/tokens.spec.ts +++ b/client/src/utils/tokens.spec.ts @@ -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); + }); +}); diff --git a/client/src/utils/tokens.ts b/client/src/utils/tokens.ts index 9e05b122e8..c5abff26c0 100644 --- a/client/src/utils/tokens.ts +++ b/client/src/utils/tokens.ts @@ -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>(); +/** + * 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>(); + +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): 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): 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(); 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(); + usageHistory.set(conversationId, history); + } + history.set(messageId, usage); + const entry = registry.get(conversationId)?.get(messageId); + if (entry) { + entry.usage = usage; + } } /** diff --git a/e2e/specs/mock/usage.spec.ts b/e2e/specs/mock/usage.spec.ts index ed376e4e04..ab7e5bc4c8 100644 --- a/e2e/specs/mock/usage.spec.ts +++ b/e2e/specs/mock/usage.spec.ts @@ -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(); }); }); diff --git a/packages/api/src/agents/usage.spec.ts b/packages/api/src/agents/usage.spec.ts index 177063b95e..8f4bd96b7e 100644 --- a/packages/api/src/agents/usage.spec.ts +++ b/packages/api/src/agents/usage.spec.ts @@ -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(); + }); +}); diff --git a/packages/api/src/agents/usage.ts b/packages/api/src/agents/usage.ts index d6252841a7..43ff590a86 100644 --- a/packages/api/src/agents/usage.ts +++ b/packages/api/src/agents/usage.ts @@ -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, +): 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): 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 = [], +): TContextUsageEvent { + const { breakdown } = snapshot; + const completedOutputTokens = finalCallOutputTokens(usageEvents); + let toolTokenCounts = breakdown.toolTokenCounts; + if (toolTokenCounts != null) { + const trimmed: Record = {}; + 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; diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index 71453a4534..227154f4c2 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -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. */