From 48decd15d7b22a90ea11292980d858ca9fcd77fc Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 17 Jun 2026 11:23:26 +0200 Subject: [PATCH] fix(coding-agent): fixed context usage tracking to keep status and selector totals in sync - Added context snapshot metadata to AssistantMessage for prompt and non-message token history. - Anchored context usage calculations on assistant snapshots and computed percent numerically. - Updated status-line, /context, selector, and interactive mode flows to share session usage totals. - Extended status-line cache fingerprinting and invalidation for assistant usage and prompt/tool/skill changes. --- packages/ai/CHANGELOG.md | 5 +- packages/ai/src/types.ts | 7 + .../test/openai-max-output-tokens-cap.test.ts | 6 +- packages/coding-agent/CHANGELOG.md | 12 +- packages/coding-agent/src/cli/bench-cli.ts | 7 +- packages/coding-agent/src/collab/host.ts | 2 +- .../src/extensibility/extensions/types.ts | 8 +- .../src/modes/components/footer.ts | 7 +- .../modes/components/status-line/component.ts | 82 +++- .../status-line/context-thresholds.ts | 1 - .../src/modes/controllers/event-controller.ts | 9 +- .../modes/controllers/selector-controller.ts | 3 +- .../src/modes/interactive-mode.ts | 4 +- .../coding-agent/src/modes/rpc/rpc-types.ts | 2 +- .../src/modes/utils/context-usage.ts | 47 +- .../coding-agent/src/session/agent-session.ts | 298 +++++++----- .../test/context-consolidation.test.ts | 445 ++++++++++++++++++ .../test/interactive-mode-plan-review.test.ts | 18 +- .../test/status-line-context-cache.test.ts | 16 +- 19 files changed, 777 insertions(+), 202 deletions(-) create mode 100644 packages/coding-agent/test/context-consolidation.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 7f64a133c..60bcc8dbe 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -1,6 +1,9 @@ # Changelog ## [Unreleased] +### Added + +- Added optional `contextSnapshot` property to `AssistantMessage` with token usage metadata via new `ContextSnapshot` interface (`promptTokens`, `nonMessageTokens`, and optional `lastMessageTimestamp`) ### Added @@ -3809,4 +3812,4 @@ _Dedicated to Peter's shoulder ([@steipete](https://twitter.com/steipete))_ ## [0.9.4] - 2025-11-26 -Initial release with multi-provider LLM support. +Initial release with multi-provider LLM support. \ No newline at end of file diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index 9bd48a5ef..b33afe0b2 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -476,12 +476,19 @@ export interface DeveloperMessage { timestamp: number; // Unix timestamp in milliseconds } +export interface ContextSnapshot { + promptTokens: number; // authoritative provider prompt/input tokens + nonMessageTokens: number; // estimated non-message total at send time + lastMessageTimestamp?: number; +} + export interface AssistantMessage { role: "assistant"; content: (TextContent | ThinkingContent | RedactedThinkingContent | ToolCall)[]; api: Api; provider: Provider; model: string; + contextSnapshot?: ContextSnapshot; responseId?: string; // Provider-specific response/message identifier when the upstream API exposes one /** * Name of the upstream provider an aggregator routed this request to, as diff --git a/packages/ai/test/openai-max-output-tokens-cap.test.ts b/packages/ai/test/openai-max-output-tokens-cap.test.ts index e50f577b3..603e0b418 100644 --- a/packages/ai/test/openai-max-output-tokens-cap.test.ts +++ b/packages/ai/test/openai-max-output-tokens-cap.test.ts @@ -84,7 +84,11 @@ async function captureCompletionsBody( return completionsSse(); }; - const result = await streamSimple(model, ctx, { apiKey: "k", ...(maxTokens === undefined ? {} : { maxTokens }), fetch: fetchMock }).result(); + const result = await streamSimple(model, ctx, { + apiKey: "k", + ...(maxTokens === undefined ? {} : { maxTokens }), + fetch: fetchMock, + }).result(); expect(result.stopReason).toBe("stop"); if (!payload) throw new Error("Expected OpenAI completions request payload"); return payload; diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c756f8fe5..dfba9b647 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,17 +1,23 @@ # Changelog ## [Unreleased] + ### Added - Added `images.describeForTextModels` option (default `true`) to control automatic image description for attachments sent to models without vision input - Added automatic vision fallback prompts to describe images for text-only models +### Changed + +- Changed context usage reporting to always return numeric token counts and percentages, so status-line and footer now show estimated values instead of `?` immediately after compaction +- Changed context usage reporting to use anchored snapshots and pending-prompts estimates, which now keeps `/context`, status line, and model selector token counts in sync + ### Fixed +- Fixed context token accounting to keep branch-local anchors during branching so sibling-branch messages no longer pollute context estimates +- Fixed context usage consistency so `/context`, status line, and idle compaction logic now report the same used-token totals +- Fixed status-line context cache invalidation when assistant reasoning signature data grows so displayed context usage updates accurately - Fixed image attachment handling for text-only models by saving attachments to `local://` and injecting generated descriptions so they are no longer lost when the target model cannot process images - -### Fixed - - Fixed the ssh tool rejecting valid Windows identity files before invoking OpenSSH by skipping Unix mode-bit key validation on native Windows ([#2850](https://github.com/can1357/oh-my-pi/issues/2850)). ## [16.0.4] - 2026-06-17 diff --git a/packages/coding-agent/src/cli/bench-cli.ts b/packages/coding-agent/src/cli/bench-cli.ts index e4a891fd4..f8833ad05 100644 --- a/packages/coding-agent/src/cli/bench-cli.ts +++ b/packages/coding-agent/src/cli/bench-cli.ts @@ -17,7 +17,12 @@ import { formatDuration, getProjectDir } from "@oh-my-pi/pi-utils"; import chalk from "chalk"; import type { ApiKeyResolverModel } from "../config/api-key-resolver"; import { type CanonicalModelQueryOptions, ModelRegistry } from "../config/model-registry"; -import { formatModelSelectorValue, formatModelString, getModelMatchPreferences, resolveCliModel } from "../config/model-resolver"; +import { + formatModelSelectorValue, + formatModelString, + getModelMatchPreferences, + resolveCliModel, +} from "../config/model-resolver"; import { Settings } from "../config/settings"; import benchPrompt from "../prompts/bench.md" with { type: "text" }; import { discoverAuthStorage } from "../sdk"; diff --git a/packages/coding-agent/src/collab/host.ts b/packages/coding-agent/src/collab/host.ts index 7f9ead967..0c52c2521 100644 --- a/packages/coding-agent/src/collab/host.ts +++ b/packages/coding-agent/src/collab/host.ts @@ -427,7 +427,7 @@ export class CollabHost { contextUsage: { tokens, contextWindow: breakdown.contextWindow, - percent: tokens !== null && breakdown.contextWindow > 0 ? (tokens / breakdown.contextWindow) * 100 : null, + percent: breakdown.contextWindow > 0 ? (tokens / breakdown.contextWindow) * 100 : 0, }, participants: this.participants, }; diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 9ab810936..b973a78bd 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -274,11 +274,11 @@ export interface ExtensionUIContext { // ============================================================================ export interface ContextUsage { - /** Estimated context tokens, or null if unknown (e.g. right after compaction, before next LLM response). */ - tokens: number | null; + /** Estimated context tokens. */ + tokens: number; contextWindow: number; - /** Context usage as percentage of context window, or null if tokens is unknown. */ - percent: number | null; + /** Context usage as percentage of context window. */ + percent: number; } export interface CompactOptions { diff --git a/packages/coding-agent/src/modes/components/footer.ts b/packages/coding-agent/src/modes/components/footer.ts index 712dbe804..57b398edd 100644 --- a/packages/coding-agent/src/modes/components/footer.ts +++ b/packages/coding-agent/src/modes/components/footer.ts @@ -186,11 +186,8 @@ export class FooterComponent implements Component { // Colorize context percentage based on usage let contextPercentStr: string; const autoIndicator = this.#autoCompactEnabled ? " (auto)" : ""; - const contextPercentDisplay = `${formatContextUsage( - contextUsage?.percent === null ? null : contextPercentValue, - contextWindow, - )}${autoIndicator}`; - if (contextUsage?.percent !== null && contextUsage?.percent !== undefined) { + const contextPercentDisplay = `${formatContextUsage(contextPercentValue, contextWindow)}${autoIndicator}`; + if (contextUsage) { const color = getContextUsageThemeColor(getContextUsageLevel(contextPercentValue, contextWindow)); contextPercentStr = color === "statusLineContext" ? contextPercentDisplay : theme.fg(color, contextPercentDisplay); diff --git a/packages/coding-agent/src/modes/components/status-line/component.ts b/packages/coding-agent/src/modes/components/status-line/component.ts index 8666962a8..75670a2b1 100644 --- a/packages/coding-agent/src/modes/components/status-line/component.ts +++ b/packages/coding-agent/src/modes/components/status-line/component.ts @@ -1,6 +1,7 @@ import * as fs from "node:fs"; import * as path from "node:path"; import type { AgentMessage } from "@oh-my-pi/pi-agent-core"; +import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { type Component, truncateToWidth, visibleWidth } from "@oh-my-pi/pi-tui"; import { getProjectDir } from "@oh-my-pi/pi-utils"; import { $ } from "bun"; @@ -55,21 +56,70 @@ function messageFingerprint(msg: AgentMessage): string { } } } else if (role === "assistant") { + const assistantMsg = msg as AssistantMessage; + const usageExt = assistantMsg.usage as unknown as { promptTokensDetails?: unknown }; + const usageTotal = assistantMsg.usage?.totalTokens ?? 0; + const promptBuckets = usageExt?.promptTokensDetails ? 1 : 0; + const stopReason = assistantMsg.stopReason ?? ""; + + let signatureLen = 0; + let redactedLen = 0; + const msgExt = assistantMsg as unknown as { + thinkingSignature?: string; + textSignature?: string; + thoughtSignature?: string; + redactedThinking?: { data?: string }; + }; + const thinkingSignature = msgExt.thinkingSignature; + if (typeof thinkingSignature === "string") { + signatureLen += thinkingSignature.length; + } + const textSignature = msgExt.textSignature; + if (typeof textSignature === "string") { + signatureLen += textSignature.length; + } + const thoughtSignature = msgExt.thoughtSignature; + if (typeof thoughtSignature === "string") { + signatureLen += thoughtSignature.length; + } + const redactedData = msgExt.redactedThinking?.data; + if (typeof redactedData === "string") { + redactedLen += redactedData.length; + } + const content = (msg as { content?: unknown }).content; if (Array.isArray(content)) { blocks = content.length; for (const block of content) { if (!block || typeof block !== "object") continue; - const b = block as { type?: string; text?: string; thinking?: string; name?: string; arguments?: unknown }; + const b = block as { + type?: string; + text?: string; + thinking?: string; + thinkingSignature?: string; + signature?: string; + textSignature?: string; + thoughtSignature?: string; + data?: string; + name?: string; + arguments?: unknown; + }; if (b.type === "text" && typeof b.text === "string") textLen += b.text.length; - else if (b.type === "thinking" && typeof b.thinking === "string") textLen += b.thinking.length; - else if (b.type === "toolCall") { + else if (b.type === "thinking") { + if (typeof b.thinking === "string") textLen += b.thinking.length; + if (typeof b.thinkingSignature === "string") signatureLen += b.thinkingSignature.length; + if (typeof b.signature === "string") signatureLen += b.signature.length; + if (typeof b.textSignature === "string") signatureLen += b.textSignature.length; + if (typeof b.thoughtSignature === "string") signatureLen += b.thoughtSignature.length; + } else if (b.type === "redactedThinking" && typeof b.data === "string") { + redactedLen += b.data.length; + } else if (b.type === "toolCall") { if (typeof b.name === "string") textLen += b.name.length; - // Argument bytes vary; a length proxy is enough to detect in-place edits. textLen += b.arguments === undefined ? 0 : JSON.stringify(b.arguments).length; } } } + return `${role}:${ts}:${textLen}:${blocks}:${images}:${signatureLen}:${redactedLen}:${usageTotal}:${promptBuckets}:${stopReason}`; } else if (role === "toolResult" || role === "hookMessage") { const content = (msg as { content?: unknown }).content; if (typeof content === "string") { @@ -95,8 +145,11 @@ interface ContextUsageMemo { length: number; lastFingerprint: string | undefined; modelContextWindow: number; - usedTokens: number | null; + usedTokens: number; contextWindow: number; + systemPromptRef: readonly string[] | undefined; + toolsRef: readonly any[] | undefined; + skillsRef: readonly any[] | undefined; } const EMPTY_MESSAGES: readonly AgentMessage[] = []; @@ -545,25 +598,32 @@ export class StatusLineComponent implements Component { * (right after compaction, before the next response). Exposed (non-private) * for unit tests and the collab host's state broadcast. */ - getCachedContextBreakdown(): { usedTokens: number | null; contextWindow: number } { + getCachedContextBreakdown(): { usedTokens: number; contextWindow: number } { const messages = this.session.messages ?? EMPTY_MESSAGES; const modelContextWindow = this.session.model?.contextWindow ?? 0; const length = messages.length; const lastFingerprint = length > 0 ? messageFingerprint(messages[length - 1]!) : undefined; + const systemPrompt = this.session.systemPrompt; + const tools = this.session.agent?.state?.tools; + const skills = this.session.skills; + const cache = this.#contextUsageCache; if ( cache && cache.messagesRef === messages && cache.length === length && cache.lastFingerprint === lastFingerprint && - cache.modelContextWindow === modelContextWindow + cache.modelContextWindow === modelContextWindow && + cache.systemPromptRef === systemPrompt && + cache.toolsRef === tools && + cache.skillsRef === skills ) { return { usedTokens: cache.usedTokens, contextWindow: cache.contextWindow }; } const usage = this.session.getContextUsage(); - const usedTokens = usage?.tokens ?? null; + const usedTokens = usage?.tokens ?? 0; const contextWindow = usage?.contextWindow ?? modelContextWindow; this.#contextUsageCache = { messagesRef: messages, @@ -572,6 +632,9 @@ export class StatusLineComponent implements Component { modelContextWindow, usedTokens, contextWindow, + systemPromptRef: systemPrompt, + toolsRef: tools, + skillsRef: skills, }; return { usedTokens, contextWindow }; } @@ -607,8 +670,7 @@ export class StatusLineComponent implements Component { if (includeContext) { const breakdown = this.getCachedContextBreakdown(); contextWindow = breakdown.contextWindow || contextWindow; - contextPercent = - breakdown.usedTokens === null ? null : contextWindow > 0 ? (breakdown.usedTokens / contextWindow) * 100 : 0; + contextPercent = contextWindow > 0 ? (breakdown.usedTokens / contextWindow) * 100 : 0; } // Collab guest: context comes from the host's state frames — the local diff --git a/packages/coding-agent/src/modes/components/status-line/context-thresholds.ts b/packages/coding-agent/src/modes/components/status-line/context-thresholds.ts index bda33a23e..87436a6fa 100644 --- a/packages/coding-agent/src/modes/components/status-line/context-thresholds.ts +++ b/packages/coding-agent/src/modes/components/status-line/context-thresholds.ts @@ -58,7 +58,6 @@ export function getContextUsageLevel(contextPercent: number, contextWindow: numb /** * Format context usage as `%/` (e.g. `5.1%/1M`), matching the * status line's context gauge so subagent and footer renderers stay in sync. - * A `null`/`undefined` percent (unknown, e.g. right after compaction) renders as `?`. */ export function formatContextUsage(contextPercent: number | null | undefined, contextWindow: number): string { const pct = contextPercent === null || contextPercent === undefined ? "?" : `${contextPercent.toFixed(1)}%`; diff --git a/packages/coding-agent/src/modes/controllers/event-controller.ts b/packages/coding-agent/src/modes/controllers/event-controller.ts index 5cd20d136..e11289a7d 100644 --- a/packages/coding-agent/src/modes/controllers/event-controller.ts +++ b/packages/coding-agent/src/modes/controllers/event-controller.ts @@ -1,6 +1,5 @@ import { INTENT_FIELD } from "@oh-my-pi/pi-agent-core"; -import { calculatePromptTokens } from "@oh-my-pi/pi-agent-core/compaction/compaction"; -import type { AssistantMessage, ImageContent } from "@oh-my-pi/pi-ai"; +import type { ImageContent } from "@oh-my-pi/pi-ai"; import { type Component, Loader, TERMINAL } from "@oh-my-pi/pi-tui"; import { extractTextContent } from "../../commit/utils"; import { settings } from "../../config/settings"; @@ -1107,11 +1106,7 @@ export class EventController { } #currentContextTokens(): number { - const lastAssistant = this.ctx.viewSession.agent.state.messages - .slice() - .reverse() - .find((m): m is AssistantMessage => m.role === "assistant" && m.stopReason !== "aborted"); - return lastAssistant?.usage ? calculatePromptTokens(lastAssistant.usage) : 0; + return this.ctx.viewSession.getContextUsage()?.tokens ?? 0; } sendCompletionNotification(): void { diff --git a/packages/coding-agent/src/modes/controllers/selector-controller.ts b/packages/coding-agent/src/modes/controllers/selector-controller.ts index 2e666561d..ba8e9acca 100644 --- a/packages/coding-agent/src/modes/controllers/selector-controller.ts +++ b/packages/coding-agent/src/modes/controllers/selector-controller.ts @@ -67,7 +67,6 @@ import { TranscriptBlock } from "../components/transcript-container"; import { TreeSelectorComponent } from "../components/tree-selector"; import { UserMessageSelectorComponent } from "../components/user-message-selector"; import type { SessionObserverRegistry } from "../session-observer-registry"; -import { computeContextBreakdown } from "../utils/context-usage"; import { buildCopyTargets } from "../utils/copy-targets"; const MANUAL_LOGIN_TIP = "Tip: You can complete pairing with /login ."; @@ -444,7 +443,7 @@ export class SelectorController { } showModelSelector(options?: { temporaryOnly?: boolean }): void { - const currentContextTokens = computeContextBreakdown(this.ctx.session).usedTokens; + const currentContextTokens = this.ctx.session.getContextUsage()?.tokens ?? 0; this.showSelector(done => { const selector = new ModelSelectorComponent( this.ctx.ui, diff --git a/packages/coding-agent/src/modes/interactive-mode.ts b/packages/coding-agent/src/modes/interactive-mode.ts index df48cc4da..04d22bc61 100644 --- a/packages/coding-agent/src/modes/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive-mode.ts @@ -2179,7 +2179,7 @@ export class InteractiveMode implements InteractiveModeContext { } #formatKeepContextLabel(contextUsage: ContextUsage | undefined): string { - if (contextUsage?.tokens == null) { + if (!contextUsage) { return "Approve and keep context"; } const tokens = formatContextTokenCount(contextUsage.tokens); @@ -2188,7 +2188,7 @@ export class InteractiveMode implements InteractiveModeContext { } #isKeepContextDisabled(contextUsage: ContextUsage | undefined): boolean { - return contextUsage?.percent != null && contextUsage.percent > PLAN_KEEP_CONTEXT_DISABLE_THRESHOLD_PERCENT; + return contextUsage !== undefined && contextUsage.percent > PLAN_KEEP_CONTEXT_DISABLE_THRESHOLD_PERCENT; } async #openPlanInExternalEditor(planFilePath: string): Promise { diff --git a/packages/coding-agent/src/modes/rpc/rpc-types.ts b/packages/coding-agent/src/modes/rpc/rpc-types.ts index e5906d787..10863bc79 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-types.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-types.ts @@ -108,7 +108,7 @@ export interface RpcSessionState { /** For session dump / export (plain-text parity with /dump). */ systemPrompt?: string[]; dumpTools?: Array<{ name: string; description: string; parameters: unknown; examples?: readonly ToolExample[] }>; - /** Current context window usage. Null tokens/percent when unknown (e.g. right after compaction). */ + /** Current context window usage. */ contextUsage?: ContextUsage; } diff --git a/packages/coding-agent/src/modes/utils/context-usage.ts b/packages/coding-agent/src/modes/utils/context-usage.ts index 4953deb41..5a11ffc58 100644 --- a/packages/coding-agent/src/modes/utils/context-usage.ts +++ b/packages/coding-agent/src/modes/utils/context-usage.ts @@ -94,7 +94,7 @@ export function computeNonMessageTokens(session: AgentSession): number { * the status-line fast path intentionally uses the equivalent collapsed total * in `computeNonMessageTokens`. */ -function computeNonMessageBreakdown(session: AgentSession): { +export function computeNonMessageBreakdown(session: AgentSession): { skillsTokens: number; toolsTokens: number; systemContextTokens: number; @@ -119,21 +119,36 @@ export function computeContextBreakdown( const model = session.model; const contextWindow = model?.contextWindow ?? 0; - let messagesTokens = 0; - const convo = session.messages; - if (convo) { - for (const message of convo) { - messagesTokens += estimateTokens(message); - } - } + const breakdown = typeof session.getContextBreakdown === "function" ? session.getContextBreakdown() : undefined; - // The rendered system prompt already contains the skill descriptions and the - // markdown tool descriptions. To present a non-overlapping breakdown: - // System prompt = total system prompt text - skills section (tool descriptions stay) - // Tools = JSON tool schema sent separately on the wire - // Skills = the skill list embedded in the system prompt - // Messages = conversation messages - const { skillsTokens, toolsTokens, systemContextTokens, systemPromptTokens } = computeNonMessageBreakdown(session); + let messagesTokens = 0; + let skillsTokens = 0; + let toolsTokens = 0; + let systemContextTokens = 0; + let systemPromptTokens = 0; + let usedTokens = 0; + + if (breakdown) { + messagesTokens = breakdown.messagesTokens; + skillsTokens = breakdown.skillsTokens; + toolsTokens = breakdown.systemToolsTokens; + systemContextTokens = breakdown.systemContextTokens; + systemPromptTokens = breakdown.systemPromptTokens; + usedTokens = breakdown.usedTokens; + } else { + const convo = session.messages; + if (convo) { + for (const message of convo) { + messagesTokens += estimateTokens(message); + } + } + const nonMessage = computeNonMessageBreakdown(session); + skillsTokens = nonMessage.skillsTokens; + toolsTokens = nonMessage.toolsTokens; + systemContextTokens = nonMessage.systemContextTokens; + systemPromptTokens = nonMessage.systemPromptTokens; + usedTokens = skillsTokens + toolsTokens + systemContextTokens + systemPromptTokens + messagesTokens; + } const categories: CategoryInfo[] = [ { id: "systemPrompt", label: "System prompt", tokens: systemPromptTokens, color: "accent", glyph: CELL_FILLED }, @@ -155,8 +170,6 @@ export function computeContextBreakdown( }, ]; - const usedTokens = categories.reduce((sum, c) => sum + c.tokens, 0); - let autoCompactBufferTokens = 0; if (contextWindow > 0) { const compactionSettings = session.settings.getGroup("compaction") as CompactionSettings; diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 32981a76b..c0a7ab907 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -202,7 +202,7 @@ import { containsOrchestrate, ORCHESTRATE_NOTICE } from "../modes/orchestrate"; import { getCurrentThemeName, theme } from "../modes/theme/theme"; import { parseTurnBudget } from "../modes/turn-budget"; import { containsUltrathink, ULTRATHINK_NOTICE } from "../modes/ultrathink"; -import { computeNonMessageTokens } from "../modes/utils/context-usage"; +import { computeNonMessageBreakdown, computeNonMessageTokens } from "../modes/utils/context-usage"; import { containsWorkflow, WORKFLOW_NOTICE } from "../modes/workflow"; import { createPlanReadMatcher } from "../plan-mode/plan-protection"; import type { PlanModeState } from "../plan-mode/state"; @@ -556,6 +556,17 @@ export interface RoleModelCycle { currentIndex: number; } +export interface ContextUsageBreakdown { + contextWindow: number; + anchored: boolean; + usedTokens: number; + systemPromptTokens: number; + systemToolsTokens: number; + systemContextTokens: number; + skillsTokens: number; + messagesTokens: number; +} + /** Session statistics for /session command */ export interface SessionStats { sessionFile: string | undefined; @@ -1229,15 +1240,13 @@ export class AgentSession { #unexpectedStopRetryCount = 0; #promptGeneration = 0; #pendingAgentEndEmit: AgentSessionEvent | undefined; - #pendingProviderRequestNonMessageTokens: number | undefined = undefined; - #lastProviderUsageNonMessage: + #pendingContextSnapshot: | { - provider: AssistantMessage["provider"]; - model: AssistantMessage["model"]; - timestamp: AssistantMessage["timestamp"]; - tokens: number; + promptTokens: number; + nonMessageTokens: number; + cutoffCount: number; } - | undefined; + | undefined = undefined; #obfuscator: SecretObfuscator | undefined; #checkpointState: CheckpointState | undefined = undefined; #pendingRewindReport: string | undefined = undefined; @@ -2359,6 +2368,15 @@ export class AgentSession { event.message.role === "fileMention" ) { // Regular LLM message - persist as SessionMessageEntry + if (event.message.role === "assistant") { + const assistantMsg = event.message as AssistantMessage; + if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { + assistantMsg.contextSnapshot = { + promptTokens: calculatePromptTokens(assistantMsg.usage), + nonMessageTokens: this.#pendingContextSnapshot?.nonMessageTokens ?? computeNonMessageTokens(this), + }; + } + } this.sessionManager.appendMessage(event.message); } // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere @@ -2367,14 +2385,6 @@ export class AgentSession { if (event.message.role === "assistant") { this.#lastAssistantMessage = event.message; const assistantMsg = event.message as AssistantMessage; - if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { - this.#lastProviderUsageNonMessage = { - provider: assistantMsg.provider, - model: assistantMsg.model, - timestamp: assistantMsg.timestamp, - tokens: this.#pendingProviderRequestNonMessageTokens ?? computeNonMessageTokens(this), - }; - } const currentGrantsAnthropicPriority = this.serviceTier === "priority" || this.serviceTier === "claude-only"; if (assistantMsg.disabledFeatures?.includes("priority") && currentGrantsAnthropicPriority) { @@ -2417,7 +2427,6 @@ export class AgentSession { this.#retryAttempt = 0; } } - if (event.message.role === "toolResult") { const { toolName, details, isError, content } = event.message as { toolName?: string; @@ -5576,11 +5585,23 @@ export class AgentSession { } const agentPromptOptions = options?.toolChoice ? { toolChoice: options.toolChoice } : undefined; - this.#pendingProviderRequestNonMessageTokens = computeNonMessageTokens(this); + const nonMessageTokens = computeNonMessageTokens(this); + const contextWindow = this.model?.contextWindow ?? 0; + const breakdown = this.getContextBreakdown({ contextWindow, pendingMessages: messages }); + const promptTokens = + breakdown?.usedTokens ?? + nonMessageTokens + + this.messages.reduce((sum, msg) => sum + estimateTokens(msg), 0) + + messages.reduce((sum, msg) => sum + estimateTokens(msg), 0); + this.#pendingContextSnapshot = { + promptTokens, + nonMessageTokens, + cutoffCount: this.messages.length + messages.length, + }; try { await this.#promptAgentWithIdleRetry(messages, agentPromptOptions); } finally { - this.#pendingProviderRequestNonMessageTokens = undefined; + this.#pendingContextSnapshot = undefined; } if (!options?.skipPostPromptRecoveryWait) { await this.#waitForPostPromptRecovery(generation); @@ -7531,39 +7552,12 @@ export class AgentSession { } } - #estimatePendingPromptTokens(messages: AgentMessage[]): number { - let tokens = computeNonMessageTokens(this); - for (const message of this.messages) { - tokens += estimateTokens(message); - } - for (const message of messages) { - tokens += estimateTokens(message); - } - return tokens; - } - #estimatePrePromptContextTokens(messages: AgentMessage[], contextWindow: number): number { - const currentUsage = this.getContextUsage({ contextWindow }); - if (typeof currentUsage?.tokens !== "number" || !Number.isFinite(currentUsage.tokens)) { - return this.#estimatePendingPromptTokens(messages); - } - - const currentEstimate = this.#estimateContextTokens(); - if (!currentEstimate.providerAnchored) { - return this.#estimatePendingPromptTokens(messages); - } - - let tokens = currentUsage.tokens; - const previousNonMessageTokens = currentEstimate.providerNonMessageTokens; - if (previousNonMessageTokens !== undefined) { - const currentNonMessageTokens = computeNonMessageTokens(this); - const nonMessageTokenGrowth = Math.max(0, currentNonMessageTokens - previousNonMessageTokens); - tokens += nonMessageTokenGrowth; - } - for (const message of messages) { - tokens += estimateTokens(message); - } - return tokens; + const breakdown = this.getContextBreakdown({ contextWindow, pendingMessages: messages }); + return ( + breakdown?.usedTokens ?? + computeNonMessageTokens(this) + messages.reduce((sum, msg) => sum + estimateTokens(msg), 0) + ); } async #runPrePromptCompactionIfNeeded(messages: AgentMessage[]): Promise { @@ -9327,7 +9321,7 @@ export class AgentSession { const recoveryBand = Math.floor(thresholdTokens * SHAKE_RECOVERY_BAND); stillOverThreshold = correctedTokens > recoveryBand; } else { - const postShakeTokens = this.#estimatePendingPromptTokens([]); + const postShakeTokens = this.getContextUsage({ contextWindow })?.tokens ?? 0; stillOverThreshold = shouldCompact(postShakeTokens, contextWindow, compactionSettings); } } @@ -11261,47 +11255,151 @@ export class AgentSession { * Uses the last assistant message's usage data when available, * otherwise estimates tokens for all messages. */ - getContextUsage(options?: { contextWindow?: number }): ContextUsage | undefined { + getContextBreakdown(options?: { + contextWindow?: number; + pendingMessages?: AgentMessage[]; + }): ContextUsageBreakdown | undefined { const model = this.model; const contextWindow = options?.contextWindow ?? model?.contextWindow ?? 0; if (!Number.isFinite(contextWindow) || contextWindow <= 0) return undefined; - // After compaction, the last assistant usage reflects pre-compaction context size. - // We can only trust usage from an assistant that responded after the latest compaction. - // If no such assistant exists, context token count is unknown until the next LLM response. + const { skillsTokens, toolsTokens, systemContextTokens, systemPromptTokens } = computeNonMessageBreakdown(this); + const categoryNonMessageTokens = skillsTokens + toolsTokens + systemContextTokens + systemPromptTokens; + const currentNonMessageTokens = computeNonMessageTokens(this); + const branchEntries = this.sessionManager.getBranch(); const latestCompaction = getLatestCompactionEntry(branchEntries); + const compactionIndex = latestCompaction ? branchEntries.lastIndexOf(latestCompaction) : -1; - if (latestCompaction) { - // Check if there's a valid assistant usage after the compaction boundary - const compactionIndex = branchEntries.lastIndexOf(latestCompaction); - let hasPostCompactionUsage = false; + let usedTokens = 0; + let anchored = false; + + const pendingMessages = options?.pendingMessages ?? []; + + let anchorEntry: SessionMessageEntry | undefined; + let isPending = false; + + if (this.#pendingContextSnapshot) { + isPending = true; + } else { for (let i = branchEntries.length - 1; i > compactionIndex; i--) { const entry = branchEntries[i]; if (entry.type === "message" && entry.message.role === "assistant") { const assistant = entry.message; - if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error") { - const contextTokens = calculateContextTokens(assistant.usage); - if (contextTokens > 0) { - hasPostCompactionUsage = true; - } + if (assistant.stopReason !== "aborted" && assistant.stopReason !== "error" && assistant.usage) { + anchorEntry = entry; break; } } } + } - if (!hasPostCompactionUsage) { - return { tokens: null, contextWindow, percent: null }; + if (isPending && this.#pendingContextSnapshot) { + const anchor = this.#pendingContextSnapshot; + anchored = true; + + const resolvedActiveMessages = this.messages; + let tailTokens = 0; + + if (resolvedActiveMessages.length > anchor.cutoffCount) { + for (let i = anchor.cutoffCount; i < resolvedActiveMessages.length; i++) { + tailTokens += estimateTokens(resolvedActiveMessages[i]); + } + } + + usedTokens = + anchor.promptTokens + + Math.max(0, currentNonMessageTokens - anchor.nonMessageTokens) + + tailTokens + + pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0); + } else if (anchorEntry) { + const anchorAssistant = anchorEntry.message as AssistantMessage; + const promptTokens = + anchorAssistant.contextSnapshot?.promptTokens ?? calculatePromptTokens(anchorAssistant.usage); + const nonMessageTokens = anchorAssistant.contextSnapshot?.nonMessageTokens ?? computeNonMessageTokens(this); + const anchor = { promptTokens, nonMessageTokens }; + anchored = true; + + const resolvedActiveMessages = this.messages; + let resolvedAnchorIndex = resolvedActiveMessages.indexOf(anchorAssistant); + if (resolvedAnchorIndex === -1) { + resolvedAnchorIndex = resolvedActiveMessages.findIndex( + msg => msg.role === "assistant" && msg.timestamp === anchorAssistant.timestamp, + ); + } + + if (resolvedAnchorIndex !== -1) { + let tailTokens = 0; + for (let i = resolvedAnchorIndex + 1; i < resolvedActiveMessages.length; i++) { + tailTokens += estimateTokens(resolvedActiveMessages[i]); + } + usedTokens = + anchor.promptTokens + + Math.max(0, currentNonMessageTokens - anchor.nonMessageTokens) + + tailTokens + + pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0); + } else { + anchored = false; } } - const estimate = this.#estimateContextTokens(); - const percent = (estimate.tokens / contextWindow) * 100; + if (!anchored && !isPending && branchEntries.length === 0) { + // Fallback: look for the latest assistant message with usage/snapshot in this.messages (for branchless/fake sessions in tests) + const resolvedActiveMessages = this.messages; + for (let i = resolvedActiveMessages.length - 1; i >= 0; i--) { + const msg = resolvedActiveMessages[i]; + if (msg.role === "assistant" && msg.stopReason !== "aborted" && msg.stopReason !== "error" && msg.usage) { + const promptTokens = msg.contextSnapshot?.promptTokens ?? calculatePromptTokens(msg.usage); + const nonMessageTokens = msg.contextSnapshot?.nonMessageTokens ?? computeNonMessageTokens(this); + + let tailTokens = 0; + for (let j = i + 1; j < resolvedActiveMessages.length; j++) { + tailTokens += estimateTokens(resolvedActiveMessages[j]); + } + + usedTokens = + promptTokens + + Math.max(0, currentNonMessageTokens - nonMessageTokens) + + tailTokens + + pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0); + anchored = true; + break; + } + } + } + if (!anchored) { + const resolvedActiveMessages = this.messages; + let messagesTokens = 0; + for (const msg of resolvedActiveMessages) { + messagesTokens += estimateTokens(msg); + } + usedTokens = + currentNonMessageTokens + + messagesTokens + + pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0); + } + + const messagesTokens = Math.max(0, usedTokens - categoryNonMessageTokens); return { - tokens: estimate.tokens, contextWindow, - percent, + anchored, + usedTokens, + systemPromptTokens, + systemToolsTokens: toolsTokens, + systemContextTokens, + skillsTokens, + messagesTokens, + }; + } + + getContextUsage(options?: { contextWindow?: number }): ContextUsage | undefined { + const breakdown = this.getContextBreakdown(options); + if (!breakdown) return undefined; + return { + tokens: breakdown.usedTokens, + contextWindow: breakdown.contextWindow, + percent: breakdown.contextWindow > 0 ? (breakdown.usedTokens / breakdown.contextWindow) * 100 : 0, }; } @@ -11485,64 +11583,6 @@ export class AgentSession { return run; } - /** - * Estimate context tokens from messages, using the last assistant usage when available. - */ - #estimateContextTokens(): { - tokens: number; - providerAnchored: boolean; - providerNonMessageTokens?: number; - } { - const messages = this.messages; - - // Find last assistant message with valid usage. - let lastUsageIndex: number | null = null; - let lastUsage: Usage | undefined; - for (let i = messages.length - 1; i >= 0; i--) { - const msg = messages[i]; - if (msg.role === "assistant") { - const assistantMsg = msg as AssistantMessage; - if (assistantMsg.stopReason !== "aborted" && assistantMsg.stopReason !== "error" && assistantMsg.usage) { - lastUsage = assistantMsg.usage; - lastUsageIndex = i; - break; - } - } - } - - if (!lastUsage || lastUsageIndex === null) { - // No usage data - estimate all messages - let estimated = 0; - for (const message of messages) { - estimated += estimateTokens(message); - } - return { - tokens: estimated, - providerAnchored: false, - }; - } - - const usageTokens = calculatePromptTokens(lastUsage); - const providerNonMessage = - this.#lastProviderUsageNonMessage && - messages[lastUsageIndex]?.role === "assistant" && - this.#lastProviderUsageNonMessage.provider === (messages[lastUsageIndex] as AssistantMessage).provider && - this.#lastProviderUsageNonMessage.model === (messages[lastUsageIndex] as AssistantMessage).model && - this.#lastProviderUsageNonMessage.timestamp === (messages[lastUsageIndex] as AssistantMessage).timestamp - ? this.#lastProviderUsageNonMessage.tokens - : undefined; - let trailingTokens = 0; - for (let i = lastUsageIndex + 1; i < messages.length; i++) { - trailingTokens += estimateTokens(messages[i]); - } - - return { - tokens: usageTokens + trailingTokens, - providerAnchored: true, - providerNonMessageTokens: providerNonMessage, - }; - } - /** * Export session to HTML. * @param outputPath Optional output path (defaults to session directory) diff --git a/packages/coding-agent/test/context-consolidation.test.ts b/packages/coding-agent/test/context-consolidation.test.ts new file mode 100644 index 000000000..735b54880 --- /dev/null +++ b/packages/coding-agent/test/context-consolidation.test.ts @@ -0,0 +1,445 @@ +import { afterAll, beforeAll, describe, expect, it, vi } from "bun:test"; +import * as path from "node:path"; +import { Agent, type AgentMessage } from "@oh-my-pi/pi-agent-core"; +import { estimateTokens } from "@oh-my-pi/pi-agent-core/compaction/compaction"; +import type { AssistantMessage, Message, Model } from "@oh-my-pi/pi-ai"; +import { createMockModel } from "@oh-my-pi/pi-ai/providers/mock"; +import { ModelRegistry } from "@oh-my-pi/pi-coding-agent/config/model-registry"; +import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import { StatusLineComponent } from "@oh-my-pi/pi-coding-agent/modes/components/status-line"; +import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; +import { computeContextBreakdown } from "@oh-my-pi/pi-coding-agent/modes/utils/context-usage"; +import { AgentSession } from "@oh-my-pi/pi-coding-agent/session/agent-session"; +import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; +import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager"; +import { TempDir } from "@oh-my-pi/pi-utils"; + +describe("Context usage consolidation", () => { + let sharedDir: TempDir; + let authStorage: AuthStorage; + let modelRegistry: ModelRegistry; + let mockModel: Model; + + beforeAll(async () => { + sharedDir = TempDir.createSync("@pi-context-shared-"); + authStorage = await AuthStorage.create(path.join(sharedDir.path(), "testauth.db")); + authStorage.setRuntimeApiKey("anthropic", "test-key"); + modelRegistry = new ModelRegistry(authStorage); + await Settings.init({ inMemory: true }); + await initTheme(); + + mockModel = createMockModel({ + id: "gpt-mock", + provider: "openai", + contextWindow: 100_000, + responses: [ + { + content: ["response text"], + stopReason: "stop", + usage: { + input: 100, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 120, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }, + ], + }); + }); + + afterAll(async () => { + authStorage.close(); + try { + await sharedDir.remove(); + } catch {} + }); + + function createSession( + tempDir: TempDir, + messages: AgentMessage[] = [], + ): { session: AgentSession; sessionManager: SessionManager; agent: Agent } { + const sessionManager = SessionManager.create(tempDir.path(), tempDir.path()); + for (const msg of messages) { + sessionManager.appendMessage(msg as unknown as Parameters[0]); + } + + const agent = new Agent({ + getApiKey: () => "test-key", + initialState: { + model: mockModel, + systemPrompt: ["You are a helpful assistant."], + tools: [], + messages, + }, + streamFn: (mockModel as unknown as { stream: unknown }).stream as unknown as NonNullable< + ConstructorParameters[0] + >["streamFn"], + }); + + const session = new AgentSession({ + agent, + sessionManager, + settings: Settings.isolated({ + "compaction.enabled": true, + "compaction.autoContinue": false, + "compaction.strategy": "context-full", + "compaction.thresholdTokens": 8000, + }), + modelRegistry, + }); + + const sessionContext = session.buildDisplaySessionContext(); + agent.replaceMessages(sessionContext.messages); + + return { session, sessionManager, agent }; + } + + function syncSession(session: AgentSession, agent: Agent): void { + const sessionContext = session.buildDisplaySessionContext(); + agent.replaceMessages(sessionContext.messages); + } + + it("keeps branch-local anchors safe from sibling branches", async () => { + const tempDir = TempDir.createSync("@branch-local-"); + const { session, sessionManager, agent } = createSession(tempDir); + + sessionManager.appendMessage({ role: "user", content: "hello", timestamp: 1000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "parent response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 100, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 120, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 100, nonMessageTokens: 10 }, + timestamp: 2000, + } as AssistantMessage); + + syncSession(session, agent); + const parentId = sessionManager.getBranch().slice(-1)[0]?.id; + + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "sibling response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 500, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 520, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 500, nonMessageTokens: 10 }, + timestamp: 3000, + } as AssistantMessage); + + sessionManager.branch(parentId!); + sessionManager.appendMessage({ role: "user", content: "active branch hello", timestamp: 4000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "active response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 200, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 220, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 200, nonMessageTokens: 10 }, + timestamp: 5000, + } as AssistantMessage); + + syncSession(session, agent); + + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(true); + expect(breakdown?.usedTokens).toBe(200); + + await tempDir.remove(); + }); + + it("recovers correct anchor after rollback", async () => { + const tempDir = TempDir.createSync("@rollback-"); + const { session, sessionManager, agent } = createSession(tempDir); + + sessionManager.appendMessage({ role: "user", content: "first", timestamp: 1000 } as Message); + const firstId = sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "first response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 150, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 170, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 150, nonMessageTokens: 10 }, + timestamp: 2000, + } as AssistantMessage); + + sessionManager.appendMessage({ role: "user", content: "second", timestamp: 3000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "second response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 300, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 320, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 300, nonMessageTokens: 10 }, + timestamp: 4000, + } as AssistantMessage); + + sessionManager.branch(firstId); + syncSession(session, agent); + + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(true); + expect(breakdown?.usedTokens).toBe(150); + + await tempDir.remove(); + }); + + it("uses speculative mode after compaction with no subsequent assistant snapshot", async () => { + const tempDir = TempDir.createSync("@compaction-speculative-"); + const { session, sessionManager, agent } = createSession(tempDir); + + const firstId = sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "pre compaction response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 150, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 170, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 150, nonMessageTokens: 10 }, + timestamp: 1000, + } as AssistantMessage); + + sessionManager.appendCompaction("compact summary", "compact summary", firstId, 50); + syncSession(session, agent); + + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(false); + + await tempDir.remove(); + }); + + it("includes custom message / summary in resolved active tail", async () => { + const tempDir = TempDir.createSync("@tail-custom-"); + const { session, sessionManager, agent } = createSession(tempDir); + + sessionManager.appendMessage({ role: "user", content: "query", timestamp: 1000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "anchor response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 150, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 170, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 150, nonMessageTokens: 10 }, + timestamp: 2000, + } as AssistantMessage); + + const customMsg: AgentMessage = { + role: "custom", + customType: "test-custom", + content: "custom content block", + display: true, + timestamp: 3000, + }; + sessionManager.appendMessage(customMsg as unknown as Parameters[0]); + syncSession(session, agent); + + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(true); + + const customEstimate = estimateTokens(customMsg); + expect(breakdown?.usedTokens).toBe(150 + customEstimate); + + await tempDir.remove(); + }); + + it("ensures /context, status line, model selector, and idle compaction all agree on used tokens", async () => { + const tempDir = TempDir.createSync("@agree-"); + const { session, sessionManager, agent } = createSession(tempDir); + + sessionManager.appendMessage({ role: "user", content: "query", timestamp: 1000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "anchor response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 250, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 270, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 250, nonMessageTokens: 10 }, + timestamp: 2000, + } as AssistantMessage); + + syncSession(session, agent); + + const breakdownVal = session.getContextBreakdown(); + const used = breakdownVal?.usedTokens; + + const cb = computeContextBreakdown(session); + expect(cb.usedTokens).toBe(used!); + + const sl = new StatusLineComponent(session); + expect(sl.getCachedContextBreakdown().usedTokens).toBe(used!); + + const cu = session.getContextUsage(); + expect(cu?.tokens).toBe(used!); + + await tempDir.remove(); + }); + + it("invalidates status-line cache on reasoning-signature growth", async () => { + const tempDir = TempDir.createSync("@cache-invalidate-"); + const { session, sessionManager, agent } = createSession(tempDir); + + sessionManager.appendMessage({ role: "user", content: "query", timestamp: 1000 } as Message); + const assistant: AssistantMessage = { + role: "assistant", + content: [{ type: "text", text: "text content" }], + usage: { + input: 250, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 270, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 250, nonMessageTokens: 10 }, + timestamp: 2000, + stopReason: "stop", + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + }; + sessionManager.appendMessage(assistant); + syncSession(session, agent); + + const sl = new StatusLineComponent(session); + const initialBreakdown = sl.getCachedContextBreakdown(); + + const assistantExt = assistant as unknown as { thinkingSignature: string }; + assistantExt.thinkingSignature = "signature_grows"; + + const nextBreakdown = sl.getCachedContextBreakdown(); + expect(nextBreakdown.usedTokens).toBe(initialBreakdown.usedTokens); + + await tempDir.remove(); + }); + + it("uses live in-flight pending snapshot when request is active", async () => { + const tempDir = TempDir.createSync("@inflight-"); + const { session, agent } = createSession(tempDir); + + const { promise, resolve } = Promise.withResolvers(); + const promptSpy = vi.spyOn(agent, "prompt").mockImplementation(async () => { + await promise; + }); + + const promptPromise = session.prompt("query"); + + // Wait for the prompt request to become active + await Bun.sleep(10); + + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(true); + expect(breakdown?.usedTokens).toBeGreaterThan(0); + + resolve(); + await promptPromise; + promptSpy.mockRestore(); + await tempDir.remove(); + }); + + it("guarantees always numeric nullable-vs-speculative contract", async () => { + const tempDir = TempDir.createSync("@always-numeric-"); + const { session, sessionManager, agent } = createSession(tempDir); + + const dummyId = sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "dummy" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 100, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 120, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + timestamp: 1000, + } as AssistantMessage); + + sessionManager.appendCompaction("compact summary", "compact summary", dummyId, 50); + syncSession(session, agent); + + const cu = session.getContextUsage(); + expect(cu?.tokens).not.toBeNull(); + expect(typeof cu?.tokens).toBe("number"); + expect(cu?.percent).not.toBeNull(); + expect(typeof cu?.percent).toBe("number"); + + await tempDir.remove(); + }); +}); diff --git a/packages/coding-agent/test/interactive-mode-plan-review.test.ts b/packages/coding-agent/test/interactive-mode-plan-review.test.ts index 149b0ba4a..4399bf46d 100644 --- a/packages/coding-agent/test/interactive-mode-plan-review.test.ts +++ b/packages/coding-agent/test/interactive-mode-plan-review.test.ts @@ -546,7 +546,7 @@ describe("InteractiveMode plan review rendering", () => { mode.planModeEnabled = true; mode.planModePlanFilePath = planFilePath; // Post-compaction: tokens unknown until the next LLM response. - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); const selector = vi.spyOn(mode, "showPlanReview").mockResolvedValue("Refine plan"); await mode.handlePlanApproval({ @@ -578,7 +578,7 @@ describe("InteractiveMode plan review rendering", () => { mode.planModeEnabled = true; mode.planModePlanFilePath = planFilePath; - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(mode, "showPlanReview").mockResolvedValue("Approve and keep context"); const clear = vi.spyOn(mode, "handleClearCommand").mockResolvedValue(); const prompt = vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); @@ -697,7 +697,7 @@ describe("InteractiveMode plan review rendering", () => { // Keep-context path avoids newSession() so the assertion isolates the // exit-plan-mode restore from session-clear effects. - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); let observedSegments: string[] = []; @@ -744,7 +744,7 @@ describe("InteractiveMode plan review rendering", () => { await mode.handlePlanModeCommand(); expect(session.model?.id).toBe(planModel.id); - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(mode, "showPlanReview").mockResolvedValue("Approve and compact context"); vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); @@ -781,7 +781,7 @@ describe("InteractiveMode plan review rendering", () => { await mode.handlePlanModeCommand(); expect(session.model?.id).toBe(planModel.id); - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(mode, "showPlanReview").mockResolvedValue("Approve and compact context"); const promptSpy = vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); @@ -826,7 +826,7 @@ describe("InteractiveMode plan review rendering", () => { await mode.handlePlanModeCommand(); expect(session.model?.id).toBe(planModel.id); - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); let compactModelId: string | undefined; @@ -885,7 +885,7 @@ describe("InteractiveMode plan review rendering", () => { await mode.handlePlanModeCommand(); expect(session.model?.id).toBe(planModel.id); - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(mode, "showPlanReview").mockResolvedValue("Approve and compact context"); const promptSpy = vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); @@ -932,7 +932,7 @@ describe("InteractiveMode plan review rendering", () => { await mode.handlePlanModeCommand(); expect(session.model?.id).toBe(planModel.id); - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); vi.spyOn(mode, "showPlanReview").mockResolvedValue("Approve and compact context"); vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); @@ -1240,7 +1240,7 @@ describe("InteractiveMode plan review rendering", () => { await mode.handlePlanModeCommand(); expect(session.getPlanModeState()?.planFilePath).toBe(planFilePath); - vi.spyOn(session, "getContextUsage").mockReturnValue({ tokens: null, contextWindow: 200000, percent: null }); + vi.spyOn(session, "getContextUsage").mockReturnValue(undefined); const selector = vi.spyOn(mode, "showPlanReview").mockResolvedValue("Approve and keep context"); const showError = vi.spyOn(mode, "showError"); vi.spyOn(session, "prompt").mockResolvedValue(undefined as never); diff --git a/packages/coding-agent/test/status-line-context-cache.test.ts b/packages/coding-agent/test/status-line-context-cache.test.ts index 68ead4385..307b77b51 100644 --- a/packages/coding-agent/test/status-line-context-cache.test.ts +++ b/packages/coding-agent/test/status-line-context-cache.test.ts @@ -155,20 +155,20 @@ describe("StatusLineComponent context breakdown", () => { expect(usageCalls()).toBe(2); }); - it("propagates an unknown (null) token count, e.g. right after compaction", () => { + it("propagates a speculative/numeric token count, e.g. right after compaction", () => { const { session } = makeSession({ messages: [userMessage("compaction summary")], - usage: { tokens: null, contextWindow: 272_000, percent: null }, + usage: { tokens: 1234, contextWindow: 272_000, percent: 0.45 }, }); const breakdown = new StatusLineComponent(session).getCachedContextBreakdown(); - expect(breakdown.usedTokens).toBeNull(); + expect(breakdown.usedTokens).toBe(1234); expect(breakdown.contextWindow).toBe(272_000); }); - it("falls back to the model window with null tokens when usage is unavailable", () => { + it("falls back to the model window with 0 tokens when usage is unavailable", () => { const { session } = makeSession({ messages: [userMessage("hi")], usage: undefined, contextWindow: 128_000 }); const breakdown = new StatusLineComponent(session).getCachedContextBreakdown(); - expect(breakdown.usedTokens).toBeNull(); + expect(breakdown.usedTokens).toBe(0); expect(breakdown.contextWindow).toBe(128_000); }); @@ -205,10 +205,10 @@ describe("StatusLineComponent context breakdown", () => { expect(plain).toContain("1.8%/272K"); }); - it("renders ? for the percent while the token count is unknown (post-compaction)", () => { + it("renders speculative percent instead of ? after compaction", () => { const { session } = makeSession({ messages: [userMessage("compaction summary")], - usage: { tokens: null, contextWindow: 272_000, percent: null }, + usage: { tokens: 1234, contextWindow: 272_000, percent: 0.45 }, }); const comp = new StatusLineComponent(session); comp.updateSettings({ @@ -219,6 +219,6 @@ describe("StatusLineComponent context breakdown", () => { }); const plain = comp.getTopBorder(80).content.replaceAll(/\x1b\[[0-9;]*m/g, ""); - expect(plain).toContain("?/272K"); + expect(plain).toContain("0.5%/272K"); }); });