From f513665ec5a0059fb17e27d18573699f42d220f1 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Thu, 11 Jun 2026 15:36:21 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20fix:=20Address=20Usage=20Review?= =?UTF-8?q?=20Findings=20=E2=80=94=20Cost=20Timing,=20Scoped=20Caches,=20F?= =?UTF-8?q?inalized=20Output?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - carry the post-snapshot output estimate into the context snapshot at finalize so the gauge keeps the last response after live resets - accumulate per-rate billable units and price the session cost at render, so usage events arriving before the token-config load still count once it resolves - pass user-scoped token-config cache keys through loadConfigModels fetches and drop the controller's unscoped fallback to prevent serving another user's resolved config - tag emitted usage events with a per-run seq so resume dedupe never drops a distinct call with an identical payload - admit the static tokenConfig override in the custom endpoint schema so it survives zod parsing into req.config --- .../controllers/TokenConfigController.js | 8 +-- api/server/controllers/agents/callbacks.js | 3 ++ client/src/hooks/Chat/useTokenUsage.ts | 29 +++++++++-- client/src/hooks/SSE/useUsageHandler.ts | 52 ++++++++++++++----- client/src/store/usage.ts | 16 ++++-- client/src/utils/tokens.spec.ts | 29 +++++++++++ client/src/utils/tokens.ts | 52 +++++++++++++------ packages/api/src/endpoints/config/models.ts | 6 +++ .../api/src/endpoints/custom/initialize.ts | 5 +- packages/data-provider/src/config.ts | 10 ++++ packages/data-provider/src/types/runs.ts | 2 + 11 files changed, 169 insertions(+), 43 deletions(-) diff --git a/api/server/controllers/TokenConfigController.js b/api/server/controllers/TokenConfigController.js index 0014485ea8..f6cf7ac554 100644 --- a/api/server/controllers/TokenConfigController.js +++ b/api/server/controllers/TokenConfigController.js @@ -30,11 +30,11 @@ async function tokenConfigController(req, res) { endpointTokenConfigs[name] = endpointConfig.tokenConfig; continue; } - /** The models-config fetch path stores under the plain endpoint name; - * chat initialization stores under the user-scoped key — accept either */ + /** Model fetches and chat initialization both store under this key — + * user-scoped whenever the fetched config can be user-specific, so a + * plain-name fallback would risk serving another user's entry */ const tokenKey = getTokenConfigKey(endpointConfig, name, req.user.id); - const cached = - (await cache.get(tokenKey)) ?? (tokenKey !== name ? await cache.get(name) : undefined); + const cached = await cache.get(tokenKey); if (cached) { endpointTokenConfigs[name] = cached; } diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js index cd74631161..a6efa8a2b8 100644 --- a/api/server/controllers/agents/callbacks.js +++ b/api/server/controllers/agents/callbacks.js @@ -132,6 +132,9 @@ class ModelEndHandler { provider: taggedUsage.provider, usage_type: taggedUsage.usage_type, runId: metadata?.run_id, + /** Per-run sequence so identical payloads from distinct calls + * stay distinguishable during resume dedupe */ + seq: this.collectedUsage.length, }); } diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts index 6ca1a4782f..6761559711 100644 --- a/client/src/hooks/Chat/useTokenUsage.ts +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -12,7 +12,8 @@ import { contextSnapshotFamily, } from '~/store/usage'; import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage'; -import { buildIndex, sumBranch } from '~/utils'; +import { buildIndex, sumBranch, costFromUnits } from '~/utils'; +import { useTokenConfigQuery } from '~/data-provider'; import useTokenLimits from './useTokenLimits'; export interface TokenUsageParams { @@ -57,6 +58,22 @@ export default function useTokenUsage({ const liveTokens = useAtomValue(liveTokensFamily(conversationKey)); const setBranchTotals = useSetAtom(branchTotalsFamily(conversationKey)); const limits = useTokenLimits(conversation); + const { data: tokenConfig } = useTokenConfigQuery(); + + /** Priced at render so events folded before the token-config load still count */ + const costUSD = useMemo(() => { + if (usageTotals.eventCount === 0) { + return undefined; + } + let total = 0; + for (const bucket of Object.values(usageTotals.byRate)) { + const rates = + (bucket.provider != null ? tokenConfig?.[bucket.provider]?.[bucket.model] : undefined) ?? + tokenConfig?.[bucket.endpoint]?.[bucket.model]; + total += costFromUnits(bucket, rates); + } + return total; + }, [usageTotals, tokenConfig]); const isSubmittingRef = useRef(isSubmitting); isSubmittingRef.current = isSubmitting; @@ -115,7 +132,9 @@ export default function useTokenUsage({ snapshot.remainingContextTokens != null ? maxTokens - snapshot.remainingContextTokens : instructionTokens + breakdown.messageTokens; - const usedTokens = Math.max(0, baseUsed) + liveTokens; + /** The snapshot is pre-invoke: in-flight output rides on `liveTokens`, + * and the last call's finalized output on `completedOutputTokens` */ + const usedTokens = Math.max(0, baseUsed) + liveTokens + (snapshot.completedOutputTokens ?? 0); return { usedTokens, maxTokens, @@ -127,7 +146,7 @@ export default function useTokenUsage({ usageTotals, liveTokens, rates: limits.rates, - costUSD: usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined, + costUSD, }; } @@ -145,7 +164,7 @@ export default function useTokenUsage({ usageTotals, liveTokens, rates: limits.rates, - costUSD: usageTotals.eventCount > 0 ? usageTotals.costUSD : undefined, + costUSD, }; - }, [snapshot, isSubmitting, branchTotals, usageTotals, liveTokens, limits]); + }, [snapshot, isSubmitting, branchTotals, usageTotals, liveTokens, limits, costUSD]); } diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index 099e0f88a6..73b418fa9f 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -1,11 +1,9 @@ import { useRef, useMemo } from 'react'; import { getDefaultStore } from 'jotai'; -import { useQueryClient } from '@tanstack/react-query'; -import { Constants, QueryKeys } from 'librechat-data-provider'; +import { Constants } from 'librechat-data-provider'; import type { TMessage, TConversation, - TTokenConfigMap, TTokenUsageEvent, TContextUsageEvent, } from 'librechat-data-provider'; @@ -18,7 +16,13 @@ import { removeUsageAtoms, contextSnapshotFamily, } from '~/store/usage'; -import { sumBranch, upsertEntries, migrateIndex, calcUsageCost, estimateTokens } from '~/utils'; +import { + sumBranch, + upsertEntries, + migrateIndex, + estimateTokens, + normalizeUsageUnits, +} from '~/utils'; const FLUSH_INTERVAL_MS = 250; @@ -72,7 +76,6 @@ function countDeltaChars(content: unknown): number { * stable and never cause re-renders themselves. */ export default function useUsageHandler(): UsageHandlers { - const queryClient = useQueryClient(); /** Streamed chars since the last snapshot or model end (current call only) */ const streamCharsRef = useRef(0); /** Provider-confirmed output tokens since the last snapshot (current run) */ @@ -102,23 +105,34 @@ export default function useUsageHandler(): UsageHandlers { const foldUsage = (data: TTokenUsageEvent, submission: UsageSubmissionLike) => { const convoKey = getConvoKey(submission); - const tokenConfig = queryClient.getQueryData([QueryKeys.tokenConfig]); const endpoint = submission.conversation?.endpoint ?? ''; const model = data.model ?? submission.conversation?.model ?? ''; - /** Agent runs report the underlying provider, where rates are keyed */ - const rates = - (data.provider != null ? tokenConfig?.[data.provider]?.[model] : undefined) ?? - tokenConfig?.[endpoint]?.[model]; + const units = normalizeUsageUnits(data); + /** Cost is priced at render from these buckets — never baked in here, + * so events arriving before the token config load still get priced */ + const bucketKey = `${data.provider ?? ''}|${endpoint}|${model}`; const totalsAtom = usageTotalsFamily(convoKey); const prev = jotai.get(totalsAtom); + const bucket = prev.byRate[bucketKey]; jotai.set(totalsAtom, { input: prev.input + (data.input_tokens ?? 0), output: prev.output + (data.output_tokens ?? 0), cacheWrite: prev.cacheWrite + (data.input_token_details?.cache_creation ?? 0), cacheRead: prev.cacheRead + (data.input_token_details?.cache_read ?? 0), - costUSD: prev.costUSD + calcUsageCost(data, rates), eventCount: prev.eventCount + 1, + byRate: { + ...prev.byRate, + [bucketKey]: bucket + ? { + ...bucket, + input: bucket.input + units.input, + output: bucket.output + units.output, + cacheWrite: bucket.cacheWrite + units.cacheWrite, + cacheRead: bucket.cacheRead + units.cacheRead, + } + : { provider: data.provider, endpoint, model, ...units }, + }, }); }; @@ -177,6 +191,10 @@ export default function useUsageHandler(): UsageHandlers { const finalizeUsage: UsageHandlers['finalizeUsage'] = (data, submission) => { const fromKey = getConvoKey(submission); const realId = data.conversation?.conversationId ?? fromKey; + /** From the refs, not the atom — the throttle may not have flushed */ + const liveAtFinalize = + confirmedRef.current + + estimateTokens(streamCharsRef.current, jotai.get(calibrationFamily(fromKey))); upsertEntries(fromKey, [data.requestMessage, data.responseMessage]); @@ -194,11 +212,19 @@ export default function useUsageHandler(): UsageHandlers { } const tailId = data.responseMessage?.messageId ?? data.requestMessage?.messageId ?? null; + const anchorId = submission.userMessage?.messageId ?? null; if (tailId) { - const anchorId = submission.userMessage?.messageId ?? null; jotai.set(branchTotalsFamily(realId), sumBranch(realId, tailId, anchorId)); } + /** The snapshot was taken pre-invoke — carry the output streamed since + * then so the gauge keeps the final response after live resets */ + const snapshotAtom = contextSnapshotFamily(realId); + const snapshot = jotai.get(snapshotAtom); + if (snapshot != null && liveAtFinalize > 0 && snapshot.anchorMessageId === anchorId) { + jotai.set(snapshotAtom, { ...snapshot, completedOutputTokens: liveAtFinalize }); + } + streamCharsRef.current = 0; confirmedRef.current = 0; setLive(realId, 0); @@ -213,5 +239,5 @@ export default function useUsageHandler(): UsageHandlers { backfillUsage, seedLive, }; - }, [queryClient]); + }, []); } diff --git a/client/src/store/usage.ts b/client/src/store/usage.ts index c122efa20d..280c1975da 100644 --- a/client/src/store/usage.ts +++ b/client/src/store/usage.ts @@ -1,12 +1,21 @@ import { atom } from 'jotai'; import { atomFamily } from 'jotai/utils'; import type { TContextUsageEvent } from 'librechat-data-provider'; -import type { BranchTotals } from '~/utils/tokens'; +import type { BranchTotals, CostUnits } from '~/utils/tokens'; import { EMPTY_BRANCH } from '~/utils/tokens'; /** Latest backend context snapshot, anchored to the run's user message for staleness checks */ export interface ContextSnapshot extends TContextUsageEvent { anchorMessageId: string | null; + /** Output tokens finalized after this pre-call snapshot (the last call's response) */ + completedOutputTokens?: number; +} + +/** Billable units accumulated per pricing lookup (provider/endpoint + model) */ +export interface RateBucket extends CostUnits { + provider?: string; + endpoint: string; + model: string; } /** Cumulative provider-reported usage for the conversation's current session */ @@ -15,8 +24,9 @@ export interface UsageTotals { output: number; cacheWrite: number; cacheRead: number; - costUSD: number; eventCount: number; + /** Cost is derived from these at render so a late token-config load still prices every event */ + byRate: Record; } export const EMPTY_USAGE_TOTALS: UsageTotals = { @@ -24,8 +34,8 @@ export const EMPTY_USAGE_TOTALS: UsageTotals = { output: 0, cacheWrite: 0, cacheRead: 0, - costUSD: 0, eventCount: 0, + byRate: {}, }; /** diff --git a/client/src/utils/tokens.spec.ts b/client/src/utils/tokens.spec.ts index 1ebee048f2..b6135f9325 100644 --- a/client/src/utils/tokens.spec.ts +++ b/client/src/utils/tokens.spec.ts @@ -9,6 +9,8 @@ import { sumBranch, estimateTokens, calcUsageCost, + costFromUnits, + normalizeUsageUnits, formatCost, groupToolTokens, countTrailingOutputChars, @@ -145,6 +147,33 @@ describe('calcUsageCost', () => { expect(calcUsageCost({ input_tokens: 100, output_tokens: 100 })).toBe(0); expect(calcUsageCost({ input_tokens: 100 }, { context: 1000 })).toBe(0); }); + + it('prices summed normalized units identically to per-event costs', () => { + const anthropicEvent = { + input_tokens: 1000, + output_tokens: 500, + input_token_details: { cache_creation: 2000, cache_read: 10000 }, + }; + const openAIEvent = { + input_tokens: 10000, + output_tokens: 500, + input_token_details: { cache_read: 4000 }, + }; + const perEvent = calcUsageCost(anthropicEvent, rates) + calcUsageCost(openAIEvent, rates); + + const a = normalizeUsageUnits(anthropicEvent); + const b = normalizeUsageUnits(openAIEvent); + const summed = costFromUnits( + { + input: a.input + b.input, + output: a.output + b.output, + cacheWrite: a.cacheWrite + b.cacheWrite, + cacheRead: a.cacheRead + b.cacheRead, + }, + rates, + ); + expect(summed).toBeCloseTo(perEvent); + }); }); describe('formatCost', () => { diff --git a/client/src/utils/tokens.ts b/client/src/utils/tokens.ts index 36e8d6437d..69574ad4ed 100644 --- a/client/src/utils/tokens.ts +++ b/client/src/utils/tokens.ts @@ -233,34 +233,56 @@ export function estimateTokens(charCount: number, calibrationRatio = 1): number return Math.round((charCount / 4) * ratio); } +/** Billable token quantities of one or more model calls, normalized for pricing */ +export interface CostUnits { + input: number; + output: number; + cacheWrite: number; + cacheRead: number; +} + /** - * USD cost of one model call. Mirrors the cache-detection heuristic used in - * token accounting: cache counts are additive when they exceed base input - * (Anthropic), otherwise cache reads are included in input tokens (OpenAI). + * Normalizes one call's usage into billable units. Mirrors the + * cache-detection heuristic used in token accounting: cache counts are + * additive when they exceed base input (Anthropic), otherwise cache reads + * are included in input tokens (OpenAI) — applied per event so units stay + * correct when summed across calls. */ -export function calcUsageCost(usage: TTokenUsageEvent, rates?: TModelTokenomics): number { - if (!rates || rates.prompt == null || rates.completion == null) { - return 0; - } +export function normalizeUsageUnits(usage: TTokenUsageEvent): CostUnits { const input = usage.input_tokens ?? 0; const output = usage.output_tokens ?? 0; const cacheWrite = usage.input_token_details?.cache_creation ?? 0; const cacheRead = usage.input_token_details?.cache_read ?? 0; + const cacheIsAdditive = cacheWrite + cacheRead > input; + return { + input: cacheIsAdditive ? input : Math.max(0, input - cacheRead - cacheWrite), + output, + cacheWrite, + cacheRead, + }; +} + +/** USD cost of normalized usage units at the given per-million-token rates */ +export function costFromUnits(units: CostUnits, rates?: TModelTokenomics): number { + if (!rates || rates.prompt == null || rates.completion == null) { + return 0; + } const writeRate = rates.cacheWrite ?? rates.prompt; const readRate = rates.cacheRead ?? rates.prompt; - - const cacheIsAdditive = cacheWrite + cacheRead > input; - const baseInput = cacheIsAdditive ? input : Math.max(0, input - cacheRead - cacheWrite); - return ( - (baseInput * rates.prompt + - cacheWrite * writeRate + - cacheRead * readRate + - output * rates.completion) / + (units.input * rates.prompt + + units.cacheWrite * writeRate + + units.cacheRead * readRate + + units.output * rates.completion) / 1e6 ); } +/** USD cost of one model call */ +export function calcUsageCost(usage: TTokenUsageEvent, rates?: TModelTokenomics): number { + return costFromUnits(normalizeUsageUnits(usage), rates); +} + export function formatTokens(count: number): string { const formatted = new Intl.NumberFormat(undefined, { notation: 'compact', diff --git a/packages/api/src/endpoints/config/models.ts b/packages/api/src/endpoints/config/models.ts index 6bd47d10ca..28ce2a3d81 100644 --- a/packages/api/src/endpoints/config/models.ts +++ b/packages/api/src/endpoints/config/models.ts @@ -11,6 +11,7 @@ import type { AppConfig } from '@librechat/data-schemas'; import type { ServerRequest, GetUserKeyValuesFunction, UserKeyValues } from '~/types'; import type { FetchModelsParams } from '~/endpoints/models'; import { fetchModels as defaultFetchModels } from '~/endpoints/models'; +import { getTokenConfigKey } from '~/endpoints/custom/initialize'; import { isUserProvided } from '~/utils'; /** @@ -185,6 +186,9 @@ export function createLoadConfigModels(deps: LoadConfigModelsDeps) { headers: endpointHeaders, direct: endpoint.directEndpoint, userIdQuery: models.userIdQuery, + /** User-scoped when configured headers resolve per user — the + * derived token config must not be cached under the shared name */ + tokenKey: getTokenConfigKey(endpoint, name, req.user?.id ?? ''), }); uniqueKeyToEndpointsMap[uniqueKey] = uniqueKeyToEndpointsMap[uniqueKey] || []; uniqueKeyToEndpointsMap[uniqueKey].push(name); @@ -216,6 +220,8 @@ export function createLoadConfigModels(deps: LoadConfigModelsDeps) { direct: endpoint.directEndpoint, userIdQuery: models.userIdQuery, skipCache: true, + /** Fetched with the user's key/URL — always user-scoped */ + tokenKey: getTokenConfigKey(endpoint, name, req.user?.id ?? ''), }); uniqueKeyToEndpointsMap[userFetchKey] = uniqueKeyToEndpointsMap[userFetchKey] || []; uniqueKeyToEndpointsMap[userFetchKey].push(name); diff --git a/packages/api/src/endpoints/custom/initialize.ts b/packages/api/src/endpoints/custom/initialize.ts index 80d0de430d..9a983e5a0f 100644 --- a/packages/api/src/endpoints/custom/initialize.ts +++ b/packages/api/src/endpoints/custom/initialize.ts @@ -27,7 +27,7 @@ export function getTokenConfigKey( endpoint: string, userId: string, ): string { - const hasTokenConfig = (endpointConfig as Record).tokenConfig != null; + const hasTokenConfig = endpointConfig.tokenConfig != null; const userProvidesKey = isUserProvided(extractEnvVariable(endpointConfig.apiKey ?? '')); const userProvidesURL = isUserProvided(extractEnvVariable(endpointConfig.baseURL ?? '')); const willForwardUserScopedHeaders = !!endpointConfig?.headers && !userProvidesURL; @@ -156,8 +156,7 @@ export async function initializeCustom({ const userId = req.user?.id ?? ''; const cache = tokenConfigCache(); - /** tokenConfig is an optional extended property on custom endpoints */ - const hasTokenConfig = (endpointConfig as Record).tokenConfig != null; + const hasTokenConfig = endpointConfig.tokenConfig != null; const tokenKey = getTokenConfigKey(endpointConfig, endpoint, userId); const cachedConfig = diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 5b50388dfd..11d2fe7890 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -846,6 +846,16 @@ export const endpointSchema = baseEndpointSchema.merge( .optional(), directEndpoint: z.boolean().optional(), titleMessageRole: z.enum(['system', 'user', 'assistant']).optional(), + /** Static per-model token config: context window and per-million-token rates */ + tokenConfig: z + .record( + z.object({ + prompt: z.number(), + completion: z.number(), + context: z.number(), + }), + ) + .optional(), }), ); diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index c68fab94b1..bbbe4a4825 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -93,6 +93,8 @@ export type TTokenUsageEvent = { provider?: string; usage_type?: 'summarization' | 'subagent'; runId?: string; + /** Per-run emission sequence; keeps identical payloads from distinct model calls unique */ + seq?: number; }; /** Lifecycle phase carried on subagent-progress envelopes (mirrors SDK SubagentUpdatePhase). */