diff --git a/apps/code/src/renderer/features/message-editor/components/ContextUsageIndicator.tsx b/apps/code/src/renderer/features/message-editor/components/ContextUsageIndicator.tsx deleted file mode 100644 index 99e121a1e..000000000 --- a/apps/code/src/renderer/features/message-editor/components/ContextUsageIndicator.tsx +++ /dev/null @@ -1,41 +0,0 @@ -import { Text, Tooltip } from "@radix-ui/themes"; -import { useContextUsageForTask } from "@renderer/features/sessions/hooks/useSession"; - -const CONTEXT_WARNING_THRESHOLD_PCT = 40; - -interface ContextUsageIndicatorProps { - taskId?: string; -} - -export function ContextUsageIndicator({ taskId }: ContextUsageIndicatorProps) { - const contextUsage = useContextUsageForTask(taskId); - if (!contextUsage || contextUsage.size <= 0) return null; - - const percent = Math.round((contextUsage.used / contextUsage.size) * 100); - - if (percent < CONTEXT_WARNING_THRESHOLD_PCT) return null; - - return ( - - - {percent}% - - - ); -} - -function getContextColor(percent: number): string { - if (percent >= 80) return "var(--red-9)"; - if (percent >= 50) return "var(--yellow-11)"; - return "var(--green-9)"; -} diff --git a/apps/code/src/renderer/features/message-editor/components/EditorToolbar.tsx b/apps/code/src/renderer/features/message-editor/components/EditorToolbar.tsx index 8149a0b0d..2f5efb8eb 100644 --- a/apps/code/src/renderer/features/message-editor/components/EditorToolbar.tsx +++ b/apps/code/src/renderer/features/message-editor/components/EditorToolbar.tsx @@ -2,7 +2,6 @@ import { ModelSelector } from "@features/sessions/components/ModelSelector"; import { Flex } from "@radix-ui/themes"; import type { FileAttachment, MentionChip } from "../utils/content"; import { AttachmentMenu } from "./AttachmentMenu"; -import { ContextUsageIndicator } from "./ContextUsageIndicator"; interface EditorToolbarProps { disabled?: boolean; @@ -43,7 +42,6 @@ export function EditorToolbar({ {!hideSelectors && ( )} - ); } diff --git a/apps/code/src/renderer/features/sessions/components/ContextUsageIndicator.tsx b/apps/code/src/renderer/features/sessions/components/ContextUsageIndicator.tsx new file mode 100644 index 000000000..ce1411a1d --- /dev/null +++ b/apps/code/src/renderer/features/sessions/components/ContextUsageIndicator.tsx @@ -0,0 +1,78 @@ +import { Tooltip } from "@components/ui/Tooltip"; +import type { ContextUsage } from "@features/sessions/hooks/useContextUsage"; +import { Flex, Text } from "@radix-ui/themes"; + +function formatTokensCompact(tokens: number): string { + if (tokens >= 1_000_000) { + return `${(tokens / 1_000_000).toFixed(1)}M`; + } + return `${Math.round(tokens / 1000)}K`; +} + +function formatTokensFull(tokens: number): string { + return tokens.toLocaleString(); +} + +function getUsageColor(percentage: number): string { + if (percentage >= 90) return "var(--red-9)"; + if (percentage >= 75) return "var(--orange-9)"; + if (percentage >= 50) return "var(--amber-9)"; + return "var(--green-9)"; +} + +const CIRCLE_SIZE = 20; +const STROKE_WIDTH = 2.5; +const RADIUS = (CIRCLE_SIZE - STROKE_WIDTH) / 2; +const CIRCUMFERENCE = 2 * Math.PI * RADIUS; + +interface ContextUsageIndicatorProps { + usage: ContextUsage | null; +} + +export function ContextUsageIndicator({ usage }: ContextUsageIndicatorProps) { + if (!usage) return null; + + const { used, size, percentage } = usage; + const strokeDashoffset = CIRCUMFERENCE - (percentage / 100) * CIRCUMFERENCE; + const color = getUsageColor(percentage); + + return ( + + + + + + + + {formatTokensCompact(used)}/{formatTokensCompact(size)} + + + + ); +} diff --git a/apps/code/src/renderer/features/sessions/components/ConversationView.tsx b/apps/code/src/renderer/features/sessions/components/ConversationView.tsx index 5ef58cb37..403e013c9 100644 --- a/apps/code/src/renderer/features/sessions/components/ConversationView.tsx +++ b/apps/code/src/renderer/features/sessions/components/ConversationView.tsx @@ -1,3 +1,4 @@ +import { useContextUsage } from "@features/sessions/hooks/useContextUsage"; import { sessionStoreSetters, useOptimisticItemsForTask, @@ -52,6 +53,7 @@ export function ConversationView({ const agentLogsEnabled = useFeatureFlag("posthog-code-background-agent-logs"); const debugLogsCloudRuns = useSettingsStore((s) => s.debugLogsCloudRuns); const showDebugLogs = agentLogsEnabled && debugLogsCloudRuns; + const contextUsage = useContextUsage(events); const { items: conversationItems, @@ -229,6 +231,7 @@ export function ConversationView({ hasPendingPermission={pendingPermissionsCount > 0} pausedDurationMs={pausedDurationMs} isCompacting={isCompacting} + usage={contextUsage} /> } diff --git a/apps/code/src/renderer/features/sessions/components/SessionFooter.tsx b/apps/code/src/renderer/features/sessions/components/SessionFooter.tsx index 4f22a07bd..f32068d56 100644 --- a/apps/code/src/renderer/features/sessions/components/SessionFooter.tsx +++ b/apps/code/src/renderer/features/sessions/components/SessionFooter.tsx @@ -1,6 +1,8 @@ +import type { ContextUsage } from "@features/sessions/hooks/useContextUsage"; import { Pause } from "@phosphor-icons/react"; import { Box, Flex, Text } from "@radix-ui/themes"; +import { ContextUsageIndicator } from "./ContextUsageIndicator"; import { formatDuration, GeneratingIndicator } from "./GeneratingIndicator"; interface SessionFooterProps { @@ -12,6 +14,7 @@ interface SessionFooterProps { hasPendingPermission?: boolean; pausedDurationMs?: number; isCompacting?: boolean; + usage?: ContextUsage | null; } export function SessionFooter({ @@ -23,20 +26,23 @@ export function SessionFooter({ hasPendingPermission = false, pausedDurationMs, isCompacting = false, + usage, }: SessionFooterProps) { if (isPromptPending && !isCompacting) { - // Show static "waiting" state when permission is pending if (hasPendingPermission) { return ( - - - Awaiting permission... + + + + Awaiting permission... + + ); @@ -44,16 +50,19 @@ export function SessionFooter({ return ( - - - {queuedCount > 0 && ( - - ({queuedCount} queued) - - )} + + + + {queuedCount > 0 && ( + + ({queuedCount} queued) + + )} + + ); @@ -69,13 +78,26 @@ export function SessionFooter({ ) { return ( - - Generated in {formatDuration(lastGenerationDuration)} - + + + Generated in {formatDuration(lastGenerationDuration)} + + + + + ); + } + + if (usage) { + return ( + + + + ); } diff --git a/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts b/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts index 38ae968dd..7c0fa9a4b 100644 --- a/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts +++ b/apps/code/src/renderer/features/sessions/components/buildConversationItems.ts @@ -344,12 +344,14 @@ function handleNotification( const params = msg.params as { trigger: "manual" | "auto"; preTokens: number; + contextSize?: number; }; markCompactingStatusComplete(b); pushItem(b, { sessionUpdate: "compact_boundary", trigger: params.trigger, preTokens: params.preTokens, + contextSize: params.contextSize, }); return; } @@ -549,6 +551,7 @@ function processSessionUpdate(b: ItemBuilder, update: SessionUpdate) { case "plan": case "available_commands_update": case "config_option_update": + case "usage_update": break; default: { diff --git a/apps/code/src/renderer/features/sessions/components/session-update/CompactBoundaryView.tsx b/apps/code/src/renderer/features/sessions/components/session-update/CompactBoundaryView.tsx index d4b8cf3f7..320ccd3fd 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/CompactBoundaryView.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/CompactBoundaryView.tsx @@ -4,13 +4,19 @@ import { Badge, Box, Flex, Text } from "@radix-ui/themes"; interface CompactBoundaryViewProps { trigger: "manual" | "auto"; preTokens: number; + contextSize?: number; } export function CompactBoundaryView({ trigger, preTokens, + contextSize, }: CompactBoundaryViewProps) { const tokensK = Math.round(preTokens / 1000); + const percent = + contextSize && contextSize > 0 + ? Math.round((preTokens / contextSize) * 100) + : null; return ( @@ -27,7 +33,9 @@ export function CompactBoundaryView({ {trigger} - (~{tokensK}K tokens summarized) + {percent !== null + ? `(${percent}% of context ยท ~${tokensK}K tokens summarized)` + : `(~${tokensK}K tokens summarized)`} diff --git a/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx b/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx index 5a34c5b21..51fa6112e 100644 --- a/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx +++ b/apps/code/src/renderer/features/sessions/components/session-update/SessionUpdateView.tsx @@ -23,6 +23,7 @@ export type RenderItem = sessionUpdate: "compact_boundary"; trigger: "manual" | "auto"; preTokens: number; + contextSize?: number; } | { sessionUpdate: "status"; @@ -101,6 +102,7 @@ export const SessionUpdateView = memo(function SessionUpdateView({ ); case "status": diff --git a/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts b/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts new file mode 100644 index 000000000..b340e9abb --- /dev/null +++ b/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts @@ -0,0 +1,59 @@ +import type { AcpMessage } from "@shared/types/session-events"; +import { useMemo } from "react"; + +export interface ContextUsage { + used: number; + size: number; + percentage: number; + cost: { amount: number; currency: string } | null; +} + +/** + * Extract the latest context window usage from session events. + * Scans backwards to find the most recent usage_update notification. + * Re-derives on each new event, giving live updates during streaming. + */ +export function useContextUsage(events: AcpMessage[]): ContextUsage | null { + return useMemo(() => extractContextUsage(events), [events]); +} + +export function extractContextUsage(events: AcpMessage[]): ContextUsage | null { + for (let i = events.length - 1; i >= 0; i--) { + const msg = events[i].message; + if ( + "method" in msg && + msg.method === "session/update" && + !("id" in msg) && + "params" in msg + ) { + const params = msg.params as + | { + update?: { + sessionUpdate?: string; + used?: number; + size?: number; + cost?: { amount: number; currency: string } | null; + }; + } + | undefined; + const update = params?.update; + if ( + update?.sessionUpdate === "usage_update" && + typeof update.used === "number" && + typeof update.size === "number" + ) { + const percentage = + update.size > 0 + ? Math.min(100, Math.round((update.used / update.size) * 100)) + : 0; + return { + used: update.used, + size: update.size, + percentage, + cost: update.cost ?? null, + }; + } + } + } + return null; +} diff --git a/apps/code/src/renderer/features/sessions/hooks/useSession.ts b/apps/code/src/renderer/features/sessions/hooks/useSession.ts index fa6f7cae0..246b2d0a1 100644 --- a/apps/code/src/renderer/features/sessions/hooks/useSession.ts +++ b/apps/code/src/renderer/features/sessions/hooks/useSession.ts @@ -147,25 +147,6 @@ export const useThoughtLevelConfigOptionForTask = ( return useConfigOptionForTask(taskId, "thought_level"); }; -/** Get context window usage for a task (used / size) */ -export const useContextUsageForTask = ( - taskId: string | undefined, -): { used: number; size: number } | undefined => { - return useSessionStore((s) => { - if (!taskId) return undefined; - const taskRunId = s.taskIdIndex[taskId]; - if (!taskRunId) return undefined; - const session = s.sessions[taskRunId]; - if ( - session?.contextUsed === undefined || - session?.contextSize === undefined - ) { - return undefined; - } - return { used: session.contextUsed, size: session.contextSize }; - }); -}; - /** Get the adapter type for a task */ export const useAdapterForTask = ( taskId: string | undefined, diff --git a/packages/agent/src/adapters/base-acp-agent.ts b/packages/agent/src/adapters/base-acp-agent.ts index 93db4e3ee..792efe416 100644 --- a/packages/agent/src/adapters/base-acp-agent.ts +++ b/packages/agent/src/adapters/base-acp-agent.ts @@ -20,6 +20,7 @@ import { DEFAULT_GATEWAY_MODEL, fetchGatewayModels, formatGatewayModelName, + type GatewayModel, isAnthropicModel, } from "../gateway-models"; import { Logger } from "../utils/logger"; @@ -33,6 +34,8 @@ export interface BaseSession { settingsManager: SettingsManager; } +const DEFAULT_CONTEXT_WINDOW = 200_000; + export abstract class BaseAcpAgent implements Agent { abstract readonly adapterName: string; protected session!: BaseSession; @@ -40,6 +43,7 @@ export abstract class BaseAcpAgent implements Agent { client: AgentSideConnection; logger: Logger; fileContentCache: { [key: string]: string } = {}; + protected gatewayModels: GatewayModel[] = []; constructor(client: AgentSideConnection) { this.client = client; @@ -119,9 +123,9 @@ export abstract class BaseAcpAgent implements Agent { currentModelId: string; options: SessionConfigSelectOption[]; }> { - const gatewayModels = await fetchGatewayModels(); + this.gatewayModels = await fetchGatewayModels(); - const options = gatewayModels + const options = this.gatewayModels .filter((model) => isAnthropicModel(model)) .map((model) => ({ value: model.id, @@ -150,4 +154,9 @@ export abstract class BaseAcpAgent implements Agent { return { currentModelId, options }; } + + getContextWindowForModel(modelId: string): number { + const match = this.gatewayModels.find((m) => m.id === modelId); + return match?.context_window ?? DEFAULT_CONTEXT_WINDOW; + } } diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index 092f4fdc9..92f9ddf1a 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -64,7 +64,6 @@ import { getAvailableSlashCommands } from "./session/commands"; import { parseMcpServers } from "./session/mcp-config"; import { DEFAULT_MODEL, - getDefaultContextWindow, getEffortOptions, resolveModelPreference, toSdkModelId, @@ -323,6 +322,16 @@ export class ClaudeAcpAgent extends BaseAcpAgent { this.session.promptRunning = true; let handedOff = false; let lastAssistantTotalUsage: number | null = null; + if (this.session.lastContextWindowSize == null) { + this.session.lastContextWindowSize = this.getContextWindowForModel( + this.session.modelId ?? "", + ); + this.logger.debug("Initial context window size from gateway", { + modelId: this.session.modelId, + contextWindowSize: this.session.lastContextWindowSize, + }); + } + let lastContextWindowSize = this.session.lastContextWindowSize; const supportsTerminalOutput = ( @@ -393,16 +402,25 @@ export class ClaudeAcpAgent extends BaseAcpAgent { this.session.accumulatedUsage.cachedWriteTokens += message.usage.cache_creation_input_tokens; - // Calculate context window size from modelUsage (minimum across all models used) + // SDK can underreport context window (e.g. 200k for 1M models). + // Use SDK value only if it's larger than what gateway reported. const contextWindows = Object.values(message.modelUsage).map( (m) => m.contextWindow, ); - const contextWindowSize = - contextWindows.length > 0 - ? Math.min(...contextWindows) - : getDefaultContextWindow(this.session.modelId ?? ""); + if (contextWindows.length > 0) { + const sdkContextWindow = Math.min(...contextWindows); + if (sdkContextWindow > lastContextWindowSize) { + lastContextWindowSize = sdkContextWindow; + } + } + this.session.lastContextWindowSize = lastContextWindowSize; + this.logger.debug("Context window size from result", { + sdkReported: contextWindows, + resolved: lastContextWindowSize, + modelId: this.session.modelId, + }); - this.session.contextSize = contextWindowSize; + this.session.contextSize = lastContextWindowSize; if (lastAssistantTotalUsage !== null) { this.session.contextUsed = lastAssistantTotalUsage; } @@ -414,7 +432,7 @@ export class ClaudeAcpAgent extends BaseAcpAgent { update: { sessionUpdate: "usage_update", used: lastAssistantTotalUsage, - size: contextWindowSize, + size: lastContextWindowSize, cost: { amount: message.total_cost_usd, currency: "USD", @@ -521,9 +539,18 @@ export class ClaudeAcpAgent extends BaseAcpAgent { }; lastAssistantTotalUsage = usage.input_tokens + - usage.output_tokens + usage.cache_read_input_tokens + usage.cache_creation_input_tokens; + + await this.client.sessionUpdate({ + sessionId: params.sessionId, + update: { + sessionUpdate: "usage_update", + used: lastAssistantTotalUsage, + size: lastContextWindowSize, + cost: null, + }, + }); } const result = await handleUserAssistantMessage(message, context); @@ -595,9 +622,11 @@ export class ClaudeAcpAgent extends BaseAcpAgent { async unstable_setSessionModel( params: SetSessionModelRequest, ): Promise { - const sdkModelId = toSdkModelId(params.modelId); - await this.session.query.setModel(sdkModelId); + await this.session.query.setModel(toSdkModelId(params.modelId)); this.session.modelId = params.modelId; + this.session.lastContextWindowSize = this.getContextWindowForModel( + params.modelId, + ); this.rebuildEffortConfigOption(params.modelId); await this.updateConfigOption("model", params.modelId); return {}; @@ -674,6 +703,8 @@ export class ClaudeAcpAgent extends BaseAcpAgent { const sdkModelId = toSdkModelId(resolvedValue); await this.session.query.setModel(sdkModelId); this.session.modelId = resolvedValue; + this.session.lastContextWindowSize = + this.getContextWindowForModel(resolvedValue); this.rebuildEffortConfigOption(resolvedValue); } else if (params.configId === "effort") { const newEffort = resolvedValue as EffortLevel; @@ -893,12 +924,12 @@ export class ClaudeAcpAgent extends BaseAcpAgent { const modelOptions = await this.getModelConfigOptions(); const resolvedModelId = settingsModel || modelOptions.currentModelId; session.modelId = resolvedModelId; + session.lastContextWindowSize = + this.getContextWindowForModel(resolvedModelId); - if (!isResume) { - const resolvedSdkModel = toSdkModelId(resolvedModelId); - if (resolvedSdkModel !== DEFAULT_MODEL) { - await this.session.query.setModel(resolvedSdkModel); - } + const resolvedSdkModel = toSdkModelId(resolvedModelId); + if (!isResume && resolvedSdkModel !== DEFAULT_MODEL) { + await this.session.query.setModel(resolvedSdkModel); } const availableModes = getAvailableModes(); diff --git a/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts b/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts index 4d0f65d00..938e7704c 100644 --- a/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts +++ b/packages/agent/src/adapters/claude/conversion/sdk-to-acp.ts @@ -544,7 +544,7 @@ export async function handleSystemMessage( message: Extract, context: MessageHandlerContext, ): Promise { - const { sessionId, client, logger } = context; + const { session, sessionId, client, logger } = context; switch (message.subtype) { case "init": @@ -554,6 +554,7 @@ export async function handleSystemMessage( sessionId, trigger: message.compact_metadata.trigger, preTokens: message.compact_metadata.pre_tokens, + contextSize: session.contextSize, }); break; case "hook_response": diff --git a/packages/agent/src/adapters/claude/session/models.ts b/packages/agent/src/adapters/claude/session/models.ts index 034cf46c1..b683bf88e 100644 --- a/packages/agent/src/adapters/claude/session/models.ts +++ b/packages/agent/src/adapters/claude/session/models.ts @@ -21,10 +21,6 @@ export function supports1MContext(modelId: string): boolean { return MODELS_WITH_1M_CONTEXT.has(modelId); } -export function getDefaultContextWindow(modelId: string): number { - return supports1MContext(modelId) ? 1_000_000 : 200_000; -} - const MODELS_WITH_EFFORT = new Set([ "claude-opus-4-5", "claude-opus-4-6", diff --git a/packages/agent/src/adapters/claude/types.ts b/packages/agent/src/adapters/claude/types.ts index 453c1a44f..6dfc6277c 100644 --- a/packages/agent/src/adapters/claude/types.ts +++ b/packages/agent/src/adapters/claude/types.ts @@ -57,6 +57,8 @@ export type Session = BaseSession & { contextUsed?: number; /** Context window size in tokens */ contextSize?: number; + /** Persists across prompt() calls so SDK-reported values survive turn boundaries */ + lastContextWindowSize?: number; promptRunning: boolean; pendingMessages: Map; nextPendingOrder: number;