🩹 fix: Address Usage Review Findings — Cost Timing, Scoped Caches, Finalized Output

- 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

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;
       }

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,
         });
       }
 

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]);
 }

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<TTokenConfigMap>([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]);
+  }, []);
 }

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<string, RateBucket>;
 }
 
 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: {},
 };
 
 /**

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', () => {

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',

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);

packages/api/src/endpoints/custom/initialize.ts

@@ -27,7 +27,7 @@ export function getTokenConfigKey(
   endpoint: string,
   userId: string,
 ): string {
-  const hasTokenConfig = (endpointConfig as Record<string, unknown>).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<string, unknown>).tokenConfig != null;
+  const hasTokenConfig = endpointConfig.tokenConfig != null;
   const tokenKey = getTokenConfigKey(endpointConfig, endpoint, userId);
 
   const cachedConfig =

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(),
   }),
 );
 

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). */