From 409196bf2e1b840956b1fb38ea5a08ce4bd77603 Mon Sep 17 00:00:00 2001 From: can1357 Date: Wed, 17 Jun 2026 12:19:49 +0200 Subject: [PATCH] fix(coding-agent/session): fixed context usage breakdown to prefer completed in-turn anchors - Fixed context breakdown to anchor estimates on the latest completed assistant usage message after compaction. - Adjusted pending-context usage selection to prefer an in-turn provider anchor when available at/after cutoff. - Added a contextUsageRevision cache token so status-line context memo invalidates after snapshot clear. --- packages/ai/src/registry/oauth/index.ts | 16 +-- packages/coding-agent/CHANGELOG.md | 3 +- packages/coding-agent/src/collab/host.ts | 2 +- .../modes/components/status-line/component.ts | 7 + .../coding-agent/src/session/agent-session.ts | 136 ++++++++++-------- .../test/context-consolidation.test.ts | 99 +++++++++++++ .../session-manager/usage-statistics.test.ts | 36 +++++ .../test/status-line-context-cache.test.ts | 27 ++++ 8 files changed, 260 insertions(+), 66 deletions(-) diff --git a/packages/ai/src/registry/oauth/index.ts b/packages/ai/src/registry/oauth/index.ts index 82cb38f66..32a6ce01e 100644 --- a/packages/ai/src/registry/oauth/index.ts +++ b/packages/ai/src/registry/oauth/index.ts @@ -142,14 +142,14 @@ export async function getOAuthApiKey( provider === "alibaba-coding-plan"; const apiKey = needsStructuredApiKey ? JSON.stringify({ - token: creds.access, - enterpriseUrl: creds.enterpriseUrl, - projectId: creds.projectId, - refreshToken: creds.refresh, - expiresAt: creds.expires, - email: creds.email, - accountId: creds.accountId, - }) + token: creds.access, + enterpriseUrl: creds.enterpriseUrl, + projectId: creds.projectId, + refreshToken: creds.refresh, + expiresAt: creds.expires, + email: creds.email, + accountId: creds.accountId, + }) : creds.access; return { newCredentials: creds, apiKey }; } diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index dfba9b647..c982c11fe 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,7 +1,6 @@ # Changelog ## [Unreleased] - ### Added - Added `images.describeForTextModels` option (default `true`) to control automatic image description for attachments sent to models without vision input @@ -14,9 +13,11 @@ ### Fixed +- Fixed context usage breakdown to use a completed assistant usage anchor from the current turn instead of a pending prompt snapshot so totals no longer overcount when a large in-turn tool step returns usage - 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 the status-line context% reading inflated during long tool turns and then dropping sharply on the next message even though no compaction ran. While a request was in flight `getContextBreakdown` summed a cl100k estimate of the entire tail on top of the stale turn-start prompt and never re-anchored to completed in-turn steps; it now prefers the real provider prompt-token count of any step that resolves at or after the pending cutoff. The status-line memo also keys on a `contextUsageRevision` that bumps when the in-flight snapshot is set/cleared, so a mid-turn estimate is invalidated on turn end/abort instead of surviving into idle until the next message - 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 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)). diff --git a/packages/coding-agent/src/collab/host.ts b/packages/coding-agent/src/collab/host.ts index 0c52c2521..55850839a 100644 --- a/packages/coding-agent/src/collab/host.ts +++ b/packages/coding-agent/src/collab/host.ts @@ -415,7 +415,7 @@ export class CollabHost { // render exactly the same anchored, provider-real count the host's own // status line shows. const breakdown = this.#ctx.statusLine.getCachedContextBreakdown(); - const tokens = breakdown.usedTokens; + const tokens = breakdown.usedTokens ?? 0; return { isStreaming: session.isStreaming, isAborting: session.isAborting, 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 75670a2b1..d8ebce02f 100644 --- a/packages/coding-agent/src/modes/components/status-line/component.ts +++ b/packages/coding-agent/src/modes/components/status-line/component.ts @@ -145,6 +145,7 @@ interface ContextUsageMemo { length: number; lastFingerprint: string | undefined; modelContextWindow: number; + contextUsageRevision: number; usedTokens: number; contextWindow: number; systemPromptRef: readonly string[] | undefined; @@ -603,6 +604,10 @@ export class StatusLineComponent implements Component { const modelContextWindow = this.session.model?.contextWindow ?? 0; const length = messages.length; const lastFingerprint = length > 0 ? messageFingerprint(messages[length - 1]!) : undefined; + // Bumps when the in-flight pending snapshot is set/cleared. Without it a + // value computed mid-turn (estimate of the active tail) would survive after + // the turn ends/aborts, since clearing the snapshot touches no message. + const contextUsageRevision = this.session.contextUsageRevision ?? 0; const systemPrompt = this.session.systemPrompt; const tools = this.session.agent?.state?.tools; @@ -615,6 +620,7 @@ export class StatusLineComponent implements Component { cache.length === length && cache.lastFingerprint === lastFingerprint && cache.modelContextWindow === modelContextWindow && + cache.contextUsageRevision === contextUsageRevision && cache.systemPromptRef === systemPrompt && cache.toolsRef === tools && cache.skillsRef === skills @@ -630,6 +636,7 @@ export class StatusLineComponent implements Component { length, lastFingerprint, modelContextWindow, + contextUsageRevision, usedTokens, contextWindow, systemPromptRef: systemPrompt, diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index c0a7ab907..d16dc10c4 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -1247,6 +1247,11 @@ export class AgentSession { cutoffCount: number; } | undefined = undefined; + // Bumped whenever the pending in-flight snapshot is set/cleared. The + // status-line context memo includes this so clearing the snapshot on + // turn-end/abort invalidates the cache even though the message list is + // unchanged — otherwise a mid-turn estimate would survive into idle. + #contextUsageRevision = 0; #obfuscator: SecretObfuscator | undefined; #checkpointState: CheckpointState | undefined = undefined; #pendingRewindReport: string | undefined = undefined; @@ -5593,15 +5598,15 @@ export class AgentSession { nonMessageTokens + this.messages.reduce((sum, msg) => sum + estimateTokens(msg), 0) + messages.reduce((sum, msg) => sum + estimateTokens(msg), 0); - this.#pendingContextSnapshot = { + this.#setPendingContextSnapshot({ promptTokens, nonMessageTokens, cutoffCount: this.messages.length + messages.length, - }; + }); try { await this.#promptAgentWithIdleRetry(messages, agentPromptOptions); } finally { - this.#pendingContextSnapshot = undefined; + this.#setPendingContextSnapshot(undefined); } if (!options?.skipPostPromptRecoveryWait) { await this.#waitForPostPromptRecovery(generation); @@ -11276,76 +11281,80 @@ export class AgentSession { const pendingMessages = options?.pendingMessages ?? []; - let anchorEntry: SessionMessageEntry | undefined; - let isPending = false; + const pending = this.#pendingContextSnapshot; - 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" && assistant.usage) { - anchorEntry = entry; - break; - } + // Always locate the latest real assistant-usage anchor after the last + // compaction. Its provider-reported promptTokens is ground truth for + // everything up to that point; only the tail after it is estimated. + let anchorEntry: SessionMessageEntry | undefined; + 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" && assistant.usage) { + anchorEntry = entry; + break; } } } - 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]); - } + const resolvedActiveMessages = this.messages; + let resolvedAnchorIndex = -1; + let anchorAssistant: AssistantMessage | undefined; + if (anchorEntry) { + const a = anchorEntry.message as AssistantMessage; + anchorAssistant = a; + resolvedAnchorIndex = resolvedActiveMessages.indexOf(a); + if (resolvedAnchorIndex === -1) { + resolvedAnchorIndex = resolvedActiveMessages.findIndex( + msg => msg.role === "assistant" && msg.timestamp === a.timestamp, + ); } + } - 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; + // A real anchor supersedes the in-flight estimate only once a step of the + // CURRENT turn has produced provider usage — i.e. it resolves at or after + // the pending cutoff. While the turn's first response is still pending (or + // the newest real anchor predates this turn) the pending snapshot is the + // only thing accounting for the just-submitted prompt, so it wins. This + // keeps a long tool turn from stacking an estimate of the entire tail on + // top of a stale turn-start prompt. + const useAnchor = + anchorAssistant !== undefined && + resolvedAnchorIndex !== -1 && + (!pending || resolvedAnchorIndex >= pending.cutoffCount); + + if (useAnchor && anchorAssistant) { 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, - ); + let tailTokens = 0; + for (let i = resolvedAnchorIndex + 1; i < resolvedActiveMessages.length; i++) { + tailTokens += estimateTokens(resolvedActiveMessages[i]); } - - if (resolvedAnchorIndex !== -1) { - let tailTokens = 0; - for (let i = resolvedAnchorIndex + 1; i < resolvedActiveMessages.length; i++) { + usedTokens = + promptTokens + + Math.max(0, currentNonMessageTokens - nonMessageTokens) + + tailTokens + + pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0); + } else if (pending) { + anchored = true; + let tailTokens = 0; + if (resolvedActiveMessages.length > pending.cutoffCount) { + for (let i = pending.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 { - anchored = false; } + usedTokens = + pending.promptTokens + + Math.max(0, currentNonMessageTokens - pending.nonMessageTokens) + + tailTokens + + pendingMessages.reduce((sum, msg) => sum + estimateTokens(msg), 0); } - if (!anchored && !isPending && branchEntries.length === 0) { + if (!anchored && !pending && 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) { @@ -11368,7 +11377,6 @@ export class AgentSession { } } if (!anchored) { - const resolvedActiveMessages = this.messages; let messagesTokens = 0; for (const msg of resolvedActiveMessages) { messagesTokens += estimateTokens(msg); @@ -11403,6 +11411,22 @@ export class AgentSession { }; } + /** + * Monotonic counter that changes whenever the in-flight pending context + * snapshot is set or cleared. Status-line context memoization keys on this so + * a value computed mid-turn cannot persist after the turn ends/aborts. + */ + get contextUsageRevision(): number { + return this.#contextUsageRevision; + } + + #setPendingContextSnapshot( + snapshot: { promptTokens: number; nonMessageTokens: number; cutoffCount: number } | undefined, + ): void { + this.#pendingContextSnapshot = snapshot; + this.#contextUsageRevision++; + } + #ingestProviderUsageHeaders(response: ProviderResponseMetadata, model?: Model): void { if (model?.provider !== "anthropic") return; this.#modelRegistry.authStorage.ingestUsageHeaders("anthropic", response.headers, { diff --git a/packages/coding-agent/test/context-consolidation.test.ts b/packages/coding-agent/test/context-consolidation.test.ts index 735b54880..7fe0e6628 100644 --- a/packages/coding-agent/test/context-consolidation.test.ts +++ b/packages/coding-agent/test/context-consolidation.test.ts @@ -409,6 +409,105 @@ describe("Context usage consolidation", () => { await tempDir.remove(); }); + it("prefers a completed in-turn provider anchor over the pending snapshot", async () => { + const tempDir = TempDir.createSync("@inturn-anchor-"); + const { session, sessionManager, agent } = createSession(tempDir); + + const { promise, resolve } = Promise.withResolvers(); + const promptSpy = vi.spyOn(agent, "prompt").mockImplementation(async () => { + await promise; + }); + + const promptPromise = session.prompt("new question"); + await Bun.sleep(10); + + // While the request hangs (pending snapshot active, cutoff at the new + // prompt), simulate the turn's user message landing plus a completed tool + // step carrying a real, large provider prompt count. + sessionManager.appendMessage({ role: "user", content: "new question", timestamp: 5000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "step" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "toolUse", + usage: { + input: 9000, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 9020, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 9000, nonMessageTokens: 10 }, + timestamp: 6000, + } as AssistantMessage); + syncSession(session, agent); + + // The in-turn anchor (index >= pending cutoff) must win: usage reflects the + // real 9000-token prompt, not the tiny turn-start pending estimate that + // would otherwise stack an estimate of the whole tail on top. + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(true); + expect(breakdown?.usedTokens).toBeGreaterThanOrEqual(9000); + + resolve(); + await promptPromise; + promptSpy.mockRestore(); + await tempDir.remove(); + }); + + it("keeps the pending snapshot (not a pre-cutoff anchor) while the first turn response is pending", async () => { + const tempDir = TempDir.createSync("@pending-precutoff-"); + const { session, sessionManager, agent } = createSession(tempDir); + + // Prior completed turn establishes a real anchor that PREDATES the new + // prompt (resolves before the pending cutoff). + sessionManager.appendMessage({ role: "user", content: "old", timestamp: 1000 } as Message); + sessionManager.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "old response" }], + api: mockModel.api, + provider: mockModel.provider, + model: mockModel.id, + stopReason: "stop", + usage: { + input: 5000, + output: 20, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 5020, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + contextSnapshot: { promptTokens: 5000, nonMessageTokens: 10 }, + timestamp: 2000, + } as AssistantMessage); + syncSession(session, agent); + + const { promise, resolve } = Promise.withResolvers(); + const promptSpy = vi.spyOn(agent, "prompt").mockImplementation(async () => { + await promise; + }); + + // Large prompt in flight, no in-turn step has produced provider usage yet. + const bigPrompt = "explain this in detail ".repeat(200); + const promptPromise = session.prompt(bigPrompt); + await Bun.sleep(10); + + // Pending must win: it accounts for the just-submitted prompt on top of the + // prior anchor. The stale pre-cutoff anchor alone (5000) would omit it, so a + // regression to "always prefer the latest real anchor" would read ~5000. + const breakdown = session.getContextBreakdown(); + expect(breakdown?.anchored).toBe(true); + expect(breakdown?.usedTokens).toBeGreaterThan(5500); + + 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); diff --git a/packages/coding-agent/test/session-manager/usage-statistics.test.ts b/packages/coding-agent/test/session-manager/usage-statistics.test.ts index dc1ddf095..0881f2ebc 100644 --- a/packages/coding-agent/test/session-manager/usage-statistics.test.ts +++ b/packages/coding-agent/test/session-manager/usage-statistics.test.ts @@ -120,4 +120,40 @@ describe("SessionManager usage statistics", () => { const usage = session.getUsageStatistics(); expect(usage.premiumRequests).toBe(0); }); + + it("accumulates the full billed cost across turns, including cache-read cost", () => { + // Contract: the session cost aggregate sums each turn's full `cost.total` + // (input+output+cacheRead+cacheWrite), not a cache-excluded "new-work" + // subset. Cache-read cost is real billed spend — the cached context is + // re-read at the cache-read rate every turn — so it must stay in the + // ledger that /usage, ACP usage_update, and hooks consume. Two turns with + // nonzero cacheRead make the readings diverge: full total = 18 vs the + // excluded subset (input+output+cacheWrite) = 8. + const session = SessionManager.inMemory(); + + session.appendMessage({ role: "user", content: "hello", timestamp: 1 }); + for (const timestamp of [2, 3]) { + session.appendMessage({ + role: "assistant", + content: [{ type: "text", text: "hi" }], + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4", + usage: { + input: 1, + output: 2, + cacheRead: 100, + cacheWrite: 10, + totalTokens: 113, + cost: { input: 1, output: 2, cacheRead: 5, cacheWrite: 1, total: 9 }, + }, + stopReason: "stop", + timestamp, + }); + } + + const usage = session.getUsageStatistics(); + expect(usage.cacheRead).toBe(200); + expect(usage.cost).toBeCloseTo(18, 8); + }); }); 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 307b77b51..26067b35f 100644 --- a/packages/coding-agent/test/status-line-context-cache.test.ts +++ b/packages/coding-agent/test/status-line-context-cache.test.ts @@ -36,12 +36,15 @@ interface Fake { usageCalls: () => number; /** Swap the value the next `getContextUsage()` query returns. */ setUsage: (usage: ContextUsage | undefined) => void; + /** Bump the in-flight pending revision the next `getCachedContextBreakdown()` reads. */ + setRevision: (n: number) => void; } function makeSession(opts: { messages: unknown[]; contextWindow?: number; usage?: ContextUsage | undefined }): Fake { const contextWindow = opts.contextWindow ?? 200_000; let usage: ContextUsage | undefined = "usage" in opts ? opts.usage : { tokens: 1234, contextWindow, percent: 0.6 }; let calls = 0; + let revision = 0; const session = { messages: opts.messages, systemPrompt: ["You are a helpful assistant."], @@ -65,6 +68,9 @@ function makeSession(opts: { messages: unknown[]; contextWindow?: number; usage? calls++; return usage; }, + get contextUsageRevision() { + return revision; + }, } as unknown as AgentSession; return { session, @@ -72,6 +78,9 @@ function makeSession(opts: { messages: unknown[]; contextWindow?: number; usage? setUsage: next => { usage = next; }, + setRevision: (n: number) => { + revision = n; + }, }; } @@ -155,6 +164,24 @@ describe("StatusLineComponent context breakdown", () => { expect(usageCalls()).toBe(2); }); + it("re-queries when only the in-flight pending revision changes (no message change)", () => { + const fake = makeSession({ + messages: [userMessage("hi")], + usage: { tokens: 190_000, contextWindow: 272_000, percent: 69.9 }, + }); + const comp = new StatusLineComponent(fake.session); + expect(comp.getCachedContextBreakdown().usedTokens).toBe(190_000); + + // Turn ends/aborts: the message list and last-message fingerprint are + // unchanged, but clearing the pending snapshot recalibrates usage to the + // real provider anchor. The memo must not keep serving the stale estimate. + fake.setUsage({ tokens: 117_000, contextWindow: 272_000, percent: 43.0 }); + fake.setRevision(1); + + expect(comp.getCachedContextBreakdown().usedTokens).toBe(117_000); + expect(fake.usageCalls()).toBe(2); + }); + it("propagates a speculative/numeric token count, e.g. right after compaction", () => { const { session } = makeSession({ messages: [userMessage("compaction summary")],