🪙 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

packages/api/src/agents/usage.spec.ts

@@ -1755,6 +1755,44 @@ describe('buildPersistedContextUsage', () => {
     expect(buildPersistedContextUsage(baseSnapshot, []).completedOutputTokens).toBeUndefined();
   });
 
+  it('reconciles the inflated estimate to the final call’s real prompt tokens', () => {
+    /** Real web-search + summarization turn: calibration pinned at 5 inflated
+     *  messageTokens to 187471 (used 213375), but the answer call's true prompt was
+     *  53702 + 2071 cache = 55773. The persisted blob must show the real context so
+     *  a reload isn't stuck several× too high. */
+    const inflated: 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,
+    };
+    const events: TTokenUsageEvent[] = [
+      {
+        input_tokens: 53702,
+        output_tokens: 3780,
+        input_token_details: { cache_read: 2071, cache_creation: 0 },
+        provider: 'anthropic',
+        runId: 'run-1',
+      },
+    ];
+    const result = buildPersistedContextUsage(inflated, events);
+    expect(237500 - (result.remainingContextTokens ?? 0)).toBe(55773);
+    expect(result.breakdown.messageTokens).toBe(55773 - 4205 - 1938);
+    expect(result.completedOutputTokens).toBe(3780);
+  });
+
   it('attributes completedOutputTokens to the snapshot run, not a parallel run', () => {
     /** Parallel/direct runs interleave: this snapshot is run-1, but run-2 emits a
      *  later primary usage. The persisted delta must be run-1's own output (40),

packages/api/src/agents/usage.ts

@@ -1,5 +1,9 @@
 import { logger } from '@librechat/data-schemas';
-import { inputTokensIncludesCache } from 'librechat-data-provider';
+import {
+  inputTokensIncludesCache,
+  reconcileContextUsage,
+  promptTokensFromUsage,
+} from 'librechat-data-provider';
 import type {
   TCustomConfig,
   TResponseUsage,
@@ -247,14 +251,14 @@ function normalizeEventUnits(event: TTokenUsageEvent): {
   };
 }
 
-/** Output tokens of the final primary model call belonging to the snapshot's
- *  run — 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. Filtering
- *  by `runId` prevents a parallel run's later usage from being attributed to this
- *  snapshot; untagged events (older lib / resume) match any run for back-compat. */
-function finalCallOutputTokens(events: ReadonlyArray<TTokenUsageEvent>, runId?: string): number {
+/** The final primary (non-tagged) model call belonging to the snapshot's run —
+ *  the call the latest pre-invoke snapshot precedes. Filtering by `runId` prevents
+ *  a parallel run's later usage from being attributed to this snapshot; untagged
+ *  events (older lib / resume) match any run for back-compat. */
+function finalPrimaryCall(
+  events: ReadonlyArray<TTokenUsageEvent>,
+  runId?: string,
+): TTokenUsageEvent | undefined {
   for (let i = events.length - 1; i >= 0; i--) {
     const event = events[i];
     if (event.usage_type != null) {
@@ -263,24 +267,32 @@ function finalCallOutputTokens(events: ReadonlyArray<TTokenUsageEvent>, runId?:
     if (runId != null && event.runId != null && event.runId !== runId) {
       continue;
     }
-    return normalizeEventUnits(event).output;
+    return event;
   }
-  return 0;
+  return undefined;
 }
 
 /**
  * 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.
+ * `responseMessage.metadata.contextUsage`. Reconciles the calibrated estimate to
+ * the final call's ACTUAL prompt tokens (the SDK multiplier over-inflates
+ * `messageTokens`, badly so when a provider injects server-side content like web
+ * search), so a reloaded turn shows the real context — not a several×-too-high
+ * number. Trims zero-valued per-tool counts (privacy/size) and records the final
+ * call's output as `completedOutputTokens` so rehydration adds the same
+ * post-snapshot delta the live gauge did. The client re-anchors the blob to the
+ * response message id on load.
  */
 export function buildPersistedContextUsage(
   snapshot: TContextUsageEvent,
   usageEvents: ReadonlyArray<TTokenUsageEvent> = [],
 ): TContextUsageEvent {
-  const { breakdown } = snapshot;
-  const completedOutputTokens = finalCallOutputTokens(usageEvents, snapshot.runId);
+  const finalCall = finalPrimaryCall(usageEvents, snapshot.runId);
+  const completedOutputTokens = finalCall ? normalizeEventUnits(finalCall).output : 0;
+  const reconciled = finalCall
+    ? reconcileContextUsage(snapshot, promptTokensFromUsage(finalCall))
+    : snapshot;
+  const { breakdown } = reconciled;
   let toolTokenCounts = breakdown.toolTokenCounts;
   if (toolTokenCounts != null) {
     const trimmed: Record<string, number> = {};
@@ -292,7 +304,7 @@ export function buildPersistedContextUsage(
     toolTokenCounts = Object.keys(trimmed).length > 0 ? trimmed : undefined;
   }
   return {
-    ...snapshot,
+    ...reconciled,
     breakdown: { ...breakdown, toolTokenCounts },
     ...(completedOutputTokens > 0 && { completedOutputTokens }),
   };

packages/api/src/stream/GenerationJobManager.ts

@@ -1,6 +1,17 @@
 import { logger, getTenantId, SYSTEM_TENANT_ID } from '@librechat/data-schemas';
-import { Constants, UsageEvents, parseTextParts } from 'librechat-data-provider';
-import type { Agents, TMessageContentParts } from 'librechat-data-provider';
+import {
+  Constants,
+  UsageEvents,
+  parseTextParts,
+  reconcileContextUsage,
+  promptTokensFromUsage,
+} from 'librechat-data-provider';
+import type {
+  TMessageContentParts,
+  TContextUsageEvent,
+  TTokenUsageEvent,
+  Agents,
+} from 'librechat-data-provider';
 import type { StandardGraph } from '@librechat/agents';
 import type {
   SerializableJobData,
@@ -12,10 +23,10 @@ import type {
 import type { GenerationJobStore } from '~/app/metrics';
 import type * as t from '~/types';
 import {
-  recordGenerationJob,
   recordGenerationStreamResumePendingEvents,
   recordGenerationStreamSubscription,
   setGenerationJobsInFlight,
+  recordGenerationJob,
 } from '~/app/metrics';
 import { InMemoryEventTransport } from './implementations/InMemoryEventTransport';
 import { InMemoryJobStore } from './implementations/InMemoryJobStore';
@@ -1167,9 +1178,17 @@ class GenerationJobManagerClass {
       return;
     }
 
-    await this.jobStore.updateJob(streamId, {
-      contextUsage: JSON.stringify((event as { data?: unknown }).data ?? null),
-    });
+    /** Share the token-usage queue so snapshot + usage writes are serialized per
+     *  stream: `persistTokenUsage` reconciles the stored snapshot (read-modify-
+     *  write), and a snapshot landing between its read and write — or a stale
+     *  reconciled write landing after a newer snapshot — would clobber the newer
+     *  run's gauge when visible calls interleave. FIFO ordering keeps each call's
+     *  pre-invoke snapshot ahead of its own usage and behind the next snapshot. */
+    await this.queueJobWrite(this.tokenUsageWriteQueues, streamId, () =>
+      this.jobStore.updateJob(streamId, {
+        contextUsage: JSON.stringify((event as { data?: unknown }).data ?? null),
+      }),
+    );
   }
 
   /**
@@ -1244,9 +1263,32 @@ class GenerationJobManagerClass {
     }
     tokenUsage.push(event.data);
 
-    await this.jobStore.updateJob(streamId, {
-      tokenUsage: JSON.stringify(tokenUsage),
-    });
+    const update: Partial<SerializableJobData> = { tokenUsage: JSON.stringify(tokenUsage) };
+
+    /** Reconcile the resume snapshot to this call's ACTUAL prompt tokens. A primary
+     *  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 a resuming client restores the real context instead of the
+     *  calibration-inflated estimate — and a mid-call resume (no usage yet) simply
+     *  keeps the raw snapshot rather than mis-applying an earlier call's tokens. */
+    const usage = event.data as TTokenUsageEvent;
+    if (usage.usage_type == null && jobData.contextUsage) {
+      try {
+        const snapshot = JSON.parse(jobData.contextUsage) as TContextUsageEvent | null;
+        if (
+          snapshot != null &&
+          (snapshot.runId == null || usage.runId == null || snapshot.runId === usage.runId)
+        ) {
+          update.contextUsage = JSON.stringify(
+            reconcileContextUsage(snapshot, promptTokensFromUsage(usage)),
+          );
+        }
+      } catch {
+        /* leave the stored snapshot as-is on parse failure */
+      }
+    }
+
+    await this.jobStore.updateJob(streamId, update);
   }
 
   private async persistReplayEvent(streamId: string, event: t.ServerSentEvent): Promise<void> {