diff --git a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx index fe62305b62..6d1f4d9afb 100644 --- a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx +++ b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx @@ -145,13 +145,28 @@ export default function Breakdown({ view, showCost, currency }: BreakdownProps) max={maxTokens} /> )} - - - {view.estimatedTokens > 0 && ( - + {view.messagesPruned ? ( + /** Over-window: the per-category split no longer describes what's + * sent, so show the pruned message total (incl. in-flight). */ + + ) : ( + <> + + + {view.estimatedTokens > 0 && ( + + )} + + )} + {view.overheadTokens > 0 && ( + )} {maxTokens == null && (

{localize('com_ui_context_unknown')}

diff --git a/client/src/hooks/Chat/useTokenUsage.ts b/client/src/hooks/Chat/useTokenUsage.ts index 39b6df9220..1f4a474d3f 100644 --- a/client/src/hooks/Chat/useTokenUsage.ts +++ b/client/src/hooks/Chat/useTokenUsage.ts @@ -6,6 +6,8 @@ import type { TMessage, TConversation, TModelTokenomics } from 'librechat-data-p import type { BranchTotals, BranchUsage } from '~/utils/tokens'; import type { ContextSnapshot } from '~/store/usage'; import { + overheadKey, + getModelOverhead, liveTokensFamily, totalUsageFamily, removeUsageAtoms, @@ -21,6 +23,7 @@ import { clearIndex, mergeUsage, sumTotalUsage, + prunedBranchTokens, findBranchSnapshotAnchor, } from '~/utils'; import { useLatestMessageId } from '~/hooks/Messages/useLatestMessage'; @@ -56,6 +59,15 @@ export interface TokenUsageView { /** Estimated tokens for count-less messages (in-flight tail excluded while * streaming); 0 on snapshots. Rendered as its own breakdown row. */ estimatedTokens: number; + /** Cached instruction + tool overhead applied to a snapshot-less estimate; 0 on + * snapshots (which carry their own breakdown) and until the agent has run. */ + overheadTokens: number; + /** Final message-token portion of a snapshot-less estimate (pruned when over + * window, excludes live); 0 on snapshots. */ + messageTokens: number; + /** True when over-window pruning replaced the raw message sum, so the breakdown + * shows a single pruned Messages row instead of input/output/estimated. */ + messagesPruned: boolean; rates?: TModelTokenomics; } @@ -238,6 +250,9 @@ export default function useTokenUsage({ totalCost: totalUsage.cost, liveTokens, estimatedTokens: 0, + overheadTokens: 0, + messageTokens: 0, + messagesPruned: false, rates: limits.rates, }; } @@ -252,23 +267,56 @@ export default function useTokenUsage({ * from re-summing the discarded pre-summary history (which otherwise pins the * gauge at 100% after a compaction). */ const maxTokens = limits.maxContextTokens; + const liveOnTail = liveTokens > 0; + /** Fixed instruction + tool-schema overhead for this agent/model (the latter is + * already folded into `instructionTokens`), cached from live usage events. The + * client can't otherwise know it for a snapshot-less branch, so reserve it from + * the prune budget and add it to used — making over-window pruning faithful and + * the gauge consistent with snapshots. Skipped when a summary baseline exists: + * `computeSummaryUsedTokens` already folds the overhead into that marker, so + * adding it again would double-count. 0 until the agent has run once this + * session (then falls back to message-only, as before). */ + const overheadTokens = + branchTotals.summaryBaseline > 0 + ? 0 + : getModelOverhead( + overheadKey( + limits.endpoint ?? conversation?.endpoint, + limits.model ?? conversation?.model, + conversation?.agent_id, + ), + ); /** When a stream is live the tail is the in-flight response, already counted * by `liveTokens`; drop its static estimate so a resumed/partial response * isn't double-counted on the estimate path. */ const estimatedTokens = Math.max( 0, - branchTotals.estTokens - (liveTokens > 0 ? branchTotals.tailEstTokens : 0), + branchTotals.estTokens - (liveOnTail ? branchTotals.tailEstTokens : 0), ); - const rawUsed = - branchTotals.input + - branchTotals.output + - estimatedTokens + - branchTotals.summaryBaseline + - liveTokens; - /** The send path prunes an over-window branch before calling the model, so the - * live gauge never actually exceeds the window; clamp the display to the - * window rather than show impossible values (e.g. 50k / 8k). */ - const usedTokens = maxTokens != null && maxTokens > 0 ? Math.min(rawUsed, maxTokens) : rawUsed; + const rawMessageTokens = branchTotals.input + branchTotals.output + estimatedTokens; + let messageTokens = rawMessageTokens; + /** The send path prunes an over-window branch oldest-first before calling the + * model, so the next call can sit well under the window even when the full + * branch exceeds it. Mirror that: when the raw sum overflows the message window + * (max minus the always-sent summary baseline and instruction overhead), report + * the newest messages that actually fit instead of clamping the whole branch to + * 100%. */ + if (maxTokens != null && maxTokens > 0) { + const messageBudget = Math.max(0, maxTokens - branchTotals.summaryBaseline - overheadTokens); + if (messageTokens > messageBudget) { + messageTokens = prunedBranchTokens( + conversationKey, + branchTotals.tailId, + messageBudget, + liveOnTail, + ); + } + } + /** When pruning replaced the raw sum, the per-category input/output/estimated + * rows no longer describe what's sent, so the breakdown collapses them into a + * single pruned Messages row to stay consistent with the gauge. */ + const messagesPruned = messageTokens < rawMessageTokens; + const usedTokens = overheadTokens + branchTotals.summaryBaseline + messageTokens + liveTokens; return { usedTokens, maxTokens, @@ -285,6 +333,9 @@ export default function useTokenUsage({ totalCost: totalUsage.cost, liveTokens, estimatedTokens, + overheadTokens, + messageTokens, + messagesPruned, rates: limits.rates, }; }, [ @@ -297,5 +348,7 @@ export default function useTokenUsage({ liveTokens, limits, branchSnapshot, + conversationKey, + conversation, ]); } diff --git a/client/src/hooks/SSE/useUsageHandler.ts b/client/src/hooks/SSE/useUsageHandler.ts index 4625282644..268d0153c1 100644 --- a/client/src/hooks/SSE/useUsageHandler.ts +++ b/client/src/hooks/SSE/useUsageHandler.ts @@ -9,10 +9,12 @@ import type { } from 'librechat-data-provider'; import type { ContextSnapshot } from '~/store/usage'; import { + overheadKey, markUsageFolded, liveTokensFamily, totalUsageFamily, removeUsageAtoms, + setModelOverhead, clearUsageFolded, calibrationFamily, pendingUsageFamily, @@ -150,6 +152,18 @@ export default function useUsageHandler(): UsageHandlers { if (data.calibrationRatio != null && data.calibrationRatio > 0) { jotai.set(calibrationFamily(convoKey), data.calibrationRatio); } + /** Cache the fixed instruction+tool overhead (already inclusive of tool + * schemas) per agent/model so a snapshot-less branch of the same config can + * reserve it in its prune budget and gauge. */ + const overhead = data.effectiveInstructionTokens ?? data.breakdown?.instructionTokens ?? 0; + setModelOverhead( + overheadKey( + submission.conversation?.endpoint, + submission.conversation?.model, + data.agentId, + ), + overhead, + ); streamCharsRef.current = 0; confirmedRef.current = 0; setLive(convoKey, 0); diff --git a/client/src/store/usage.spec.ts b/client/src/store/usage.spec.ts new file mode 100644 index 0000000000..79a1d26bf5 --- /dev/null +++ b/client/src/store/usage.spec.ts @@ -0,0 +1,34 @@ +import { overheadKey, setModelOverhead, getModelOverhead } from './usage'; + +describe('model overhead cache', () => { + it('keys agents by agentId so the resolved reader and raw writer agree', () => { + /** The writer stores under the raw `agents` submission (no resolved model); + * the reader resolves to the agent's real provider/model. Both must produce + * the same key, or the cache misses for the main agents case. */ + const writerKey = overheadKey('agents', '', 'agent_123'); + const readerKey = overheadKey('openAI', 'gpt-4', 'agent_123'); + expect(writerKey).toBe('agent:agent_123'); + expect(readerKey).toBe('agent:agent_123'); + + setModelOverhead(writerKey, 1500); + expect(getModelOverhead(readerKey)).toBe(1500); + }); + + it('keys non-agent configs by endpoint:model and round-trips overhead', () => { + const key = overheadKey('openAI', 'gpt-4', null); + expect(key).toBe('openAI::gpt-4'); + + /** Unknown config defaults to 0 (estimate falls back to message-only). */ + expect(getModelOverhead(key)).toBe(0); + + setModelOverhead(key, 800); + expect(getModelOverhead(key)).toBe(800); + + /** Non-positive overhead is ignored, keeping the last good value. */ + setModelOverhead(key, 0); + expect(getModelOverhead(key)).toBe(800); + + /** A different config is isolated. */ + expect(getModelOverhead(overheadKey('google', 'gemini', null))).toBe(0); + }); +}); diff --git a/client/src/store/usage.ts b/client/src/store/usage.ts index a4bd064a58..8d5e666baf 100644 --- a/client/src/store/usage.ts +++ b/client/src/store/usage.ts @@ -127,6 +127,43 @@ export function clearUsageFolded(conversationId: string): void { foldedUsageKeys.delete(conversationId); } +/** + * Per-agent/model instruction+tool overhead — the system prompt + tool-schema + * tokens the next call always sends — cached from the breakdown the live + * `ON_CONTEXT_USAGE` event emits. Keyed by agent/model (NOT per conversation), + * so a snapshot-less branch (import / never-generated) can reuse the overhead + * once the same agent has run anywhere this session, making its prune budget and + * gauge account for the fixed overhead the client can't otherwise know. Never + * cleared on convo switch; bounded by the number of distinct configs used. + */ +const modelOverhead = new Map(); + +/** Stable cache key the writer (usage handler) and reader (estimate) build + * identically. An agent resolves to its real provider/model only after its data + * loads, so the reader (resolved) and writer (raw `agents` submission) would key + * differently — key by `agentId` when present so both agree; non-agent configs + * key by endpoint:model. */ +export function overheadKey( + endpoint?: string | null, + model?: string | null, + agentId?: string | null, +): string { + if (agentId != null && agentId !== '') { + return `agent:${agentId}`; + } + return `${endpoint ?? ''}::${model ?? ''}`; +} + +export function setModelOverhead(key: string, tokens: number): void { + if (tokens > 0) { + modelOverhead.set(key, tokens); + } +} + +export function getModelOverhead(key: string): number { + return modelOverhead.get(key) ?? 0; +} + /** Jotai atomFamily entries are never GC'd — call on conversation switch/cleanup */ export function removeUsageAtoms(conversationId: string): void { branchTotalsFamily.remove(conversationId); diff --git a/client/src/utils/tokens.spec.ts b/client/src/utils/tokens.spec.ts index eddc7a2c89..ed4c865a4a 100644 --- a/client/src/utils/tokens.spec.ts +++ b/client/src/utils/tokens.spec.ts @@ -10,6 +10,7 @@ import { mergeUsage, setEntryUsage, sumTotalUsage, + prunedBranchTokens, findBranchSnapshotAnchor, estimateTokens, normalizeUsageUnits, @@ -218,6 +219,34 @@ describe('token index', () => { expect(totals.tailEstTokens).toBe(5); }); + describe('prunedBranchTokens (over-window mirror of getMessagesWithinTokenLimit)', () => { + /** u1 ← a1(huge, old) ← u2 ← a2(tail). */ + const buildChain = () => + buildIndex(CONVO, [ + msg('u1', Constants.NO_PARENT, true, 2), + msg('a1', 'u1', false, 10), + msg('u2', 'a1', true, 2), + msg('a2', 'u2', false, 2), + ]); + + it('keeps the newest messages that fit and stops at the first overflow', () => { + buildChain(); + /** Budget 8: a2(2)+u2(2)=4 fit; a1(10) would overflow → pruned. */ + expect(prunedBranchTokens(CONVO, 'a2', 8, false)).toBe(4); + }); + + it('returns the full branch sum when it fits the budget', () => { + buildChain(); + expect(prunedBranchTokens(CONVO, 'a2', 100, false)).toBe(16); + }); + + it('skips the in-flight tail when excludeTail is set', () => { + buildChain(); + /** Skip a2; a1(10)+u2(2)+u1(2)=14 all fit under 100. */ + expect(prunedBranchTokens(CONVO, 'a2', 100, true)).toBe(14); + }); + }); + it('caps the branch at a summary marker instead of re-summing compacted history', () => { const summarized = { messageId: 'a2', diff --git a/client/src/utils/tokens.ts b/client/src/utils/tokens.ts index 1372aad111..8ee65cc3dc 100644 --- a/client/src/utils/tokens.ts +++ b/client/src/utils/tokens.ts @@ -387,6 +387,59 @@ export function sumBranch( return { ...totals, tailEstTokens, tailId, usage, summaryBaseline }; } +/** + * Message tokens that would actually be sent for an over-window branch. The send + * path prunes oldest-first to fit (`getMessagesWithinTokenLimit`), so walk the + * branch newest→oldest and stop once the next message would exceed `budget`, + * mirroring its "newest-that-fits" behavior for the gauge. Approximation: it omits + * the instruction/tool-schema overhead and tool-call pairing the real pruner also + * accounts for, which the client can't know for a snapshot-less branch — close + * enough for an estimate, and superseded by an exact snapshot once generated. + * `budget` is the message window (max minus the always-sent summary baseline); + * when `excludeTail`, the in-flight tail response is skipped (it rides on + * `liveTokens`). Per-message contribution matches `sumBranch`: stored `tokenCount` + * when counted, else the char-based `estTokens`. + */ +export function prunedBranchTokens( + conversationId: string, + tailId: string | null | undefined, + budget: number, + excludeTail: boolean, +): number { + const index = registry.get(conversationId); + if (!index || !tailId || budget <= 0) { + return 0; + } + + let total = 0; + let currentId: string | null = tailId; + let guard = index.size; + let isTail = true; + + while (currentId && currentId !== Constants.NO_PARENT && guard-- > 0) { + const entry: TokenEntry | undefined = index.get(currentId); + if (!entry) { + break; + } + const skip = isTail && excludeTail; + isTail = false; + if (!skip) { + const contribution = entry.tokenCount > 0 ? entry.tokenCount : entry.estTokens; + if (total + contribution > budget) { + break; + } + total += contribution; + } + /** Pre-summary turns are subsumed by the baseline the caller already reserved, + * so stop after counting the summarizing turn — mirrors `sumBranch`. */ + if (entry.summaryUsedTokens != null && entry.summaryUsedTokens > 0) { + break; + } + currentId = entry.parentMessageId; + } + return total; +} + /** * Sums provider usage/cost across EVERY message in the conversation (all * branches, including regenerated/abandoned responses) — the conversation