diff --git a/apps/server/src/provider/Layers/ClaudeAdapter.test.ts b/apps/server/src/provider/Layers/ClaudeAdapter.test.ts index 68ef37d1a..7d97ef558 100644 --- a/apps/server/src/provider/Layers/ClaudeAdapter.test.ts +++ b/apps/server/src/provider/Layers/ClaudeAdapter.test.ts @@ -247,9 +247,14 @@ function makeContextUsageResponse(input: { readonly maxTokens: number; readonly rawMaxTokens?: number; readonly isAutoCompactEnabled?: boolean; + readonly categories?: SDKControlGetContextUsageResponse["categories"]; }): SDKControlGetContextUsageResponse { return { - categories: [], + // A category-less response is treated as a partially started process and + // ignored by the adapter, so a usable fixture needs at least one category. + categories: input.categories ?? [ + { name: "Messages", tokens: input.totalTokens, color: "#000000" }, + ], totalTokens: input.totalTokens, maxTokens: input.maxTokens, rawMaxTokens: input.rawMaxTokens ?? input.maxTokens, @@ -5122,6 +5127,66 @@ describe("ClaudeAdapterLive", () => { ); }); + it.effect( + "ignores a category-less context usage report instead of emitting a bare snapshot", + () => { + const harness = makeHarness(); + return Effect.gen(function* () { + const adapter = yield* ClaudeAdapter; + + const runtimeEventsFiber = yield* Stream.take(adapter.streamEvents, 6).pipe( + Stream.runCollect, + Effect.forkChild, + ); + + yield* adapter.startSession({ + threadId: THREAD_ID, + provider: ProviderDriverKind.make("claudeAgent"), + runtimeMode: "full-access", + }); + + yield* adapter.sendTurn({ + threadId: THREAD_ID, + input: "hello", + attachments: [], + }); + + // A freshly resumed SDK process reports a token estimate before it has + // rebuilt its context breakdown. Emitting that would replace the last + // rich snapshot with a bare one, so no usage event may be produced. + harness.query.setContextUsageResponse( + makeContextUsageResponse({ + totalTokens: 93169, + maxTokens: 1000000, + categories: [], + }), + ); + harness.query.emit({ + type: "result", + subtype: "success", + is_error: false, + duration_ms: 100, + duration_api_ms: 90, + num_turns: 1, + result: "done", + stop_reason: "end_turn", + session_id: "sdk-session-bare-context", + } as unknown as SDKMessage); + harness.query.finish(); + + const runtimeEvents = Array.from(yield* Fiber.join(runtimeEventsFiber)); + assert.deepEqual( + runtimeEvents.filter((event) => event.type === "thread.token-usage.updated"), + [], + ); + assert.equal(harness.query.getContextUsageCalls.length, 1); + }).pipe( + Effect.provideService(Random.Random, makeDeterministicRandomService()), + Effect.provide(harness.layer), + ); + }, + ); + it.effect("emits Claude context window on result completion usage snapshots", () => { const harness = makeHarness(); return Effect.gen(function* () { @@ -5187,6 +5252,7 @@ describe("ClaudeAdapterLive", () => { cachedInputTokens: 21144, outputTokens: 679, maxTokens: 200000, + contextCategories: [{ name: "Messages", tokens: 24542 }], lastInputTokens: 23863, lastCachedInputTokens: 21144, lastOutputTokens: 679, @@ -5274,6 +5340,7 @@ describe("ClaudeAdapterLive", () => { usedTokens: 22000, lastUsedTokens: 22000, maxTokens: 200000, + contextCategories: [{ name: "Messages", tokens: 22000 }], compactsAutomatically: true, }, }); @@ -5285,6 +5352,7 @@ describe("ClaudeAdapterLive", () => { cachedInputTokens: 20, outputTokens: 80, maxTokens: 200000, + contextCategories: [{ name: "Messages", tokens: 24000 }], lastInputTokens: 120, lastCachedInputTokens: 20, lastOutputTokens: 80, @@ -5332,14 +5400,25 @@ describe("ClaudeAdapterLive", () => { attachments: [], }); + // Both snapshots carry an equal-but-not-identical category array: the + // dedupe has to compare the entries, not the array reference. + const categories = [ + { name: "System prompt", tokens: 3000, color: "#111111" }, + { name: "MCP tools", tokens: 700, color: "#444444", isDeferred: true }, + { name: "System tools (deferred)", tokens: 5000, color: "#555555" }, + { name: "Messages", tokens: 19000, color: "#222222" }, + { name: "Free space", tokens: 178000, color: "#333333" }, + ] satisfies SDKControlGetContextUsageResponse["categories"]; harness.query.setContextUsageResponses([ makeContextUsageResponse({ totalTokens: 22000, maxTokens: 200000, + categories: categories.map((category) => ({ ...category })), }), makeContextUsageResponse({ totalTokens: 22000, maxTokens: 200000, + categories: categories.map((category) => ({ ...category })), }), ]); @@ -5377,6 +5456,12 @@ describe("ClaudeAdapterLive", () => { usedTokens: 22000, lastUsedTokens: 22000, maxTokens: 200000, + // Free space is derived client-side and deferred categories are not + // part of the used total, so both are filtered out here. + contextCategories: [ + { name: "System prompt", tokens: 3000 }, + { name: "Messages", tokens: 19000 }, + ], compactsAutomatically: true, }, }); @@ -5755,6 +5840,7 @@ describe("ClaudeAdapterLive", () => { lastOutputTokens: 70, lastReasoningOutputTokens: 50, maxTokens: 1000000, + contextCategories: [{ name: "Messages", tokens: 80 }], compactsAutomatically: true, }, }); diff --git a/apps/server/src/provider/Layers/ClaudeAdapter.ts b/apps/server/src/provider/Layers/ClaudeAdapter.ts index 9e1ad53bf..a94130694 100644 --- a/apps/server/src/provider/Layers/ClaudeAdapter.ts +++ b/apps/server/src/provider/Layers/ClaudeAdapter.ts @@ -777,6 +777,49 @@ function readClaudeUsageTotals(value: unknown): ClaudeUsageTotals | undefined { }; } +/** Free space is derived client-side from maxTokens - usedTokens, so keeping the + * SDK's own free-space category would double-count it in the breakdown. */ +function isClaudeFreeSpaceCategory(name: string): boolean { + return name.trim().toLowerCase() === "free space"; +} + +/** The SDK marks deferred categories with `isDeferred`, but has also been seen + * tagging the name itself ("MCP tools (deferred)"); match either. */ +function isClaudeDeferredCategory( + category: SDKControlGetContextUsageResponse["categories"][number], +): boolean { + return category.isDeferred === true || category.name.trim().toLowerCase().endsWith("(deferred)"); +} + +/** Keeps the SDK's category order: legend colors are assigned by index, so + * reordering would make them flicker between updates. */ +function normalizeClaudeContextCategories( + value: SDKControlGetContextUsageResponse["categories"] | undefined, +): ThreadTokenUsageSnapshot["contextCategories"] { + if (!Array.isArray(value)) { + return undefined; + } + + const categories = value.flatMap((category) => { + const name = typeof category?.name === "string" ? category.name.trim() : ""; + const tokens = asPositiveFiniteInteger(category?.tokens); + // Deferred categories (tool schemas not yet loaded) are excluded from the + // SDK's own used-token total, so listing them would make the breakdown sum + // past the reported usage. + if ( + name.length === 0 || + tokens === undefined || + isClaudeDeferredCategory(category) || + isClaudeFreeSpaceCategory(name) + ) { + return []; + } + return [{ name, tokens }]; + }); + + return categories.length > 0 ? categories : undefined; +} + function normalizeClaudeContextUsage( value: SDKControlGetContextUsageResponse | undefined, ): ThreadTokenUsageSnapshot | undefined { @@ -791,15 +834,27 @@ function normalizeClaudeContextUsage( const maxTokens = asPositiveFiniteInteger(value.maxTokens) ?? asPositiveFiniteInteger(value.rawMaxTokens); + const contextCategories = normalizeClaudeContextCategories(value.categories); + // A fully initialized Claude session always reports at least a system-prompt + // and messages category. A response without any is a partially started + // process (e.g. queried during session resume) whose token estimate would + // replace the last known rich snapshot with a bare one — treat it as + // unavailable and let the caller fall back instead. + if (contextCategories === undefined) { + return undefined; + } return { usedTokens, lastUsedTokens: usedTokens, ...(maxTokens !== undefined ? { maxTokens } : {}), + contextCategories, compactsAutomatically: value.isAutoCompactEnabled, }; } +/** Scalar keys only: `contextCategories` is an array, so it needs the + * element-wise comparison below instead of an identity check. */ const THREAD_TOKEN_USAGE_SNAPSHOT_KEYS = [ "usedTokens", "totalProcessedTokens", @@ -818,6 +873,22 @@ const THREAD_TOKEN_USAGE_SNAPSHOT_KEYS = [ "compactsAutomatically", ] as const satisfies ReadonlyArray; +function areClaudeContextCategoriesEqual( + left: ThreadTokenUsageSnapshot["contextCategories"], + right: ThreadTokenUsageSnapshot["contextCategories"], +): boolean { + if (left === undefined || right === undefined) { + return left === right; + } + if (left.length !== right.length) { + return false; + } + return left.every((category, index) => { + const other = right[index]; + return other !== undefined && category.name === other.name && category.tokens === other.tokens; + }); +} + function areThreadTokenUsageSnapshotsEqual( left: ThreadTokenUsageSnapshot | undefined, right: ThreadTokenUsageSnapshot | undefined, @@ -825,7 +896,10 @@ function areThreadTokenUsageSnapshotsEqual( if (!left || !right) { return false; } - return THREAD_TOKEN_USAGE_SNAPSHOT_KEYS.every((key) => left[key] === right[key]); + return ( + THREAD_TOKEN_USAGE_SNAPSHOT_KEYS.every((key) => left[key] === right[key]) && + areClaudeContextCategoriesEqual(left.contextCategories, right.contextCategories) + ); } function normalizeClaudeCompactBoundaryUsage( diff --git a/apps/server/src/provider/Layers/CodexAdapter.test.ts b/apps/server/src/provider/Layers/CodexAdapter.test.ts index 68c8cc591..10816965b 100644 --- a/apps/server/src/provider/Layers/CodexAdapter.test.ts +++ b/apps/server/src/provider/Layers/CodexAdapter.test.ts @@ -2430,12 +2430,14 @@ lifecycleLayer("CodexAdapterLive lifecycle", (it) => { return; } + // Cumulative fields track the session-wide `total`; the `last*` fields and + // usedTokens (the current context size) track the latest turn. assert.deepEqual(firstEvent.value.payload.usage, { usedTokens: 126, totalProcessedTokens: 11_839, maxTokens: 258_400, - inputTokens: 120, - cachedInputTokens: 0, + inputTokens: 11_833, + cachedInputTokens: 3456, outputTokens: 6, reasoningOutputTokens: 0, lastUsedTokens: 126, diff --git a/apps/server/src/provider/Layers/CodexAdapter.ts b/apps/server/src/provider/Layers/CodexAdapter.ts index 29c87bde3..b8a4a6c93 100644 --- a/apps/server/src/provider/Layers/CodexAdapter.ts +++ b/apps/server/src/provider/Layers/CodexAdapter.ts @@ -654,10 +654,17 @@ function normalizeCodexTokenUsage( } const maxTokens = usage.modelContextWindow ?? undefined; - const inputTokens = usage.last.inputTokens; - const cachedInputTokens = usage.last.cachedInputTokens; - const outputTokens = usage.last.outputTokens; - const reasoningOutputTokens = usage.last.reasoningOutputTokens; + // Cumulative fields come from `total` (session-wide, matching the Claude + // adapter's semantics); the `last*` fields stay on the latest turn. usedTokens + // remains `last.totalTokens` because that is the current context size. + const inputTokens = usage.total.inputTokens; + const cachedInputTokens = usage.total.cachedInputTokens; + const outputTokens = usage.total.outputTokens; + const reasoningOutputTokens = usage.total.reasoningOutputTokens; + const lastInputTokens = usage.last.inputTokens; + const lastCachedInputTokens = usage.last.cachedInputTokens; + const lastOutputTokens = usage.last.outputTokens; + const lastReasoningOutputTokens = usage.last.reasoningOutputTokens; return { usedTokens, @@ -670,12 +677,10 @@ function normalizeCodexTokenUsage( ...(outputTokens !== undefined ? { outputTokens } : {}), ...(reasoningOutputTokens !== undefined ? { reasoningOutputTokens } : {}), ...(usedTokens !== undefined ? { lastUsedTokens: usedTokens } : {}), - ...(inputTokens !== undefined ? { lastInputTokens: inputTokens } : {}), - ...(cachedInputTokens !== undefined ? { lastCachedInputTokens: cachedInputTokens } : {}), - ...(outputTokens !== undefined ? { lastOutputTokens: outputTokens } : {}), - ...(reasoningOutputTokens !== undefined - ? { lastReasoningOutputTokens: reasoningOutputTokens } - : {}), + ...(lastInputTokens !== undefined ? { lastInputTokens } : {}), + ...(lastCachedInputTokens !== undefined ? { lastCachedInputTokens } : {}), + ...(lastOutputTokens !== undefined ? { lastOutputTokens } : {}), + ...(lastReasoningOutputTokens !== undefined ? { lastReasoningOutputTokens } : {}), compactsAutomatically: true, }; } diff --git a/apps/web/src/components/chat/ContextWindowMeter.browser.tsx b/apps/web/src/components/chat/ContextWindowMeter.browser.tsx index 1389104b1..fa3cacbeb 100644 --- a/apps/web/src/components/chat/ContextWindowMeter.browser.tsx +++ b/apps/web/src/components/chat/ContextWindowMeter.browser.tsx @@ -27,9 +27,16 @@ const TEST_CONTEXT_WINDOW: ContextWindowSnapshot = { toolUses: null, durationMs: null, compactsAutomatically: false, + contextCategories: null, updatedAt: "2026-06-19T12:00:00.000Z", }; +function contextSegmentNames(): ReadonlyArray { + return Array.from(document.querySelectorAll("[data-context-segment]")).map( + (segment) => segment.getAttribute("data-context-segment") ?? "", + ); +} + const TEST_ACCOUNT_USAGE: ProviderAccountUsagePresentation = { label: "Codex usage", reachedLimit: false, @@ -106,6 +113,68 @@ describe("ContextWindowMeter", () => { } }); + it("breaks the context window down by category behind a collapsed toggle", async () => { + const screen = await render( + , + ); + + try { + await page.getByRole("button", { name: /Context window/ }).click(); + + await expect.element(page.getByText(/1.2m of 1.3m input tokens/)).toBeVisible(); + expect(contextSegmentNames()).toEqual(["Messages", "System tools"]); + + const breakdownToggle = page.getByRole("button", { name: "Show context breakdown" }); + await expect.element(page.getByText("System tools")).not.toBeInTheDocument(); + + await breakdownToggle.click(); + + await expect.element(page.getByText("Messages")).toBeVisible(); + await expect.element(page.getByText("System tools")).toBeVisible(); + await expect.element(page.getByText("Free space")).toBeVisible(); + + await breakdownToggle.click(); + await expect.element(page.getByText("Free space")).not.toBeInTheDocument(); + } finally { + await screen.unmount(); + } + }); + + it("shows a single used segment when the provider reports no categories", async () => { + const screen = await render( + , + ); + + try { + await page.getByRole("button", { name: /Context window/ }).click(); + + expect(contextSegmentNames()).toEqual(["Used"]); + await expect + .element(page.getByRole("button", { name: "Show context breakdown" })) + .not.toBeInTheDocument(); + await expect.element(page.getByText(/29.2% ⋅ 3.4k of 11.8k input tokens/)).toBeVisible(); + } finally { + await screen.unmount(); + } + }); + it("keeps manual compaction enabled at low context usage", async () => { const onCompactContext = vi.fn(); const usedPercentage = (44_272 / 258_400) * 100; diff --git a/apps/web/src/components/chat/ContextWindowMeter.tsx b/apps/web/src/components/chat/ContextWindowMeter.tsx index 8d3818d01..618de0f17 100644 --- a/apps/web/src/components/chat/ContextWindowMeter.tsx +++ b/apps/web/src/components/chat/ContextWindowMeter.tsx @@ -1,7 +1,13 @@ +import { ChevronDownIcon } from "lucide-react"; import { useCallback, useLayoutEffect, useRef, useState } from "react"; import { cn } from "~/lib/utils"; -import { type ContextWindowSnapshot, formatContextWindowTokens } from "~/lib/contextWindow"; +import { + type ContextWindowSnapshot, + deriveCachedInputRate, + formatContextWindowPercentage, + formatContextWindowTokens, +} from "~/lib/contextWindow"; import { isProviderUsageNearLimit, type ProviderAccountUsagePresentation, @@ -9,16 +15,6 @@ import { import { Popover, PopoverPopup, PopoverTrigger } from "../ui/popover"; import { Tooltip, TooltipPopup, TooltipTrigger } from "../ui/tooltip"; -function formatPercentage(value: number | null): string | null { - if (value === null || !Number.isFinite(value)) { - return null; - } - if (value < 10) { - return `${value.toFixed(1).replace(/\.0$/, "")}%`; - } - return `${Math.round(value)}%`; -} - function isDomNode(value: unknown): value is Node { return typeof value === "object" && value !== null && "nodeType" in value; } @@ -59,6 +55,92 @@ function AccountUsageBar(props: { ); } +/** Categorical swatches assigned by index (cycled), so a category keeps its + * color as long as the provider keeps its order. */ +const CONTEXT_CATEGORY_COLOR_CLASS_NAMES = [ + "bg-context-cat-1", + "bg-context-cat-2", + "bg-context-cat-3", + "bg-context-cat-4", + "bg-context-cat-5", + "bg-context-cat-6", + "bg-context-cat-7", + "bg-context-cat-8", +] as const; + +type ContextBreakdownSegment = { + readonly key: string; + readonly name: string; + readonly tokens: number; + readonly percentage: number; + readonly colorClassName: string; +}; + +/** Providers without a category breakdown (Codex) collapse to one used-vs-free + * segment so the bar renders through the same path. */ +function buildContextBreakdownSegments( + usage: ContextWindowSnapshot | null, +): ReadonlyArray { + const maxTokens = usage?.maxTokens ?? null; + if (!usage || maxTokens === null || maxTokens <= 0) { + return []; + } + + const toPercentage = (tokens: number) => Math.max(0, Math.min(100, (tokens / maxTokens) * 100)); + const categories = usage.contextCategories ?? null; + if (categories && categories.length > 0) { + // Colors are keyed to the provider's order (stable across updates), then + // the bar and legend display largest-first. + return categories + .map((category, index) => ({ + key: `${index}-${category.name}`, + name: category.name, + tokens: category.tokens, + percentage: toPercentage(category.tokens), + colorClassName: + CONTEXT_CATEGORY_COLOR_CLASS_NAMES[index % CONTEXT_CATEGORY_COLOR_CLASS_NAMES.length] ?? + CONTEXT_CATEGORY_COLOR_CLASS_NAMES[0], + })) + .sort((left, right) => right.tokens - left.tokens); + } + + if (usage.usedTokens <= 0) { + return []; + } + return [ + { + key: "used", + name: "Used", + tokens: usage.usedTokens, + percentage: usage.usedPercentage ?? toPercentage(usage.usedTokens), + colorClassName: "bg-primary", + }, + ]; +} + +function ContextBreakdownLegendRow(props: { + swatchClassName: string; + name: string; + tokens: number | null; + percentage: number | null; +}) { + return ( +
+
+ ); +} + export function ContextWindowMeter(props: { usage: ContextWindowSnapshot | null; accountUsage?: ProviderAccountUsagePresentation | null; @@ -77,9 +159,18 @@ export function ContextWindowMeter(props: { const [isOpen, setIsOpen] = useState(false); const [isPinnedOpen, setIsPinnedOpen] = useState(false); const [isHoverOpenSuppressed, setIsHoverOpenSuppressed] = useState(false); + const [isBreakdownOpen, setIsBreakdownOpen] = useState(false); const accountUsage = props.accountUsage ?? null; + const breakdownSegments = buildContextBreakdownSegments(usage); + const hasContextCategories = (usage?.contextCategories?.length ?? 0) > 0; + const cachedInputRate = deriveCachedInputRate(usage); + const cachedInputDetail = cachedInputRate + ? `${formatContextWindowPercentage(cachedInputRate.percentage) ?? "0%"} ⋅ ${formatContextWindowTokens( + cachedInputRate.cachedTokens, + )} of ${formatContextWindowTokens(cachedInputRate.inputTokens)} input tokens` + : null; const usageNearLimit = isProviderUsageNearLimit(accountUsage); - const usedPercentage = formatPercentage(usage?.usedPercentage ?? null); + const usedPercentage = formatContextWindowPercentage(usage?.usedPercentage ?? null); const normalizedPercentage = Math.max(0, Math.min(100, usage?.usedPercentage ?? 0)); const radius = 9.75; const circumference = 2 * Math.PI * radius; @@ -222,7 +313,7 @@ export function ContextWindowMeter(props: { )} > {usage.usedPercentage !== null - ? Math.round(usage.usedPercentage) + ? Math.trunc(usage.usedPercentage) : formatContextWindowTokens(usage.usedTokens)} ) : null} @@ -248,29 +339,82 @@ export function ContextWindowMeter(props: {
Context window
- {!usage ? ( -
- No tokens used yet - {contextWindowLabel ? ( - <> - ⋅ - {contextWindowLabel} window - - ) : null} +
+ {!usage ? ( +
+ No tokens used yet + {contextWindowLabel ? ( + <> + ⋅ + {contextWindowLabel} window + + ) : null} +
+ ) : usage.maxTokens !== null && usedPercentage ? ( +
+ {usedPercentage} + ⋅ + {formatContextWindowTokens(usage.usedTokens)} + / + {formatContextWindowTokens(usage.maxTokens ?? null)} context used +
+ ) : ( +
+ {formatContextWindowTokens(usage.usedTokens)} tokens used so far +
+ )} + {hasContextCategories ? ( + + ) : null} +
+ {breakdownSegments.length > 0 ? ( +