🪙 refactor: Reconcile Context Gauge to Actual Provider Tokens (#13780)

* 🪙 fix: Reconcile Context Gauge to Actual Provider Tokens

The context gauge could read several× too high (e.g. 213K when the real prompt
was 56K) and stay there across reloads. Root cause: the SDK's calibrationRatio is
`cumulativeProviderReported / cumulativeRawSent`, but a provider's server-side
web search injects large fetched content into the prompt that the SDK never sent
or counted — pinning the ratio at its cap (5) and multiplying every later message
estimate, including post-summary ones. The gauge rendered (and persisted) that
inflated estimate, never the provider's actual token count.

Fix: reconcile the snapshot to the call's ACTUAL prompt tokens (input + cache),
which already arrive in on_token_usage. Only messageTokens is calibration-scaled
(instructions/summary are raw tiktoken), so keep those and set messageTokens to
the remainder, recomputing free space. Shared `promptTokensFromUsage` +
`reconcileContextUsage` in data-provider; applied server-side in
buildPersistedContextUsage (reload-stable) and client-side in useUsageHandler on
each primary usage (corrects at turn-end, no follow-up needed). Also drop the
summary double-count from the Breakdown Messages row.

Deferred (separate agents PR): the SDK over-calibration also fires summarization
prematurely; fixing it needs decoupling real-content estimation from server-side
injection headroom without weakening pruning-overflow safety.

* 🪙 fix: Harden Token Reconciliation for Provider-less + Resume Paths

Codex review on the reconciliation:
- promptTokensFromUsage: when the provider is absent (custom/OpenAI-compatible
  payloads), fall back to the same magnitude heuristic normalizeUsageUnits uses
  (cache ≤ input ⇒ already included) so cached events aren't re-inflated.
- Resume: backfillUsage restores a primary call's usage without replaying a live
  on_token_usage (Redis mode), so the live reconcile never ran and a reconnected
  session stayed on the inflated estimate. New reconcileBackfill reconciles the
  restored snapshot from the final primary call after contextHandler installs it.

* 🪙 fix: Reconcile Resume Snapshot Server-Side, Not via Backfill

Codex: the client reconcileBackfill scanned the resumed run's collectedUsage and
applied the final primary to the latest snapshot — but on a mid-call resume that
usage belongs to an EARLIER call, corrupting the restored gauge.

Move the resume reconciliation server-side: GenerationJobManager.persistTokenUsage
reconciles the stored contextUsage to a primary usage's actual prompt tokens as it
arrives. That usage is the post-invoke truth for the call the latest stored
snapshot precedes (no snapshot is captured between a call's pre-invoke dispatch
and its usage), so it's correct by construction and run-matched. A mid-call resume
(no usage yet) keeps the raw snapshot instead of mis-applying an earlier call's
tokens; it reconciles once the call completes. Removed client reconcileBackfill;
the live-path reconcile (non-resume) stays.

* 🪙 fix: Guard Reconciliation Against Replays and Snapshot Races

Two Codex concurrency findings on the reconciliation:
- Client: reconcile only on a NEWLY folded primary usage. A replayed duplicate
  (folded=false on resume) can be an earlier tool-loop call sharing the run id,
  which would overwrite the latest snapshot with an earlier, smaller prompt. Moved
  the reconcile after the folded guard.
- Server: serialize the context-usage write through the same per-stream queue as
  the token-usage write. persistTokenUsage reconciles the stored snapshot
  (read-modify-write); an unserialized trackContextUsage could store a newer
  snapshot between the read and write — or a stale reconciled write could land
  after a newer snapshot — clobbering the newer run's gauge when calls interleave.
  FIFO keeps each call's snapshot ahead of its own usage and behind the next.

* chore: import order in GenerationJobManager.ts

client/src/components/Chat/Input/TokenUsage/Breakdown.tsx

@@ -46,7 +46,12 @@ export default function Breakdown({ view, showCost, currency }: BreakdownProps)
     snapshot?.effectiveInstructionTokens ?? breakdown?.instructionTokens ?? 0;
   const systemTokens =
     (breakdown?.systemMessageTokens ?? 0) + (breakdown?.dynamicInstructionTokens ?? 0);
-  const messageTokens = Math.max(0, usedTokens - instructionTokens);
+  /** Summary has its own row, so exclude it (it's part of `usedTokens`) to avoid
+   *  double-counting it inside the Messages row on a summarized turn. */
+  const messageTokens = Math.max(
+    0,
+    usedTokens - instructionTokens - (breakdown?.summaryTokens ?? 0),
+  );
   const freeTokens = maxTokens != null ? Math.max(0, maxTokens - usedTokens) : null;
 
   const groups =

client/src/hooks/SSE/__tests__/useUsageHandler.spec.tsx

@@ -0,0 +1,116 @@
+import { getDefaultStore } from 'jotai';
+import { renderHook } from '@testing-library/react';
+import type { TContextUsageEvent, TTokenUsageEvent } from 'librechat-data-provider';
+import useUsageHandler from '~/hooks/SSE/useUsageHandler';
+import { contextSnapshotFamily } from '~/store/usage';
+
+/** Mirrors a real web-search + summarization turn: calibration pinned at 5
+ *  inflated messageTokens to 187471 (used 213375), while the call's true prompt
+ *  was 53702 + 2071 cache = 55773. */
+const inflatedSnapshot = (over?: Partial<TContextUsageEvent>): TContextUsageEvent => ({
+  runId: 'run-1',
+  breakdown: {
+    maxContextTokens: 250000,
+    instructionTokens: 4205,
+    systemMessageTokens: 384,
+    dynamicInstructionTokens: 1525,
+    toolSchemaTokens: 2296,
+    summaryTokens: 1938,
+    toolCount: 1,
+    messageCount: 2,
+    messageTokens: 187471,
+    availableForMessages: 233295,
+  },
+  contextBudget: 237500,
+  remainingContextTokens: 24125,
+  calibrationRatio: 5,
+  ...over,
+});
+
+const primaryUsage = (over?: Partial<TTokenUsageEvent>): TTokenUsageEvent => ({
+  input_tokens: 53702,
+  output_tokens: 3780,
+  total_tokens: 57482,
+  input_token_details: { cache_read: 2071, cache_creation: 0 },
+  provider: 'anthropic',
+  runId: 'run-1',
+  seq: 1,
+  ...over,
+});
+
+describe('useUsageHandler — live snapshot reconciliation', () => {
+  it('reconciles the live snapshot to the primary call’s actual prompt tokens', () => {
+    const convo = 'convo-recon-1';
+    const submission = {
+      userMessage: { messageId: 'u1', conversationId: convo },
+      conversation: { conversationId: convo },
+    };
+    const { result } = renderHook(() => useUsageHandler());
+    const store = getDefaultStore();
+
+    result.current.contextHandler(inflatedSnapshot(), submission);
+    expect(store.get(contextSnapshotFamily(convo))?.breakdown.messageTokens).toBe(187471);
+
+    result.current.usageHandler(primaryUsage(), submission);
+
+    const snap = store.get(contextSnapshotFamily(convo));
+    /** used = budget − remaining = real prompt (55773), down from 213375 */
+    expect(237500 - (snap?.remainingContextTokens ?? 0)).toBe(55773);
+    expect(snap?.breakdown.messageTokens).toBe(55773 - 4205 - 1938);
+    /** instructions/summary stay raw; the anchor is preserved */
+    expect(snap?.breakdown.instructionTokens).toBe(4205);
+    expect(snap?.anchorMessageId).toBe('u1');
+  });
+
+  it('does not reconcile a replayed (already-folded) primary usage', () => {
+    /** On resume, backfill marks the run's collected usages folded; a replayed
+     *  `on_token_usage` then arrives folded=false. Since tool-loop calls share the
+     *  run id, reconciling with such a duplicate could overwrite the latest
+     *  snapshot with an earlier call's prompt — so it must be skipped. */
+    const convo = 'convo-recon-replay';
+    const submission = {
+      userMessage: { messageId: 'ur', conversationId: convo },
+      conversation: { conversationId: convo },
+    };
+    const { result } = renderHook(() => useUsageHandler());
+    const store = getDefaultStore();
+
+    result.current.contextHandler(inflatedSnapshot(), submission);
+    /** Mark the usage folded (as resume backfill would), then replay it live. */
+    result.current.backfillUsage([primaryUsage()], submission);
+    result.current.usageHandler(primaryUsage(), submission);
+
+    /** Duplicate (folded=false) → no reconcile → snapshot stays the raw estimate. */
+    expect(store.get(contextSnapshotFamily(convo))?.breakdown.messageTokens).toBe(187471);
+  });
+
+  it('does not reconcile when the usage belongs to a different run', () => {
+    const convo = 'convo-recon-2';
+    const submission = {
+      userMessage: { messageId: 'u2', conversationId: convo },
+      conversation: { conversationId: convo },
+    };
+    const { result } = renderHook(() => useUsageHandler());
+    const store = getDefaultStore();
+
+    result.current.contextHandler(inflatedSnapshot({ runId: 'run-A' }), submission);
+    result.current.usageHandler(primaryUsage({ runId: 'run-Z', seq: 9 }), submission);
+
+    expect(store.get(contextSnapshotFamily(convo))?.breakdown.messageTokens).toBe(187471);
+  });
+
+  it('does not let a tagged (summarization) usage touch the gauge', () => {
+    const convo = 'convo-recon-3';
+    const submission = {
+      userMessage: { messageId: 'u3', conversationId: convo },
+      conversation: { conversationId: convo },
+    };
+    const { result } = renderHook(() => useUsageHandler());
+    const store = getDefaultStore();
+
+    result.current.contextHandler(inflatedSnapshot(), submission);
+    result.current.usageHandler(primaryUsage({ usage_type: 'summarization', seq: 2 }), submission);
+
+    expect(store.get(contextSnapshotFamily(convo))?.breakdown.messageTokens).toBe(187471);
+  });
+});

client/src/hooks/SSE/useResumableSSE.ts

@@ -660,6 +660,9 @@ export default function useResumableSSE(
              *  this conversation keep their usage and gap events still count */
             backfillUsage(data.resumeState?.collectedUsage ?? [], resumeSubmission);
             if (data.resumeState?.contextUsage) {
+              /** Already reconciled to the call's real prompt tokens server-side
+               *  (GenerationJobManager.persistTokenUsage) when the snapshot's call
+               *  completed, so install it as-is — no client backfill reconcile. */
               contextHandler(data.resumeState.contextUsage, resumeSubmission);
             }
             /** Output streamed before this resume is not re-delivered as deltas

client/src/hooks/SSE/useUsageHandler.ts

@@ -1,6 +1,6 @@
 import { useRef, useMemo } from 'react';
 import { getDefaultStore } from 'jotai';
-import { Constants } from 'librechat-data-provider';
+import { Constants, reconcileContextUsage, promptTokensFromUsage } from 'librechat-data-provider';
 import type {
   TMessage,
   TConversation,
@@ -191,15 +191,41 @@ export default function useUsageHandler(): UsageHandlers {
       return true;
     };
 
+    /** Reconcile the live snapshot's calibrated estimate to a primary call's
+     *  ACTUAL prompt tokens. The snapshot is pre-invoke for that call; this usage
+     *  is its post-invoke truth, so the gauge stops showing the SDK multiplier's
+     *  inflation (severe when a provider injects server-side content like web
+     *  search). Resume is handled server-side (the job-store snapshot is reconciled
+     *  when the call's usage is persisted), so no backfill reconcile is needed. */
+    const reconcileLiveSnapshot = (data: TTokenUsageEvent, submission: UsageSubmissionLike) => {
+      const convoKey = getConvoKey(submission);
+      const snapshotAtom = contextSnapshotFamily(convoKey);
+      const snapshot = jotai.get(snapshotAtom);
+      if (snapshot == null) {
+        return;
+      }
+      /** The snapshot precedes this call, so it must belong to the same run */
+      if (snapshot.runId != null && data.runId != null && snapshot.runId !== data.runId) {
+        return;
+      }
+      const reconciled = reconcileContextUsage(snapshot, promptTokensFromUsage(data));
+      jotai.set(snapshotAtom, { ...reconciled, anchorMessageId: snapshot.anchorMessageId });
+    };
+
     const usageHandler: UsageHandlers['usageHandler'] = (data, submission) => {
       const folded = foldUsage(data, submission);
 
-      /** Only primary-call usage drives the live context estimate; tagged
-       *  buckets (summarization, subagent) fold into totals/cost only. Skip
-       *  the bump for an already-counted event replayed on resume. */
+      /** Only a NEWLY folded primary call drives the live gauge. A replayed
+       *  duplicate (folded=false on resume) can be an EARLIER tool-loop call that
+       *  shares this run's id — reconciling with it would overwrite the latest
+       *  snapshot with an earlier, smaller prompt. Tagged buckets (summarization,
+       *  subagent) never touch the context gauge. */
       if (!folded || data.usage_type != null) {
         return;
       }
+      /** This primary's provider-reported prompt is the post-invoke truth for the
+       *  call the latest snapshot precedes — reconcile the gauge to it. */
+      reconcileLiveSnapshot(data, submission);
       /** Use the repaired completion count (not raw output_tokens) so the
        *  snapshot gauge keeps the full response for under-reporting providers */
       confirmedRef.current += normalizeUsageUnits(data).output;