mirror of
https://github.com/danny-avila/LibreChat.git
synced 2026-08-04 14:57:42 +00:00
🪙 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
This commit is contained in:
Danny Avila
GitHub
parent
055585f9f1
commit
d18d62e7c1
10 changed files with 445 additions and 32 deletions
|
|
@ -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);
|
||||
}
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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 =
|
||||
|
|
|
|||
116
client/src/hooks/SSE/__tests__/useUsageHandler.spec.tsx
Normal file
116
client/src/hooks/SSE/__tests__/useUsageHandler.spec.tsx
Normal file
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 }),
|
||||
};
|
||||
|
|
|
|||
|
|
@ -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> {
|
||||
|
|
|
|||
112
packages/data-provider/src/types/runs.spec.ts
Normal file
112
packages/data-provider/src/types/runs.spec.ts
Normal file
|
|
@ -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);
|
||||
});
|
||||
});
|
||||
|
|
@ -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'
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue