From d18d62e7c1118b99a7a2b05b4dc986504f2810ce Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Tue, 16 Jun 2026 11:05:44 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=AA=99=20refactor:=20Reconcile=20Context?= =?UTF-8?q?=20Gauge=20to=20Actual=20Provider=20Tokens=20(#13780)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 🪙 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 --- .../__tests__/usageEvents.integration.spec.js | 6 + .../Chat/Input/TokenUsage/Breakdown.tsx | 7 +- .../SSE/__tests__/useUsageHandler.spec.tsx | 116 ++++++++++++++++++ client/src/hooks/SSE/useResumableSSE.ts | 3 + client/src/hooks/SSE/useUsageHandler.ts | 34 ++++- packages/api/src/agents/usage.spec.ts | 38 ++++++ packages/api/src/agents/usage.ts | 48 +++++--- .../api/src/stream/GenerationJobManager.ts | 60 +++++++-- packages/data-provider/src/types/runs.spec.ts | 112 +++++++++++++++++ packages/data-provider/src/types/runs.ts | 53 ++++++++ 10 files changed, 445 insertions(+), 32 deletions(-) create mode 100644 client/src/hooks/SSE/__tests__/useUsageHandler.spec.tsx create mode 100644 packages/data-provider/src/types/runs.spec.ts diff --git a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js index 046b236936..caa14f6ea7 100644 --- a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js +++ b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js @@ -375,6 +375,12 @@ describe('usage events through the real agents pipeline', () => { expect(resumeState.contextUsage.breakdown.maxContextTokens).toBe(MAX_CONTEXT_TOKENS); /** Latest-wins: the persisted snapshot is the second call's */ expect(resumeState.contextUsage.prePruneContextTokens).toBeGreaterThan(0); + /** Reconciled to the final primary call's actual prompt: openAI folds cache + * into input_tokens (150), so the resume snapshot's used = 150 — the real + * context, not the calibrated estimate. */ + const used = + resumeState.contextUsage.contextBudget - resumeState.contextUsage.remainingContextTokens; + expect(used).toBe(SECOND_CALL_USAGE.input_tokens); } }); diff --git a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx index bc1c5ae3c8..f8c479b68a 100644 --- a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx +++ b/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 = diff --git a/client/src/hooks/SSE/__tests__/useUsageHandler.spec.tsx b/client/src/hooks/SSE/__tests__/useUsageHandler.spec.tsx new file mode 100644 index 0000000000..59f3bf2b85 --- /dev/null +++ b/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 => ({ + 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 => ({ + 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); + }); +}); diff --git a/client/src/hooks/SSE/useResumableSSE.ts b/client/src/hooks/SSE/useResumableSSE.ts index 160efb81b1..1eb7267044 100644 --- a/client/src/hooks/SSE/useResumableSSE.ts +++ b/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 diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index 152890252e..4625282644 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/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; diff --git a/packages/api/src/agents/usage.spec.ts b/packages/api/src/agents/usage.spec.ts index 2a1ef75fec..6c293f7768 100644 --- a/packages/api/src/agents/usage.spec.ts +++ b/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), diff --git a/packages/api/src/agents/usage.ts b/packages/api/src/agents/usage.ts index a2fca29e4a..31e8498867 100644 --- a/packages/api/src/agents/usage.ts +++ b/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, 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, + 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, 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 = [], ): 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 = {}; @@ -292,7 +304,7 @@ export function buildPersistedContextUsage( toolTokenCounts = Object.keys(trimmed).length > 0 ? trimmed : undefined; } return { - ...snapshot, + ...reconciled, breakdown: { ...breakdown, toolTokenCounts }, ...(completedOutputTokens > 0 && { completedOutputTokens }), }; diff --git a/packages/api/src/stream/GenerationJobManager.ts b/packages/api/src/stream/GenerationJobManager.ts index 9b1980a6b8..4445417c5c 100644 --- a/packages/api/src/stream/GenerationJobManager.ts +++ b/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 = { 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 { diff --git a/packages/data-provider/src/types/runs.spec.ts b/packages/data-provider/src/types/runs.spec.ts new file mode 100644 index 0000000000..28593d95f3 --- /dev/null +++ b/packages/data-provider/src/types/runs.spec.ts @@ -0,0 +1,112 @@ +import type { TContextUsageEvent, TTokenUsageEvent } from './runs'; +import { promptTokensFromUsage, reconcileContextUsage } from './runs'; + +describe('promptTokensFromUsage', () => { + it('adds cache reads/writes for additive providers (Anthropic)', () => { + const event: TTokenUsageEvent = { + input_tokens: 53702, + input_token_details: { cache_read: 2071, cache_creation: 0 }, + provider: 'anthropic', + }; + expect(promptTokensFromUsage(event)).toBe(55773); + }); + + it('treats input_tokens as the full prompt for subset providers (OpenAI)', () => { + const event: TTokenUsageEvent = { + input_tokens: 1000, + input_token_details: { cache_read: 200, cache_creation: 100 }, + provider: 'openAI', + }; + expect(promptTokensFromUsage(event)).toBe(1000); + }); + + it('handles missing fields', () => { + expect(promptTokensFromUsage({ provider: 'anthropic' })).toBe(0); + }); + + it('uses the magnitude heuristic when the provider is absent (cache ≤ input ⇒ included)', () => { + /** OpenAI-compatible/custom payload with no provider: cache already folded + * into input_tokens, so it must NOT be re-added. */ + const event: TTokenUsageEvent = { + input_tokens: 1000, + input_token_details: { cache_read: 400, cache_creation: 0 }, + }; + expect(promptTokensFromUsage(event)).toBe(1000); + }); + + it('adds cache when provider is absent and cache exceeds input (additive shape)', () => { + const event: TTokenUsageEvent = { + input_tokens: 100, + input_token_details: { cache_read: 900, cache_creation: 0 }, + }; + expect(promptTokensFromUsage(event)).toBe(1000); + }); +}); + +describe('reconcileContextUsage', () => { + /** The exact over-reporting case from a real web-search + summarization turn: + * calibrationRatio pinned at 5 inflated messageTokens to 187471 → used 213375, + * while the provider's real prompt for that call was 55773. */ + const inflatedSnapshot: 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, + }; + + it('reconciles the inflated estimate to the real prompt tokens', () => { + const result = reconcileContextUsage(inflatedSnapshot, 55773); + /** used = budget − remaining = real prompt, down from the inflated 213375 */ + expect(237500 - (result.remainingContextTokens ?? 0)).toBe(55773); + /** only messageTokens is corrected; instructions/summary stay raw */ + expect(result.breakdown.messageTokens).toBe(55773 - 4205 - 1938); + expect(result.breakdown.instructionTokens).toBe(4205); + expect(result.breakdown.summaryTokens).toBe(1938); + /** rows still sum to the real total */ + expect( + result.breakdown.messageTokens + + result.breakdown.instructionTokens + + result.breakdown.summaryTokens, + ).toBe(55773); + }); + + it('clamps messageTokens to zero when the prompt is smaller than the overhead', () => { + const result = reconcileContextUsage(inflatedSnapshot, 3000); + expect(result.breakdown.messageTokens).toBe(0); + expect(result.remainingContextTokens).toBe(237500 - 3000); + }); + + it('is a no-op for an unusable prompt count', () => { + expect(reconcileContextUsage(inflatedSnapshot, 0)).toBe(inflatedSnapshot); + expect(reconcileContextUsage(inflatedSnapshot, -5)).toBe(inflatedSnapshot); + expect(reconcileContextUsage(inflatedSnapshot, NaN)).toBe(inflatedSnapshot); + }); + + it('end-to-end: promptTokensFromUsage feeds reconcileContextUsage (followup turn)', () => { + const followupSnapshot: TContextUsageEvent = { + ...inflatedSnapshot, + breakdown: { ...inflatedSnapshot.breakdown, messageTokens: 30480 }, + remainingContextTokens: 202815, + }; + const usage: TTokenUsageEvent = { + input_tokens: 7804, + input_token_details: { cache_read: 2071, cache_creation: 0 }, + provider: 'anthropic', + }; + const result = reconcileContextUsage(followupSnapshot, promptTokensFromUsage(usage)); + expect(237500 - (result.remainingContextTokens ?? 0)).toBe(9875); + expect(result.breakdown.messageTokens).toBe(9875 - 4205 - 1938); + }); +}); diff --git a/packages/data-provider/src/types/runs.ts b/packages/data-provider/src/types/runs.ts index 227154f4c2..280cd11cb2 100644 --- a/packages/data-provider/src/types/runs.ts +++ b/packages/data-provider/src/types/runs.ts @@ -1,3 +1,5 @@ +import { inputTokensIncludesCache } from '../schemas'; + export enum ContentTypes { TEXT = 'text', THINK = 'think', @@ -126,6 +128,57 @@ export type TTokenUsageEvent = { cost?: number; }; +/** + * Full prompt token count for one completed model call — the EXACT context the + * model saw, provider-aware: additive providers (Anthropic/Bedrock) report + * `input_tokens` excluding cache, so cache reads/writes are added back; subset + * providers (OpenAI/…) already fold cache into `input_tokens`. When the provider + * is absent (custom/OpenAI-compatible payloads), fall back to the same magnitude + * heuristic `normalizeUsageUnits` uses — cache ≤ input means it's already + * included — so cached events aren't re-inflated. The ground truth the gauge + * reconciles its calibrated estimate to. + */ +export const promptTokensFromUsage = (event: TTokenUsageEvent): number => { + const input = event.input_tokens ?? 0; + const details = event.input_token_details ?? {}; + const cacheRead = details.cache_read ?? 0; + const cacheCreation = details.cache_creation ?? 0; + const includesCache = + event.provider != null + ? inputTokensIncludesCache(event.provider) + : cacheRead + cacheCreation <= input; + return includesCache ? input : input + cacheRead + cacheCreation; +}; + +/** + * Reconciles a pre-invoke context snapshot's CALIBRATED estimate to a call's + * ACTUAL prompt tokens. The SDK's calibration multiplier scales only + * `messageTokens` (instructions/summary are raw tiktoken counts), and it can + * over-shoot badly when a provider injects server-side content the SDK never + * counted (e.g. Anthropic web search) — pinning the gauge several× too high and + * persisting it. Trust the provider's own prompt count: keep the raw + * instruction/summary rows, set `messageTokens` to the remainder, and recompute + * the free space. No-op when `promptTokens` is unusable. + */ +export const reconcileContextUsage = ( + snapshot: TContextUsageEvent, + promptTokens: number, +): TContextUsageEvent => { + if (!Number.isFinite(promptTokens) || promptTokens <= 0) { + return snapshot; + } + const { breakdown } = snapshot; + const budget = snapshot.contextBudget ?? breakdown.maxContextTokens; + const nonMessageTokens = (breakdown.instructionTokens ?? 0) + (breakdown.summaryTokens ?? 0); + const messageTokens = Math.max(0, promptTokens - nonMessageTokens); + return { + ...snapshot, + breakdown: { ...breakdown, messageTokens }, + remainingContextTokens: + budget != null ? Math.max(0, budget - promptTokens) : snapshot.remainingContextTokens, + }; +}; + /** Lifecycle phase carried on subagent-progress envelopes (mirrors SDK SubagentUpdatePhase). */ export type SubagentUpdatePhase = | 'start'