🪟 feat: Faithful Over-Window Context Estimate via Prune Mirror and Overhead Reserve (#13959)
Some checks are pending
Docker Dev Branch Images Build / build (Dockerfile, lc-dev, node) (push) Waiting to run
Docker Dev Branch Images Build / build (Dockerfile.multi, lc-dev-api, api-build) (push) Waiting to run
GitNexus Index / index (push) Waiting to run
GitNexus Index / post-index (push) Blocked by required conditions

*  feat: Mirror send-path pruning in the over-window context estimate

For a snapshot-less branch whose tokens exceed the window, the send path
prunes oldest-first (getMessagesWithinTokenLimit), so the next call can sit
well under the window. The gauge previously clamped the full sum to 100%,
hiding that headroom. Add prunedBranchTokens — a newest->oldest walk that
keeps messages until the next would overflow the message budget (max minus the
summary baseline), mirroring the pruner — and use it on the estimate path in
place of the clamp. Approximation: omits the instruction/tool overhead and
tool-call pairing the real pruner accounts for (unknowable for a snapshot-less
branch); superseded by an exact snapshot once the branch is generated.

*  feat: Reserve cached instruction/tool overhead in the snapshot-less estimate

The over-window prune mirror and the gauge couldn't account for the fixed
instruction + tool-schema overhead the next call always sends, because a
snapshot-less branch has no breakdown. The backend already emits that overhead
in the ON_CONTEXT_USAGE breakdown, so cache it per agent/model (keyed
endpoint::model::agentId, already inclusive of tool schemas) from the live
usage events, then reserve it from the prune budget and add it to used so the
estimate is consistent with snapshots. Falls back to message-only until the
agent has run once this session. Surfaced as a System row in the estimate
breakdown.

* 🩹 fix: Address Codex review on the over-window estimate

- Key the overhead cache by agentId when present. useTokenLimits resolves an
  agent to its real provider/model, so the reader keyed `provider::model::agent`
  while the writer stored `agents::::agent` — a cache miss for the main agents
  case. Both sides now resolve to `agent:<id>` (non-agent configs: endpoint:model).
- Skip the overhead reserve when a summary baseline exists: computeSummaryUsedTokens
  already folds instruction/tool overhead into that marker, so adding it again
  double-counted on summarized branches.
- Collapse the breakdown's input/output/estimated rows into one pruned Messages
  row when over-window pruning ran, so the popover matches the gauge instead of
  summing to the discarded pre-prune history.
This commit is contained in:
Danny Avila 2026-06-25 17:12:53 -04:00 committed by GitHub
parent 0d1103f62f
commit 0789a04d11
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 253 additions and 18 deletions

View file

@ -145,13 +145,28 @@ export default function Breakdown({ view, showCost, currency }: BreakdownProps)
max={maxTokens}
/>
)}
<Row label={localize('com_ui_input')} value={view.branchTotals.input} />
<Row
label={localize('com_ui_output')}
value={view.branchTotals.output + view.liveTokens}
/>
{view.estimatedTokens > 0 && (
<Row label={localize('com_ui_context_estimated')} value={view.estimatedTokens} />
{view.messagesPruned ? (
/** Over-window: the per-category split no longer describes what's
* sent, so show the pruned message total (incl. in-flight). */
<Row
label={localize('com_ui_context_messages')}
value={view.messageTokens + view.liveTokens}
max={maxTokens}
/>
) : (
<>
<Row label={localize('com_ui_input')} value={view.branchTotals.input} />
<Row
label={localize('com_ui_output')}
value={view.branchTotals.output + view.liveTokens}
/>
{view.estimatedTokens > 0 && (
<Row label={localize('com_ui_context_estimated')} value={view.estimatedTokens} />
)}
</>
)}
{view.overheadTokens > 0 && (
<Row label={localize('com_ui_context_system')} value={view.overheadTokens} />
)}
{maxTokens == null && (
<p className="text-xs text-text-secondary">{localize('com_ui_context_unknown')}</p>

View file

@ -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,
]);
}

View file

@ -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);

View file

@ -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);
});
});

View file

@ -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<string, number>();
/** 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);

View file

@ -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',

View file

@ -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 newestoldest 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