From 994a820e43f607746fc26678f453a49785fdbf99 Mon Sep 17 00:00:00 2001 From: SSU-WEI HUANG Date: Sun, 31 May 2026 19:36:13 +0800 Subject: [PATCH] fix(opencode): surface ACP context usage live to web status bar (#756) --- .../agent/backends/acp/AcpSdkBackend.test.ts | 101 ++++++++++++++++++ cli/src/agent/backends/acp/AcpSdkBackend.ts | 51 +++++++-- 2 files changed, 146 insertions(+), 6 deletions(-) diff --git a/cli/src/agent/backends/acp/AcpSdkBackend.test.ts b/cli/src/agent/backends/acp/AcpSdkBackend.test.ts index 7fb579d1..59a0b45e 100644 --- a/cli/src/agent/backends/acp/AcpSdkBackend.test.ts +++ b/cli/src/agent/backends/acp/AcpSdkBackend.test.ts @@ -644,4 +644,105 @@ describe('AcpSdkBackend', () => { expect(stragglerIdx).toBeGreaterThanOrEqual(0); expect(turnCompleteIdx).toBeGreaterThan(stragglerIdx); }); + + it('emits a context-only usage mid-turn so the status bar updates live', async () => { + backendStatics.UPDATE_QUIET_PERIOD_MS = 25; + backendStatics.UPDATE_DRAIN_TIMEOUT_MS = 200; + backendStatics.PRE_PROMPT_UPDATE_QUIET_PERIOD_MS = 1; + backendStatics.PRE_PROMPT_UPDATE_DRAIN_TIMEOUT_MS = 50; + backendStatics.LATE_FLUSH_INTERVAL_MS = 5; + backendStatics.LATE_FLUSH_QUIET_PERIOD_MS = 10; + backendStatics.LATE_FLUSH_WINDOW_MS = 50; + + const backend = new AcpSdkBackend({ command: 'opencode' }); + const backendInternal = backend as unknown as { + transport: { + sendRequest: (...args: unknown[]) => Promise; + close: () => Promise; + } | null; + handleSessionUpdate: (params: unknown) => void; + }; + + const messages: AgentMessage[] = []; + backendInternal.transport = { + sendRequest: async () => { + backendInternal.handleSessionUpdate({ + sessionId: 'session-1', + update: { sessionUpdate: 'usage_update', used: 1_000, size: 200_000 } + }); + backendInternal.handleSessionUpdate({ + sessionId: 'session-1', + update: { sessionUpdate: 'usage_update', used: 1_000, size: 200_000 } + }); + backendInternal.handleSessionUpdate({ + sessionId: 'session-1', + update: { sessionUpdate: 'usage_update', used: 2_500, size: 200_000 } + }); + await sleep(5); + return { + stopReason: 'end_turn', + usage: { inputTokens: 100, outputTokens: 50 } + }; + }, + close: async () => {} + }; + + await backend.prompt('session-1', [{ type: 'text', text: 'hello' }], (m) => messages.push(m)); + + const usageMessages = messages.filter((m): m is Extract => m.type === 'usage'); + // Two mid-turn ticks (deduped second 1_000) + one final emit with the + // prompt-level input/output totals. + expect(usageMessages.length).toBe(3); + expect(usageMessages[0]).toMatchObject({ inputTokens: 0, outputTokens: 0, contextTokens: 1_000, contextWindow: 200_000 }); + expect(usageMessages[1]).toMatchObject({ inputTokens: 0, outputTokens: 0, contextTokens: 2_500, contextWindow: 200_000 }); + expect(usageMessages[2]).toMatchObject({ inputTokens: 100, outputTokens: 50, contextTokens: 2_500, contextWindow: 200_000 }); + }); + + it('emits a context-only usage on finalize when the prompt response carries no usage', async () => { + backendStatics.UPDATE_QUIET_PERIOD_MS = 25; + backendStatics.UPDATE_DRAIN_TIMEOUT_MS = 200; + backendStatics.PRE_PROMPT_UPDATE_QUIET_PERIOD_MS = 1; + backendStatics.PRE_PROMPT_UPDATE_DRAIN_TIMEOUT_MS = 50; + backendStatics.LATE_FLUSH_INTERVAL_MS = 5; + backendStatics.LATE_FLUSH_QUIET_PERIOD_MS = 10; + backendStatics.LATE_FLUSH_WINDOW_MS = 50; + + const backend = new AcpSdkBackend({ command: 'opencode' }); + const backendInternal = backend as unknown as { + transport: { + sendRequest: (...args: unknown[]) => Promise; + close: () => Promise; + } | null; + handleSessionUpdate: (params: unknown) => void; + }; + + const messages: AgentMessage[] = []; + backendInternal.transport = { + sendRequest: async () => { + backendInternal.handleSessionUpdate({ + sessionId: 'session-1', + update: { sessionUpdate: 'usage_update', used: 4_200, size: 200_000 } + }); + await sleep(5); + // No `usage` field on the response: simulates slash-handled + // turns or errored turns that skip the model. + return { stopReason: 'end_turn' }; + }, + close: async () => {} + }; + + await backend.prompt('session-1', [{ type: 'text', text: 'hi' }], (m) => messages.push(m)); + + const usageMessages = messages.filter((m): m is Extract => m.type === 'usage'); + // One mid-turn emit and one finalize fallback — both context-only. + expect(usageMessages.length).toBe(2); + for (const usage of usageMessages) { + expect(usage).toMatchObject({ + inputTokens: 0, + outputTokens: 0, + contextTokens: 4_200, + contextWindow: 200_000 + }); + } + }); }); diff --git a/cli/src/agent/backends/acp/AcpSdkBackend.ts b/cli/src/agent/backends/acp/AcpSdkBackend.ts index 234e37d9..c7b0e41a 100644 --- a/cli/src/agent/backends/acp/AcpSdkBackend.ts +++ b/cli/src/agent/backends/acp/AcpSdkBackend.ts @@ -55,6 +55,7 @@ export class AcpSdkBackend implements AgentBackend { private responseCompleteResolvers: Array<() => void> = []; private lastSessionUpdateAt = 0; private latestUsageUpdate: AcpUsageUpdate | null = null; + private activeOnUpdate: ((msg: AgentMessage) => void) | null = null; /** Retry configuration for ACP initialization */ private static readonly INIT_RETRY_OPTIONS = { @@ -283,6 +284,7 @@ export class AcpSdkBackend implements AgentBackend { ); this.messageHandler?.drainBuffers(); this.messageHandler = new AcpMessageHandler(onUpdate); + this.activeOnUpdate = onUpdate; this.isProcessingMessage = true; this.lastSessionUpdateAt = Date.now(); this.latestUsageUpdate = null; @@ -323,11 +325,27 @@ export class AcpSdkBackend implements AgentBackend { contextTokens: latestUsageUpdate ? latestUsageUpdate.contextTokens : undefined, contextWindow: latestUsageUpdate ? latestUsageUpdate.contextWindow : undefined }); + } else if ( + latestUsageUpdate + && (latestUsageUpdate.contextTokens !== undefined || latestUsageUpdate.contextWindow !== undefined) + ) { + // Agent did not return prompt usage (slash-handled turns, + // errored turns), but we did see ACP usage updates during + // the turn. Emit a context-only usage so the status bar + // reflects the current context size. + onUpdate({ + type: 'usage', + inputTokens: 0, + outputTokens: 0, + contextTokens: latestUsageUpdate.contextTokens, + contextWindow: latestUsageUpdate.contextWindow + }); } if (stopReason) { onUpdate({ type: 'turn_complete', stopReason }); } } finally { + this.activeOnUpdate = null; this.isProcessingMessage = false; this.notifyResponseComplete(); } @@ -407,6 +425,7 @@ export class AcpSdkBackend implements AgentBackend { if (!this.transport) return; this.messageHandler?.drainBuffers(); this.messageHandler = null; + this.activeOnUpdate = null; this.activeSessionId = null; this.isProcessingMessage = false; this.sessionModelsMetadata.clear(); @@ -431,12 +450,32 @@ export class AcpSdkBackend implements AgentBackend { if (!isObject(update)) return; if (asString(update.sessionUpdate) !== ACP_SESSION_UPDATE_TYPES.usageUpdate) return; - const contextTokens = this.asFiniteNumber(update.used); - const contextWindow = this.asFiniteNumber(update.size); - this.latestUsageUpdate = { - contextTokens: contextTokens ?? undefined, - contextWindow: contextWindow ?? undefined - }; + const contextTokens = this.asFiniteNumber(update.used) ?? undefined; + const contextWindow = this.asFiniteNumber(update.size) ?? undefined; + const prev = this.latestUsageUpdate; + const changed = !prev + || prev.contextTokens !== contextTokens + || prev.contextWindow !== contextWindow; + this.latestUsageUpdate = { contextTokens, contextWindow }; + + // Surface context updates mid-turn so the web status bar shows live + // ctx N/M (X%) instead of staying blank until the final prompt usage + // arrives. ACP usage_update only carries context tokens, so I/O is + // sent as 0; the final prompt-finalize emit overwrites with the real + // input/output totals. + if ( + changed + && this.activeOnUpdate + && (contextTokens !== undefined || contextWindow !== undefined) + ) { + this.activeOnUpdate({ + type: 'usage', + inputTokens: 0, + outputTokens: 0, + contextTokens, + contextWindow + }); + } } private readLatestUsageUpdate(): AcpUsageUpdate | null {