diff --git a/api/server/controllers/agents/__tests__/client.contextMetadata.spec.js b/api/server/controllers/agents/__tests__/client.contextMetadata.spec.js
new file mode 100644
index 0000000000..6df38efb6c
--- /dev/null
+++ b/api/server/controllers/agents/__tests__/client.contextMetadata.spec.js
@@ -0,0 +1,139 @@
+const AgentClient = require('../client');
+
+/** Minimal post-(maybe-)summary snapshot. baseUsed = maxContextTokens(1000) -
+ * remainingContextTokens(700) = 300, so the marker (summaryUsedTokens) is 300. */
+const snapshot = (summaryTokens) => ({
+ runId: 'run-1',
+ agentId: 'agent-1',
+ breakdown: {
+ maxContextTokens: 1000,
+ instructionTokens: 50,
+ systemMessageTokens: 50,
+ dynamicInstructionTokens: 0,
+ toolSchemaTokens: 0,
+ summaryTokens,
+ toolCount: 0,
+ messageCount: 1,
+ messageTokens: 20,
+ availableForMessages: 900,
+ },
+ contextBudget: 1000,
+ remainingContextTokens: 700,
+ prePruneContextTokens: 300,
+ effectiveInstructionTokens: 50,
+ calibrationRatio: 1,
+});
+
+const primary = { input_tokens: 10, output_tokens: 5, total_tokens: 15 };
+const summarizationUsage = { ...primary, usage_type: 'summarization' };
+const primaryFor = (runId, output_tokens) => ({
+ input_tokens: 10,
+ output_tokens,
+ total_tokens: 10 + output_tokens,
+ provider: 'openAI',
+ runId,
+});
+
+function buildMeta({ snap, latestUsageIndex, usageEvents }) {
+ const self = {
+ collectedThoughtSignatures: null,
+ usageEmitSink: usageEvents,
+ contextUsageSink: snap
+ ? { latest: snap, count: 1, latestUsageIndex }
+ : { latest: null, count: 0 },
+ };
+ return AgentClient.prototype.buildResponseMetadata.call(self);
+}
+
+describe('AgentClient.buildResponseMetadata — snapshot persistence + summary marker', () => {
+ it('persists the snapshot when a primary usage follows it (normal turn)', () => {
+ const meta = buildMeta({ snap: snapshot(0), latestUsageIndex: 0, usageEvents: [primary] });
+ expect(meta.contextUsage).toBeDefined();
+ expect(meta.summaryUsedTokens).toBeUndefined();
+ });
+
+ it('persists the post-summary snapshot when the only pre-primary usage is the summarization', () => {
+ /** A summarized turn: the summarization usage precedes the post-summary
+ * snapshot (index 1), then the model's primary usage follows it. The old
+ * count guard miscounted and dropped this; the new guard keeps it. The
+ * marker subtracts the summarization output (5): the generated summary is in
+ * the snapshot baseline (summaryTokens) AND the response tokenCount, so
+ * 300 − 5 = 295 keeps the client estimate from counting it twice. */
+ const meta = buildMeta({
+ snap: snapshot(80),
+ latestUsageIndex: 1,
+ usageEvents: [summarizationUsage, primary],
+ });
+ expect(meta.contextUsage).toBeDefined();
+ expect(meta.summaryUsedTokens).toBe(295);
+ });
+
+ it('still emits the summary marker when the final call emitted no usage', () => {
+ /** Interrupted summarized turn: no primary usage follows the latest snapshot,
+ * so the snapshot is (correctly) not persisted — but the coarse marker
+ * survives so the client estimate still caps the discarded history. The
+ * summarization output (5) is subtracted (300 − 5 = 295). */
+ const meta = buildMeta({
+ snap: snapshot(80),
+ latestUsageIndex: 1,
+ usageEvents: [summarizationUsage],
+ });
+ expect(meta.contextUsage).toBeUndefined();
+ expect(meta.summaryUsedTokens).toBe(295);
+ });
+
+ it('drops the snapshot and emits no marker when the final call had no usage and no summary', () => {
+ const meta = buildMeta({ snap: snapshot(0), latestUsageIndex: 1, usageEvents: [primary] });
+ expect(meta.contextUsage).toBeUndefined();
+ expect(meta.summaryUsedTokens).toBeUndefined();
+ });
+
+ it('does not persist the snapshot when only a parallel run produced post-snapshot usage', () => {
+ /** A snapshot (run-1) → B snapshot (run-1 is latest) but the only following
+ * usage belongs to a sibling run (run-2). The guard must NOT persist run-1's
+ * snapshot with run-2's output — it falls back to the per-message estimate. */
+ const meta = buildMeta({
+ snap: snapshot(0),
+ latestUsageIndex: 0,
+ usageEvents: [primaryFor('run-2', 99)],
+ });
+ expect(meta.contextUsage).toBeUndefined();
+ });
+
+ it('persists with the snapshot run output when its own primary usage follows', () => {
+ const meta = buildMeta({
+ snap: snapshot(0),
+ latestUsageIndex: 0,
+ usageEvents: [primaryFor('run-2', 99), primaryFor('run-1', 7)],
+ });
+ expect(meta.contextUsage).toBeDefined();
+ expect(meta.contextUsage.completedOutputTokens).toBe(7);
+ });
+
+ it('subtracts earlier tool-loop output from the summary marker (interrupted turn)', () => {
+ /** Multi-call summarized turn stopped before the final usage: the earlier
+ * call (output 40) is baked into baseUsed (300), so the marker is 300 − 40 =
+ * 260. No primary follows the snapshot, so the full snapshot is not persisted
+ * and the client uses this marker — which must not double-count the 40 that
+ * the response tokenCount also carries. */
+ const meta = buildMeta({
+ snap: snapshot(80),
+ latestUsageIndex: 1,
+ usageEvents: [primaryFor('run-1', 40)],
+ });
+ expect(meta.contextUsage).toBeUndefined();
+ expect(meta.summaryUsedTokens).toBe(260);
+ });
+
+ it('subtracts only this run’s earlier output, not a parallel run’s', () => {
+ const meta = buildMeta({
+ snap: snapshot(80),
+ latestUsageIndex: 2,
+ usageEvents: [primaryFor('run-2', 999), primaryFor('run-1', 40), primaryFor('run-1', 5)],
+ });
+ /** baseUsed 300 − run-1's earlier 40 = 260; run-2's 999 is ignored. */
+ expect(meta.summaryUsedTokens).toBe(260);
+ /** run-1's own primary follows the snapshot → snapshot persisted with output 5. */
+ expect(meta.contextUsage.completedOutputTokens).toBe(5);
+ });
+});
diff --git a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js
index 9023d210cd..046b236936 100644
--- a/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js
+++ b/api/server/controllers/agents/__tests__/usageEvents.integration.spec.js
@@ -1,7 +1,7 @@
const { z } = require('zod');
const { tool } = require('@langchain/core/tools');
const { ChatGenerationChunk } = require('@langchain/core/outputs');
-const { HumanMessage, AIMessageChunk } = require('@langchain/core/messages');
+const { HumanMessage, AIMessage, AIMessageChunk } = require('@langchain/core/messages');
const {
Run,
Providers,
@@ -377,4 +377,95 @@ describe('usage events through the real agents pipeline', () => {
expect(resumeState.contextUsage.prePruneContextTokens).toBeGreaterThan(0);
}
});
+
+ /** Drives a real summarization (tight context + padded history); self-summarize
+ * reuses the overridden fake model so no API key is needed. */
+ async function runSummarizationLoop({ res, collectedUsage, contextUsageSink, usageEmitSink }) {
+ const { aggregateContent } = createContentAggregator();
+ const handlers = getDefaultHandlers({
+ res,
+ aggregateContent,
+ toolEndCallback: () => {},
+ collectedUsage,
+ contextUsageSink,
+ usageEmitSink,
+ summarizationOptions: { enabled: true },
+ });
+
+ const pad = 'context detail to overflow the tiny budget. '.repeat(40);
+ const history = [
+ new HumanMessage(`Turn 1 question. ${pad}`),
+ new AIMessage(`Turn 1 answer. ${pad}`),
+ new HumanMessage(`Turn 2 question. ${pad}`),
+ new AIMessage(`Turn 2 answer. ${pad}`),
+ new HumanMessage(`Final question after a lot of prior history. ${pad}`),
+ ];
+ const indexTokenCountMap = {};
+ history.forEach((message, i) => {
+ indexTokenCountMap[i] = charCounter(message);
+ });
+
+ const run = await Run.create({
+ runId: `summ-e2e-${Date.now()}`,
+ graphConfig: {
+ type: 'standard',
+ llmConfig: {
+ provider: Providers.OPENAI,
+ model: 'gpt-4o-mini',
+ streaming: true,
+ streamUsage: false,
+ },
+ instructions: 'You are a helpful assistant.',
+ maxContextTokens: 700,
+ summarizationEnabled: true,
+ summarizationConfig: { provider: Providers.OPENAI, model: 'gpt-4o-mini' },
+ },
+ returnContent: true,
+ customHandlers: handlers,
+ tokenCounter: charCounter,
+ indexTokenCountMap,
+ });
+
+ run.Graph.overrideModel = new UsageFakeModel(
+ { responses: ['## Summary\nPrior turns compacted.', 'Here is the final answer.'] },
+ [{ input_tokens: 40, output_tokens: 8, total_tokens: 48 }],
+ );
+
+ await run.processStream(
+ { messages: history },
+ {
+ configurable: { thread_id: 'summ-e2e-thread', user_id: 'user-1' },
+ streamMode: 'values',
+ version: 'v2',
+ },
+ );
+ return run;
+ }
+
+ /** A summarized turn compacts the context (summary tokens replace the older
+ * turns) and the reduced snapshot is persisted — the latest snapshot is
+ * followed by a primary usage, so the save guard keeps it and the client
+ * uses the snapshot (not the inflated whole-history estimate). */
+ test('persists the reduced (compacted) snapshot after summarization', async () => {
+ if (!hasContextUsageEvent) {
+ return;
+ }
+ const res = createMockRes();
+ const contextUsageSink = { latest: null, count: 0 };
+ const usageEmitSink = [];
+ await runSummarizationLoop({ res, collectedUsage: [], contextUsageSink, usageEmitSink });
+
+ const snapshot = contextUsageSink.latest;
+ /** Summarization fired: a summary exists and the kept message tokens are
+ * small (the compacted context, not the full history). */
+ expect(snapshot?.breakdown?.summaryTokens).toBeGreaterThan(0);
+ expect(snapshot?.breakdown?.messageTokens).toBeLessThan(snapshot?.breakdown?.summaryTokens);
+
+ /** The save guard keeps it: a primary usage follows the latest snapshot. */
+ const afterLatest = usageEmitSink.slice(contextUsageSink.latestUsageIndex ?? 0);
+ expect(afterLatest.some((e) => e.usage_type == null)).toBe(true);
+ expect(
+ buildPersistedContextUsage(snapshot, usageEmitSink).breakdown.summaryTokens,
+ ).toBeGreaterThan(0);
+ });
});
diff --git a/api/server/controllers/agents/callbacks.js b/api/server/controllers/agents/callbacks.js
index e7cb3dfdaa..d0b2ea4aaa 100644
--- a/api/server/controllers/agents/callbacks.js
+++ b/api/server/controllers/agents/callbacks.js
@@ -546,16 +546,23 @@ function getDefaultHandlers({
checkIfLastAgent(metadata?.last_agent_id, metadata?.langgraph_node) ||
!metadata?.hide_sequential_outputs
) {
- await emitEvent(res, streamId, { event, data });
- /** Capture the latest visible snapshot (last-wins) + count visible
- * snapshots (one per model call). The count lets the save path persist
- * the breakdown only when the FINAL call emitted usage (primary usage
- * events === snapshots), so completedOutputTokens is a real
- * post-snapshot delta and reload doesn't over-report. */
+ /** Capture the latest visible snapshot (last-wins) and how many usage
+ * events preceded it BEFORE awaiting the emit. `emitEvent` can yield
+ * (resumable SSE / Redis publish); with parallel runs active this
+ * call's own primary usage could land in `usageEmitSink` during that
+ * yield, pushing `latestUsageIndex` past the very event that proves the
+ * snapshot completed — the save path would then slice it away and drop
+ * a valid breakdown. The recorded index lets the save path persist only
+ * when a PRIMARY usage follows this snapshot (the snapshot's call
+ * actually invoked the model); a summarization detour emits a snapshot
+ * whose only following usage is tagged `summarization`, which a plain
+ * snapshot-count would over-count and wrongly drop. */
if (contextUsageSink) {
contextUsageSink.latest = data;
contextUsageSink.count = (contextUsageSink.count ?? 0) + 1;
+ contextUsageSink.latestUsageIndex = usageEmitSink?.length ?? 0;
}
+ await emitEvent(res, streamId, { event, data });
}
},
};
diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js
index 0b18529e38..94d17fa3f5 100644
--- a/api/server/controllers/agents/client.js
+++ b/api/server/controllers/agents/client.js
@@ -26,6 +26,8 @@ const {
aggregateEmittedUsage,
resolveAgentTokenConfig,
buildPersistedContextUsage,
+ computeSummaryUsedTokens,
+ priorRunOutputTokens,
createSubagentUsageSink,
isDeepSeekReasoningProvider,
GenerationJobManager,
@@ -869,16 +871,49 @@ class AgentClient extends BaseClient {
metadata.thoughtSignatures = signatures;
}
const usageEvents = this.usageEmitSink ?? [];
- /** Persist the breakdown only when the FINAL visible call (the one the latest
- * snapshot precedes) emitted usage — i.e. as many primary usage events as
- * visible snapshots. If the final call emitted no usage_metadata (provider
- * gap, or interrupted after an earlier call did emit), `completedOutputTokens`
- * would be an earlier call's output the latest snapshot already counts, so
- * reload would over-report; fall back to the coarse per-message estimate. */
- const primaryUsageCount = usageEvents.filter((event) => event.usage_type == null).length;
- const snapshotCount = this.contextUsageSink?.count ?? 0;
- if (this.contextUsageSink?.latest && snapshotCount > 0 && primaryUsageCount >= snapshotCount) {
- metadata.contextUsage = buildPersistedContextUsage(this.contextUsageSink.latest, usageEvents);
+ /** Persist the breakdown only when the latest snapshot's OWN run completed —
+ * i.e. a PRIMARY usage event (usage_type == null) from that run's id arrived
+ * AFTER the snapshot. Matching by run id keeps `completedOutputTokens` a real
+ * post-snapshot delta even when parallel/direct runs interleave (A snapshot →
+ * B snapshot → A usage must NOT persist B's snapshot with A's output); an
+ * interrupted final call that emits no usage falls back to the per-message
+ * estimate. It still keeps the post-summary snapshot: the summarization detour
+ * emits an extra snapshot whose following primary usage shares that run's id,
+ * which the old snapshot-count guard miscounted and wrongly dropped. Events
+ * without a run id (older lib / resume) match any snapshot for back-compat. */
+ const latestSnapshot = this.contextUsageSink?.latest;
+ const latestSnapshotUsageIndex = this.contextUsageSink?.latestUsageIndex ?? 0;
+ const latestSnapshotRunId = latestSnapshot?.runId;
+ const hasPrimaryAfterSnapshot = usageEvents
+ .slice(latestSnapshotUsageIndex)
+ .some(
+ (event) =>
+ event.usage_type == null &&
+ (latestSnapshotRunId == null ||
+ event.runId == null ||
+ event.runId === latestSnapshotRunId),
+ );
+ if (latestSnapshot && hasPrimaryAfterSnapshot) {
+ metadata.contextUsage = buildPersistedContextUsage(latestSnapshot, usageEvents);
+ }
+ /** Lightweight summarization marker — persisted whenever this turn compacted
+ * the context, INDEPENDENT of the snapshot guard above. When the client has
+ * no usable snapshot on the branch and falls back to the per-message
+ * estimate, it caps the discarded pre-summary history at this baseline
+ * instead of re-summing it (the gauge otherwise reads 100% forever). Shared
+ * with the abort save path via `computeSummaryUsedTokens`. Subtract the
+ * response's earlier tool-loop outputs (the primaries that preceded the
+ * latest snapshot, same run): those tokens are inside the snapshot baseline
+ * AND in the response `tokenCount` the client estimate adds on top, so
+ * leaving them in the marker double-counts them on a multi-call turn. */
+ const priorOutputTokens = priorRunOutputTokens(
+ usageEvents,
+ latestSnapshotUsageIndex,
+ latestSnapshotRunId,
+ );
+ const summaryUsedTokens = computeSummaryUsedTokens(latestSnapshot, priorOutputTokens);
+ if (summaryUsedTokens != null) {
+ metadata.summaryUsedTokens = summaryUsedTokens;
}
const usage = aggregateEmittedUsage(usageEvents);
if (usage) {
diff --git a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx
index cdc5a138ae..bc1c5ae3c8 100644
--- a/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx
+++ b/client/src/components/Chat/Input/TokenUsage/Breakdown.tsx
@@ -133,6 +133,13 @@ export default function Breakdown({ view, showCost, currency }: BreakdownProps)
) : (
<>
+ {view.branchTotals.summaryBaseline > 0 && (
+ |
+ )}
|
{
expect(altTotals.output).toBe(1019);
});
+ it('caps the branch at a summary marker instead of re-summing compacted history', () => {
+ const summarized = {
+ messageId: 'a2',
+ parentMessageId: 'u2',
+ isCreatedByUser: false,
+ tokenCount: 40,
+ conversationId: CONVO,
+ text: '',
+ /** a2's turn compacted the history; pre-invoke context was 500 tokens. */
+ metadata: { summaryUsedTokens: 500 },
+ } as TMessage;
+ buildIndex(CONVO, [
+ msg('u1', Constants.NO_PARENT, true, 100),
+ msg('a1', 'u1', false, 9000) /** huge pre-summary history, now discarded */,
+ msg('u2', 'a1', true, 200),
+ summarized,
+ msg('u3', 'a2', true, 15),
+ msg('a3', 'u3', false, 25),
+ ]);
+
+ const totals = sumBranch(CONVO, 'a3');
+ /** Walk stops at a2: only its output + the post-summary turn are summed. */
+ expect(totals.summaryBaseline).toBe(500);
+ expect(totals.input).toBe(15);
+ expect(totals.output).toBe(65); // a3 (25) + a2 (40)
+ /** Estimate used = post-summary + compacted baseline = 580, not the 9380
+ * raw history sum that pinned the gauge at 100%. */
+ expect(totals.input + totals.output + totals.summaryBaseline).toBe(580);
+ });
+
+ it('keeps provider usage/cost across the full branch even past a summary marker', () => {
+ const summarized = {
+ messageId: 'a2',
+ parentMessageId: 'u2',
+ isCreatedByUser: false,
+ tokenCount: 40,
+ conversationId: CONVO,
+ text: '',
+ metadata: { usage: USAGE_B, summaryUsedTokens: 500 },
+ } as TMessage;
+ buildIndex(CONVO, [
+ msg('u1', Constants.NO_PARENT, true, 10),
+ responseMsg('a1', 'u1', 9000, USAGE_A) /** pre-summary spend */,
+ msg('u2', 'a1', true, 200),
+ summarized,
+ msg('u3', 'a2', true, 15),
+ responseMsg('a3', 'u3', 25, USAGE_A) /** post-summary spend */,
+ ]);
+
+ const totals = sumBranch(CONVO, 'a3');
+ /** Context is capped at the marker... */
+ expect(totals.summaryBaseline).toBe(500);
+ /** ...but cost/usage is cumulative spend and spans the WHOLE branch
+ * (a1 + a2 + a3 = 0.01 + 0.02 + 0.01), not truncated at the summary boundary. */
+ expect(totals.usage.cost).toBeCloseTo(0.04);
+ expect(totals.usage.input).toBe(400); // 100 + 200 + 100
+ });
+
it('flags whether the anchor message is on the branch', () => {
buildIndex(CONVO, [
msg('u1', Constants.NO_PARENT, true, 10),
@@ -99,6 +157,34 @@ describe('token index', () => {
expect(sumBranch(CONVO, 'a1-alt', 'a1').containsAnchor).toBe(false);
});
+ it('stops matching the anchor once the context is capped at a summary marker', () => {
+ const summarized = {
+ messageId: 'a2',
+ parentMessageId: 'u2',
+ isCreatedByUser: false,
+ tokenCount: 40,
+ conversationId: CONVO,
+ text: '',
+ metadata: { summaryUsedTokens: 500 },
+ } as TMessage;
+ buildIndex(CONVO, [
+ msg('u1', Constants.NO_PARENT, true, 100),
+ msg('a1', 'u1', false, 20),
+ msg('u2', 'a1', true, 200),
+ summarized,
+ msg('u3', 'a2', true, 15),
+ msg('a3', 'u3', false, 25),
+ ]);
+ /** a1 is older than the summary marker (a2): a snapshot anchored there is
+ * pre-summary, so it must NOT count as on-branch — else useTokenUsage revives
+ * that stale breakdown over the summary-baseline estimate. */
+ expect(sumBranch(CONVO, 'a3', 'a1').containsAnchor).toBe(false);
+ /** The summarized response's own (post-summary) snapshot still matches... */
+ expect(sumBranch(CONVO, 'a3', 'a2').containsAnchor).toBe(true);
+ /** ...as does a snapshot from a newer turn. */
+ expect(sumBranch(CONVO, 'a3', 'a3').containsAnchor).toBe(true);
+ });
+
it('tracks uncounted messages and tolerates missing parents', () => {
buildIndex(CONVO, [msg('u2', 'missing-parent', true, undefined), msg('a2', 'u2', false, 15)]);
@@ -159,6 +245,31 @@ describe('findBranchSnapshotAnchor', () => {
expect(findBranchSnapshotAnchor(CONVO, 'a1', new Map())).toBeNull();
expect(findBranchSnapshotAnchor('unknown', 'a1', new Map([['a1', 1]]))).toBeNull();
});
+
+ it('does not cross a summary marker to recover a stale pre-summary snapshot', () => {
+ const summarized = {
+ messageId: 'a2',
+ parentMessageId: 'u2',
+ isCreatedByUser: false,
+ tokenCount: 40,
+ conversationId: CONVO,
+ text: '',
+ metadata: { summaryUsedTokens: 500 },
+ } as TMessage;
+ buildIndex(CONVO, [
+ msg('u1', Constants.NO_PARENT, true, 10),
+ msg('a1', 'u1', false, 20) /** has a snapshot, but pre-summary */,
+ msg('u2', 'a1', true, 30),
+ summarized /** compacted here, no snapshot of its own */,
+ msg('u3', 'a2', true, 15),
+ msg('a3', 'u3', false, 25),
+ ]);
+ /** a1 is the only stored anchor, but it sits before the summary — the walk
+ * must stop at a2 and return null so the summary-baseline estimate is used. */
+ expect(findBranchSnapshotAnchor(CONVO, 'a3', new Map([['a1', 1]]))).toBeNull();
+ /** When the summarized response itself has a snapshot, return it. */
+ expect(findBranchSnapshotAnchor(CONVO, 'a3', new Map([['a2', 1]]))).toBe('a2');
+ });
});
describe('estimateTokens', () => {
diff --git a/client/src/utils/tokens.ts b/client/src/utils/tokens.ts
index 9bb129c6c9..adbf65a097 100644
--- a/client/src/utils/tokens.ts
+++ b/client/src/utils/tokens.ts
@@ -31,6 +31,10 @@ export interface TokenEntry {
parentMessageId: string | null;
/** Per-response provider usage from `metadata.usage` (response messages only) */
usage?: BranchUsage;
+ /** Pre-invoke compacted context size (`metadata.summaryUsedTokens`) for a
+ * response whose turn summarized. Caps the estimate so it stops re-summing
+ * the now-discarded pre-summary history. */
+ summaryUsedTokens?: number;
}
export interface BranchTotals {
@@ -47,6 +51,11 @@ export interface BranchTotals {
containsAnchor: boolean;
/** Provider usage/cost summed along the active branch */
usage: BranchUsage;
+ /** Compacted-context baseline from the deepest summarized response on the
+ * branch (0 if none). The branch walk stops there, so `input`/`output` cover
+ * only the post-summary messages; the estimate adds this to avoid counting
+ * the discarded pre-summary history. */
+ summaryBaseline: number;
}
export const EMPTY_BRANCH: BranchTotals = {
@@ -57,6 +66,7 @@ export const EMPTY_BRANCH: BranchTotals = {
tailId: null,
containsAnchor: false,
usage: EMPTY_USAGE,
+ summaryBaseline: 0,
};
/** Module-level token index: conversationId → messageId → entry. Not render state. */
@@ -128,11 +138,16 @@ function addUsage(target: BranchUsage, usage?: BranchUsage): void {
}
function toEntry(message: Partial): TokenEntry {
+ const summaryUsedTokens = message.metadata?.summaryUsedTokens;
return {
tokenCount: typeof message.tokenCount === 'number' ? message.tokenCount : 0,
isCreatedByUser: message.isCreatedByUser === true,
parentMessageId: message.parentMessageId ?? null,
usage: readPersistedUsage(message),
+ summaryUsedTokens:
+ typeof summaryUsedTokens === 'number' && summaryUsedTokens > 0
+ ? summaryUsedTokens
+ : undefined,
};
}
@@ -218,6 +233,11 @@ export function sumBranch(
const totals = { input: 0, output: 0, counted: 0, total: 0, containsAnchor: false };
const usage: BranchUsage = { ...EMPTY_USAGE };
+ let summaryBaseline = 0;
+ /** Once a summary marker is crossed, older turns are out of the CONTEXT WINDOW
+ * (subsumed by the baseline) — but their provider spend still happened, so the
+ * usage/cost walk continues to the root while context counting stops. */
+ let contextCapped = false;
let currentId: string | null = tailId;
let guard = index.size;
@@ -227,10 +247,15 @@ export function sumBranch(
break;
}
totals.total += 1;
- if (anchorId != null && currentId === anchorId) {
+ /** Only match the anchor while still inside the active context window. An
+ * anchor OLDER than the deepest summary marker belongs to a pre-summary
+ * snapshot; treating it as on-branch would let `useTokenUsage` revive that
+ * stale breakdown (discarded history) over the summary-baseline estimate
+ * that `findBranchSnapshotAnchor` correctly refuses to recover. */
+ if (!contextCapped && anchorId != null && currentId === anchorId) {
totals.containsAnchor = true;
}
- if (entry.tokenCount > 0) {
+ if (!contextCapped && entry.tokenCount > 0) {
totals.counted += 1;
if (entry.isCreatedByUser) {
totals.input += entry.tokenCount;
@@ -238,11 +263,19 @@ export function sumBranch(
totals.output += entry.tokenCount;
}
}
+ /** Cost/usage is cumulative spend — never truncated at the summary boundary. */
addUsage(usage, entry.usage);
+ /** This response's turn compacted the history: its own output is counted
+ * above; record the pre-invoke compacted baseline and stop counting context
+ * tokens for older (summarized-away) turns, but keep walking for cost. */
+ if (!contextCapped && entry.summaryUsedTokens != null) {
+ summaryBaseline = entry.summaryUsedTokens;
+ contextCapped = true;
+ }
currentId = entry.parentMessageId;
}
- return { ...totals, tailId, usage };
+ return { ...totals, tailId, usage, summaryBaseline };
}
/**
@@ -311,6 +344,12 @@ export function findBranchSnapshotAnchor(
if (!entry) {
break;
}
+ /** Stop at a summarized response that has no snapshot of its own: crossing it
+ * would recover an older PRE-summary snapshot (discarded history), which the
+ * summary-baseline estimate is meant to replace. */
+ if (entry.summaryUsedTokens != null) {
+ return null;
+ }
currentId = entry.parentMessageId;
}
diff --git a/packages/api/src/agents/usage.spec.ts b/packages/api/src/agents/usage.spec.ts
index 732db1da19..2a1ef75fec 100644
--- a/packages/api/src/agents/usage.spec.ts
+++ b/packages/api/src/agents/usage.spec.ts
@@ -10,6 +10,8 @@ import {
resolveAgentTokenConfig,
buildPersistedContextUsage,
buildAbortedResponseMetadata,
+ computeSummaryUsedTokens,
+ priorRunOutputTokens,
} from './usage';
describe('recordCollectedUsage', () => {
@@ -1752,6 +1754,157 @@ describe('buildPersistedContextUsage', () => {
it('omits completedOutputTokens when there are no primary calls', () => {
expect(buildPersistedContextUsage(baseSnapshot, []).completedOutputTokens).toBeUndefined();
});
+
+ 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),
+ * never run-2's trailing output (99). */
+ const events: TTokenUsageEvent[] = [
+ {
+ input_tokens: 100,
+ output_tokens: 40,
+ total_tokens: 140,
+ provider: 'openAI',
+ runId: 'run-1',
+ },
+ {
+ input_tokens: 200,
+ output_tokens: 99,
+ total_tokens: 299,
+ provider: 'openAI',
+ runId: 'run-2',
+ },
+ ];
+ const result = buildPersistedContextUsage(baseSnapshot, events);
+ expect(result.completedOutputTokens).toBe(40);
+ });
+
+ it('omits completedOutputTokens when only other runs emitted usage', () => {
+ /** The snapshot run never completed (no matching primary); a sibling run's
+ * output must not be borrowed — fall back to the per-message estimate. */
+ const events: TTokenUsageEvent[] = [
+ {
+ input_tokens: 200,
+ output_tokens: 99,
+ total_tokens: 299,
+ provider: 'openAI',
+ runId: 'run-2',
+ },
+ ];
+ expect(buildPersistedContextUsage(baseSnapshot, events).completedOutputTokens).toBeUndefined();
+ });
+
+ it('matches untagged usage events for back-compat (older lib / resume)', () => {
+ /** Events without a runId predate run tagging; they match any snapshot so the
+ * last primary (25) is still recorded. */
+ const events: TTokenUsageEvent[] = [
+ { input_tokens: 100, output_tokens: 40, total_tokens: 140, provider: 'openAI' },
+ { input_tokens: 200, output_tokens: 25, total_tokens: 225, provider: 'openAI' },
+ ];
+ expect(buildPersistedContextUsage(baseSnapshot, events).completedOutputTokens).toBe(25);
+ });
+});
+
+describe('computeSummaryUsedTokens', () => {
+ const summarized = (over?: Partial): TContextUsageEvent => ({
+ runId: 'run-1',
+ breakdown: {
+ maxContextTokens: 1000,
+ instructionTokens: 50,
+ systemMessageTokens: 50,
+ dynamicInstructionTokens: 0,
+ toolSchemaTokens: 0,
+ summaryTokens: 80,
+ toolCount: 0,
+ messageCount: 1,
+ messageTokens: 20,
+ availableForMessages: 900,
+ },
+ contextBudget: 1000,
+ remainingContextTokens: 700,
+ effectiveInstructionTokens: 50,
+ ...over,
+ });
+
+ it('uses contextBudget − remainingContextTokens when available', () => {
+ expect(computeSummaryUsedTokens(summarized())).toBe(300);
+ });
+
+ it('falls back to instructions + summary + messages, including summaryTokens', () => {
+ /** summaryTokens is a separate breakdown field, so the no-remaining fallback
+ * must add it: 50 + 80 + 20 = 150, not 70. */
+ expect(computeSummaryUsedTokens(summarized({ remainingContextTokens: undefined }))).toBe(150);
+ });
+
+ it('returns undefined when the turn did not summarize', () => {
+ expect(
+ computeSummaryUsedTokens(
+ summarized({ breakdown: { ...summarized().breakdown, summaryTokens: 0 } }),
+ ),
+ ).toBeUndefined();
+ expect(computeSummaryUsedTokens(null)).toBeUndefined();
+ expect(computeSummaryUsedTokens(undefined)).toBeUndefined();
+ });
+
+ it('subtracts the response’s earlier tool-loop outputs from the marker', () => {
+ /** 300 baseUsed − 90 earlier outputs = 210, so the client estimate
+ * (summaryBaseline + full response tokenCount) doesn’t double-count them. */
+ expect(computeSummaryUsedTokens(summarized(), 90)).toBe(210);
+ });
+
+ it('clamps to undefined when the prior outputs exceed the baseline', () => {
+ expect(computeSummaryUsedTokens(summarized(), 5000)).toBeUndefined();
+ });
+});
+
+describe('priorRunOutputTokens', () => {
+ const ev = (over: Partial): TTokenUsageEvent => ({
+ input_tokens: 10,
+ output_tokens: 0,
+ total_tokens: 10,
+ provider: 'openAI',
+ ...over,
+ });
+
+ it('sums primary outputs before the index for the matching run', () => {
+ const events = [
+ ev({ output_tokens: 20, runId: 'run-1' }),
+ ev({ output_tokens: 30, runId: 'run-1' }),
+ ev({ output_tokens: 99, runId: 'run-1' }), // at/after the index — excluded
+ ];
+ expect(priorRunOutputTokens(events, 2, 'run-1')).toBe(50);
+ });
+
+ it('counts run-matched primary + summarization, skips subagent/sequential and other runs', () => {
+ /** Both the primary and the summarization output are in this run's tokenCount
+ * AND baseline, so both are subtracted; subagent/sequential are excluded from
+ * the reported output total; a parallel run's primary is not this snapshot's. */
+ const events = [
+ ev({ output_tokens: 20, runId: 'run-1' }), // primary, matches
+ ev({ output_tokens: 8, runId: 'run-1', usage_type: 'summarization' }), // counted
+ ev({ output_tokens: 5, runId: 'run-1', usage_type: 'subagent' }), // skipped
+ ev({ output_tokens: 7, runId: 'run-1', usage_type: 'sequential' }), // skipped
+ ev({ output_tokens: 40, runId: 'run-2' }), // other-run primary — skipped
+ ];
+ expect(priorRunOutputTokens(events, 5, 'run-1')).toBe(28);
+ });
+
+ it('does not subtract a parallel sibling run’s summarization output', () => {
+ /** The summarize detour inherits the graph run id (traceConfig), so a sibling
+ * run's summary carries a DIFFERENT runId; its summary is in the sibling's
+ * baseline, not this snapshot's, so subtracting it would under-report. */
+ const events = [
+ ev({ output_tokens: 20, runId: 'run-1' }),
+ ev({ output_tokens: 8, runId: 'run-2', usage_type: 'summarization' }), // sibling — skipped
+ ];
+ expect(priorRunOutputTokens(events, 2, 'run-1')).toBe(20);
+ });
+
+ it('matches untagged events for back-compat and returns 0 with no prior calls', () => {
+ const events = [ev({ output_tokens: 20 }), ev({ output_tokens: 30 })];
+ expect(priorRunOutputTokens(events, 2, 'run-1')).toBe(50);
+ expect(priorRunOutputTokens(events, 0, 'run-1')).toBe(0);
+ });
});
describe('buildAbortedResponseMetadata', () => {
@@ -1783,6 +1936,65 @@ describe('buildAbortedResponseMetadata', () => {
expect(result).toEqual({ usage: { input: 100, output: 20, cacheWrite: 0, cacheRead: 0 } });
expect((result as { contextUsage?: unknown }).contextUsage).toBeUndefined();
});
+
+ const abortSnapshot: TContextUsageEvent = {
+ runId: 'run-1',
+ breakdown: {
+ maxContextTokens: 1000,
+ instructionTokens: 50,
+ systemMessageTokens: 50,
+ dynamicInstructionTokens: 0,
+ toolSchemaTokens: 0,
+ summaryTokens: 80,
+ toolCount: 0,
+ messageCount: 1,
+ messageTokens: 20,
+ availableForMessages: 900,
+ },
+ contextBudget: 1000,
+ remainingContextTokens: 700,
+ };
+
+ it('persists the summary marker (but not the full snapshot) for a stopped summarized turn', () => {
+ /** A single-call stopped turn: the interrupted call emitted no usage, so the
+ * marker is the full baseUsed (300) with nothing to subtract. */
+ const result = buildAbortedResponseMetadata({
+ tokenUsage: JSON.stringify([]),
+ contextUsage: JSON.stringify(abortSnapshot),
+ });
+ expect(result?.summaryUsedTokens).toBe(300);
+ /** The full snapshot stays off the abort path (completedOutputTokens ambiguity). */
+ expect((result as { contextUsage?: unknown }).contextUsage).toBeUndefined();
+ });
+
+ it('does not subtract any output from the marker on a stopped turn', () => {
+ /** The abort tokenCount comes from countTokens(text) or is absent — it does
+ * NOT fold in summarization/earlier-call output the way recordCollectedUsage
+ * does. So the marker is the full baseUsed (300); subtracting the summarization
+ * (8) or the primary (20) here would under-report after reload. */
+ const events: TTokenUsageEvent[] = [
+ {
+ input_tokens: 100,
+ output_tokens: 20,
+ total_tokens: 120,
+ provider: 'openAI',
+ runId: 'run-1',
+ },
+ {
+ input_tokens: 60,
+ output_tokens: 8,
+ total_tokens: 68,
+ provider: 'openAI',
+ runId: 'run-1',
+ usage_type: 'summarization',
+ },
+ ];
+ const result = buildAbortedResponseMetadata({
+ tokenUsage: JSON.stringify(events),
+ contextUsage: JSON.stringify(abortSnapshot),
+ });
+ expect(result?.summaryUsedTokens).toBe(300);
+ });
});
describe('resolveAgentTokenConfig', () => {
diff --git a/packages/api/src/agents/usage.ts b/packages/api/src/agents/usage.ts
index e8b4583cff..a2fca29e4a 100644
--- a/packages/api/src/agents/usage.ts
+++ b/packages/api/src/agents/usage.ts
@@ -247,16 +247,23 @@ function normalizeEventUnits(event: TTokenUsageEvent): {
};
}
-/** Output tokens of the response's final primary model call — 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. */
-function finalCallOutputTokens(events: ReadonlyArray): number {
+/** 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 {
for (let i = events.length - 1; i >= 0; i--) {
- if (events[i].usage_type == null) {
- return normalizeEventUnits(events[i]).output;
+ const event = events[i];
+ if (event.usage_type != null) {
+ continue;
}
+ if (runId != null && event.runId != null && event.runId !== runId) {
+ continue;
+ }
+ return normalizeEventUnits(event).output;
}
return 0;
}
@@ -273,7 +280,7 @@ export function buildPersistedContextUsage(
usageEvents: ReadonlyArray = [],
): TContextUsageEvent {
const { breakdown } = snapshot;
- const completedOutputTokens = finalCallOutputTokens(usageEvents);
+ const completedOutputTokens = finalCallOutputTokens(usageEvents, snapshot.runId);
let toolTokenCounts = breakdown.toolTokenCounts;
if (toolTokenCounts != null) {
const trimmed: Record = {};
@@ -291,6 +298,83 @@ export function buildPersistedContextUsage(
};
}
+/**
+ * Sum of this response's output tokens already folded into a later snapshot's
+ * pre-invoke baseline that the response message's `tokenCount` ALSO carries — the
+ * overlap `computeSummaryUsedTokens` subtracts from the marker so the live-path
+ * client estimate (`summaryBaseline + responseTokenCount`) doesn't double-count:
+ * - earlier tool-loop PRIMARY calls: a multi-call turn's first output sits in the
+ * kept-message context of the next call's snapshot AND in `tokenCount`.
+ * - the SUMMARIZATION call's generated summary: it sits in the snapshot baseline
+ * as `summaryTokens` AND in `tokenCount` (`recordCollectedUsage` folds
+ * summarization completion into the reported output total; subagent/sequential
+ * are kept out of that total, so they are excluded here too).
+ *
+ * Both are matched by `runId` and bounded by `beforeIndex` to the calls that
+ * preceded the snapshot. The summarize detour inherits the graph run id
+ * (`traceConfig` spreads `config.metadata.run_id`), so it shares the snapshot's
+ * `runId`; a parallel sibling run's summary carries a DIFFERENT `runId` and must
+ * NOT be subtracted (its summary lives in the sibling's baseline, not this one).
+ * Untagged events (older lib / resume) match any run for back-compat.
+ *
+ * Only the live path (which builds `tokenCount` via `recordCollectedUsage`) calls
+ * this; the abort path subtracts nothing — see {@link buildAbortedResponseMetadata}.
+ */
+export function priorRunOutputTokens(
+ events: ReadonlyArray,
+ beforeIndex: number,
+ runId?: string,
+): number {
+ let total = 0;
+ const end = Math.min(beforeIndex, events.length);
+ for (let i = 0; i < end; i++) {
+ const event = events[i];
+ if (event.usage_type != null && event.usage_type !== 'summarization') {
+ continue;
+ }
+ if (runId != null && event.runId != null && event.runId !== runId) {
+ continue;
+ }
+ total += normalizeEventUnits(event).output;
+ }
+ return total;
+}
+
+/**
+ * Pre-invoke compacted context size for a summarized turn (instructions +
+ * summary + kept messages), or `undefined` when the turn did not summarize.
+ * Persisted as the lightweight `summaryUsedTokens` marker so the client estimate
+ * fallback caps the discarded pre-summary history instead of re-summing it (the
+ * gauge otherwise reads 100% in perpetuity after a compaction). Pre-invoke, so
+ * it carries none of the `completedOutputTokens` ambiguity that keeps the full
+ * snapshot off some save paths. `summaryTokens` is a SEPARATE breakdown field, so
+ * the non-`remainingContextTokens` fallback adds it explicitly.
+ *
+ * `priorOutputTokens` (this response's earlier tool-loop outputs, see
+ * {@link priorRunOutputTokens}) is subtracted: those tokens are inside the
+ * baseline's kept messages AND in the response message's `tokenCount` the client
+ * adds on top, so leaving them in the marker double-counts them on a tool-loop
+ * summarized turn. Single-call turns pass 0 and are unaffected.
+ */
+export function computeSummaryUsedTokens(
+ snapshot: TContextUsageEvent | null | undefined,
+ priorOutputTokens = 0,
+): number | undefined {
+ const summaryTokens = snapshot?.breakdown?.summaryTokens ?? 0;
+ if (!snapshot || summaryTokens <= 0) {
+ return undefined;
+ }
+ const maxTokens = snapshot.contextBudget ?? snapshot.breakdown.maxContextTokens ?? 0;
+ const baseUsed =
+ snapshot.remainingContextTokens != null
+ ? maxTokens - snapshot.remainingContextTokens
+ : (snapshot.effectiveInstructionTokens ?? snapshot.breakdown.instructionTokens ?? 0) +
+ summaryTokens +
+ (snapshot.breakdown.messageTokens ?? 0);
+ const adjusted = baseUsed - Math.max(0, priorOutputTokens);
+ return adjusted > 0 ? Math.round(adjusted) : undefined;
+}
+
function parseUsageEvents(value?: string | null): TTokenUsageEvent[] {
if (typeof value !== 'string' || value.length === 0) {
return [];
@@ -309,21 +393,50 @@ function parseUsageEvents(value?: string | null): TTokenUsageEvent[] {
* reload (finding: stopped responses otherwise lose cost). Shared by every abort
* save path (agents abort route + legacy abort middleware).
*
- * Deliberately persists ONLY `usage`, not `contextUsage`: unlike the live path,
- * the abort path can't tell whether the FINAL call (the one the latest snapshot
- * precedes) emitted usage — the job stores only the latest snapshot, not the
- * snapshot count. If the final call emitted none, `completedOutputTokens` would
- * reuse an earlier call's output the snapshot already counts → reload
- * over-reports. A stopped/incomplete response therefore falls back to the coarse
- * per-message gauge estimate on reload, which is both safe and apt for an
- * interrupted turn that never reached a clean pre-invoke breakdown.
+ * Deliberately omits the full `contextUsage`: unlike the live path, the abort
+ * path can't tell whether the FINAL call (the one the latest snapshot precedes)
+ * emitted usage — the job stores only the latest snapshot, not the snapshot
+ * count. If the final call emitted none, `completedOutputTokens` would reuse an
+ * earlier call's output the snapshot already counts → reload over-reports. So a
+ * stopped response falls back to the per-message gauge estimate on reload.
+ *
+ * It DOES persist the `summaryUsedTokens` marker when the stopped turn had
+ * summarized: that marker is pre-invoke (no `completedOutputTokens` ambiguity),
+ * and without it the fallback estimate re-sums the history the compaction
+ * discarded — leaving a stopped summarized turn pinned at 100%. Unlike the live
+ * path, the abort `tokenCount` comes from `countTokens(text)` (abortMiddleware) or
+ * is absent (agents abort route) — it does NOT fold in the summarization or
+ * earlier-call output the way `recordCollectedUsage` does. So the marker subtracts
+ * NOTHING: the full pre-invoke baseline is correct, and the client adds only the
+ * partial answer text on top (no overlap to cancel).
*/
export function buildAbortedResponseMetadata(
- job: { tokenUsage?: string | null } | null | undefined,
-): { usage?: TResponseUsage } | undefined {
+ job: { tokenUsage?: string | null; contextUsage?: string | null } | null | undefined,
+): { usage?: TResponseUsage; summaryUsedTokens?: number } | undefined {
const events = parseUsageEvents(job?.tokenUsage);
const usage = aggregateEmittedUsage(events);
- return usage ? { usage } : undefined;
+
+ let snapshot: TContextUsageEvent | null = null;
+ if (typeof job?.contextUsage === 'string' && job.contextUsage.length > 0) {
+ try {
+ snapshot = JSON.parse(job.contextUsage) as TContextUsageEvent;
+ } catch {
+ snapshot = null;
+ }
+ }
+ /** Subtract nothing: the abort `tokenCount` (countTokens(text) or absent) does
+ * not fold in summarization/earlier-call output, so the full baseline is the
+ * marker and the client's partial-text addition has no overlap to cancel. */
+ const summaryUsedTokens = computeSummaryUsedTokens(snapshot);
+
+ const metadata: { usage?: TResponseUsage; summaryUsedTokens?: number } = {};
+ if (usage) {
+ metadata.usage = usage;
+ }
+ if (summaryUsedTokens != null) {
+ metadata.summaryUsedTokens = summaryUsedTokens;
+ }
+ return Object.keys(metadata).length > 0 ? metadata : undefined;
}
/**