diff --git a/.changeset/acp-context-window-usage.md b/.changeset/acp-context-window-usage.md new file mode 100644 index 0000000000..f7c768dc98 --- /dev/null +++ b/.changeset/acp-context-window-usage.md @@ -0,0 +1,5 @@ +--- +"@moonshot-ai/kimi-code": minor +--- + +ACP agents now push live context window usage via the `usage_update` session notification so compatible clients (Zed, JetBrains) can render a real-time token consumption indicator. diff --git a/packages/acp-adapter/src/events-map.ts b/packages/acp-adapter/src/events-map.ts index 0448f2eb9c..85648cc688 100644 --- a/packages/acp-adapter/src/events-map.ts +++ b/packages/acp-adapter/src/events-map.ts @@ -525,3 +525,34 @@ export function configOptionUpdateNotification( }, }; } + +/** + * Build an ACP `usage_update` session notification from contextual + * token usage data (e.g. an `agent.status.updated` SDK event). + * + * Wire shape ({@link https://agentclientprotocol.com/ protocol} UNSTABLE / + * `@experimental`): + * - `size` = total context window capacity (maxContextTokens) + * - `used` = tokens currently in context (contextTokens) + * - `cost` = cumulative cost snapshot (unset for now — can be wired + * later from the session cost tracker). + * + * Returns `null` when either value is missing so callers can silently + * skip rather than pushing a degenerate update. + */ +export function contextUsageToUsageUpdate( + sessionId: string, + contextTokens: number | undefined, + maxContextTokens: number | undefined, +): SessionNotification | null { + if (contextTokens === undefined || maxContextTokens === undefined) return null; + if (maxContextTokens <= 0) return null; + return { + sessionId, + update: { + sessionUpdate: 'usage_update', + size: maxContextTokens, + used: contextTokens, + }, + }; +} diff --git a/packages/acp-adapter/src/index.ts b/packages/acp-adapter/src/index.ts index 4b0af447f9..5dbf6e9174 100644 --- a/packages/acp-adapter/src/index.ts +++ b/packages/acp-adapter/src/index.ts @@ -19,6 +19,7 @@ export { export { acpToolCallId, assistantDeltaToSessionUpdate, + contextUsageToUsageUpdate, inferToolKind, stringifyArgs, thinkingDeltaToSessionUpdate, diff --git a/packages/acp-adapter/src/server.ts b/packages/acp-adapter/src/server.ts index 6707fd4cae..5710f9830b 100644 --- a/packages/acp-adapter/src/server.ts +++ b/packages/acp-adapter/src/server.ts @@ -421,6 +421,7 @@ export class AcpServer implements Agent { DEFAULT_MODE_ID, ); this.scheduleAvailableCommandsUpdate(session.id); + void acpSession.pushContextUsage(); return { sessionId: session.id, configOptions, @@ -464,6 +465,7 @@ export class AcpServer implements Agent { // `resumeSession`, which intentionally omits this step. await acpSession.replayHistory(); this.scheduleAvailableCommandsUpdate(session.id); + void acpSession.pushContextUsage(); return { configOptions }; } @@ -494,6 +496,7 @@ export class AcpServer implements Agent { mode: 'resume', }); this.scheduleAvailableCommandsUpdate(session.id); + void this.sessions.get(session.id)?.pushContextUsage(); return { configOptions }; } diff --git a/packages/acp-adapter/src/session.ts b/packages/acp-adapter/src/session.ts index 747b44ea9c..018d29d9d1 100644 --- a/packages/acp-adapter/src/session.ts +++ b/packages/acp-adapter/src/session.ts @@ -44,6 +44,7 @@ import { acpToolCallId, assistantDeltaToSessionUpdate, configOptionUpdateNotification, + contextUsageToUsageUpdate, planFromDisplayBlock, stringifyArgs, thinkingDeltaToSessionUpdate, @@ -550,6 +551,41 @@ export class AcpSession { } } + /** + * Push a one-shot `usage_update` snapshot derived from the SDK + * `getStatus()` call. Best-effort — failures are logged and never + * surface to the client as errors. + * + * Called on session creation and after history replay so the client + * sees context usage immediately instead of waiting for the next + * `agent.status.updated` event (which only fires inside an active + * turn, via `runTurnBody`). + */ + async pushContextUsage(): Promise { + if (typeof this.session.getStatus !== 'function') return; + try { + const status = await this.session.getStatus(); + const update = contextUsageToUsageUpdate( + this.id, + status.contextTokens, + status.maxContextTokens, + ); + if (update !== null) { + this.conn.sessionUpdate(update).catch((err) => { + log.warn('acp: failed to push usage_update', { + sessionId: this.id, + error: err instanceof Error ? err.message : String(err), + }); + }); + } + } catch (err) { + log.warn('acp: failed to push usage_update', { + sessionId: this.id, + error: err instanceof Error ? err.message : String(err), + }); + } + } + /** * Replay the underlying SDK session's persisted history as a stream * of ACP `session/update` notifications. @@ -964,6 +1000,34 @@ export class AcpSession { const outcome = await completion; if (outcome.kind === 'completed') { await this.emitLocalCommandMessage(formatCompactionCompleted(outcome.result)); + // After compaction, the core updates context token counts via + // emitStatusUpdated(), but the agent.status.updated handler is + // only registered in runTurnBody(). Push a one-shot usage_update + // so the client reflects the compacted context size immediately + // rather than waiting for the next normal prompt. + // Best-effort only — a failed status refresh must not turn a + // successful compaction into a failure. + try { + const status = await this.session.getStatus(); + const update = contextUsageToUsageUpdate( + this.id, + status.contextTokens, + status.maxContextTokens, + ); + if (update !== null) { + this.conn.sessionUpdate(update).catch((err) => { + log.warn('acp: failed to push usage_update after compaction', { + sessionId: this.id, + error: err instanceof Error ? err.message : String(err), + }); + }); + } + } catch (err) { + log.warn('acp: failed to refresh usage after compaction', { + sessionId: this.id, + error: err instanceof Error ? err.message : String(err), + }); + } } else { await this.emitLocalCommandMessage('Compaction cancelled.'); } @@ -1097,6 +1161,23 @@ export class AcpSession { }); return; } + if (event.type === 'agent.status.updated') { + if (!isFromMainAgent(event)) return; + const update = contextUsageToUsageUpdate( + sessionId, + event.contextTokens, + event.maxContextTokens, + ); + if (update !== null) { + conn.sessionUpdate(update).catch((err) => { + log.warn('acp: failed to push usage_update', { + sessionId, + error: err instanceof Error ? err.message : String(err), + }); + }); + } + return; + } if (event.type === 'tool.call.started') { if (!isFromMainAgent(event)) return; // Seed the accumulator with the **stringified initial args**. diff --git a/packages/acp-adapter/test/plan-and-commands.test.ts b/packages/acp-adapter/test/plan-and-commands.test.ts index cc5463d736..a5e119d103 100644 --- a/packages/acp-adapter/test/plan-and-commands.test.ts +++ b/packages/acp-adapter/test/plan-and-commands.test.ts @@ -20,6 +20,7 @@ import { AcpServer } from '../src/server'; import { AUTHED_STATUS } from './_helpers/harness-stubs'; import { availableCommandsUpdateNotification, + contextUsageToUsageUpdate, planFromDisplayBlock, todoListToSessionUpdate, } from '../src/events-map'; @@ -359,3 +360,162 @@ describe('Phase 9.3 e2e · todo_list display block becomes a plan session_update expect(planUpdates).toHaveLength(0); }); }); + +describe('contextUsageToUsageUpdate', () => { + it('builds a usage_update notification with size and used', () => { + const note = contextUsageToUsageUpdate('sess-u', 12345, 200000); + expect(note).not.toBeNull(); + expect(note?.sessionId).toBe('sess-u'); + expect(note?.update).toEqual({ + sessionUpdate: 'usage_update', + size: 200000, + used: 12345, + }); + }); + + it('returns null when contextTokens is undefined', () => { + expect(contextUsageToUsageUpdate('sess-u', undefined, 200000)).toBeNull(); + }); + + it('returns null when maxContextTokens is undefined', () => { + expect(contextUsageToUsageUpdate('sess-u', 12345, undefined)).toBeNull(); + }); + + it('returns null when both values are undefined', () => { + expect(contextUsageToUsageUpdate('sess-u', undefined, undefined)).toBeNull(); + }); + + it('returns null when maxContextTokens is zero or negative', () => { + expect(contextUsageToUsageUpdate('sess-u', 100, 0)).toBeNull(); + expect(contextUsageToUsageUpdate('sess-u', 100, -1)).toBeNull(); + }); +}); + +describe('e2e · agent.status.updated emits usage_update', () => { + it('forwards agent.status.updated events as usage_update session notifications', async () => { + const sessionId = 'sess-usage'; + const turnId = 1; + const session = makeScriptedSession(sessionId, [ + { + type: 'agent.status.updated', + sessionId, + agentId: 'main', + contextTokens: 5000, + maxContextTokens: 100000, + contextUsage: 0.05, + } as Event, + { + type: 'turn.ended', + sessionId, + agentId: 'main', + turnId, + reason: 'completed', + } as Event, + ]); + const harness = { + auth: { status: async () => AUTHED_STATUS }, + createSession: async () => session, + } as unknown as KimiHarness; + + const { agentStream, clientStream } = makeInMemoryStreamPair(); + new AgentSideConnection((c) => new AcpServer(harness, c), agentStream); + const collecting = new CollectingClient(); + const client = new ClientSideConnection(() => collecting, clientStream); + + await client.newSession({ cwd: '/tmp/x', mcpServers: [] }); + await client.prompt({ sessionId, prompt: [textBlock('hello')] }); + await flushNdjson(); + + const usageUpdates = collecting.updates.filter( + (n) => + (n.update as { sessionUpdate: string }).sessionUpdate === 'usage_update', + ); + expect(usageUpdates).toHaveLength(1); + expect(usageUpdates[0]?.sessionId).toBe(sessionId); + expect(usageUpdates[0]?.update).toEqual({ + sessionUpdate: 'usage_update', + size: 100000, + used: 5000, + }); + }); + + it('does not emit usage_update for sub-agent status events', async () => { + const sessionId = 'sess-usage-sub'; + const turnId = 1; + const session = makeScriptedSession(sessionId, [ + { + type: 'agent.status.updated', + sessionId, + agentId: 'sub-1', + contextTokens: 100, + maxContextTokens: 50000, + } as Event, + { + type: 'turn.ended', + sessionId, + agentId: 'main', + turnId, + reason: 'completed', + } as Event, + ]); + const harness = { + auth: { status: async () => AUTHED_STATUS }, + createSession: async () => session, + } as unknown as KimiHarness; + + const { agentStream, clientStream } = makeInMemoryStreamPair(); + new AgentSideConnection((c) => new AcpServer(harness, c), agentStream); + const collecting = new CollectingClient(); + const client = new ClientSideConnection(() => collecting, clientStream); + + await client.newSession({ cwd: '/tmp/x', mcpServers: [] }); + await client.prompt({ sessionId, prompt: [textBlock('hello')] }); + await flushNdjson(); + + const usageUpdates = collecting.updates.filter( + (n) => + (n.update as { sessionUpdate: string }).sessionUpdate === 'usage_update', + ); + expect(usageUpdates).toHaveLength(0); + }); + + it('skips status events with missing contextTokens', async () => { + const sessionId = 'sess-usage-skip'; + const turnId = 1; + const session = makeScriptedSession(sessionId, [ + { + type: 'agent.status.updated', + sessionId, + agentId: 'main', + maxContextTokens: 100000, + // contextTokens deliberately omitted + } as Event, + { + type: 'turn.ended', + sessionId, + agentId: 'main', + turnId, + reason: 'completed', + } as Event, + ]); + const harness = { + auth: { status: async () => AUTHED_STATUS }, + createSession: async () => session, + } as unknown as KimiHarness; + + const { agentStream, clientStream } = makeInMemoryStreamPair(); + new AgentSideConnection((c) => new AcpServer(harness, c), agentStream); + const collecting = new CollectingClient(); + const client = new ClientSideConnection(() => collecting, clientStream); + + await client.newSession({ cwd: '/tmp/x', mcpServers: [] }); + await client.prompt({ sessionId, prompt: [textBlock('hello')] }); + await flushNdjson(); + + const usageUpdates = collecting.updates.filter( + (n) => + (n.update as { sessionUpdate: string }).sessionUpdate === 'usage_update', + ); + expect(usageUpdates).toHaveLength(0); + }); +}); diff --git a/packages/acp-adapter/test/session-new.test.ts b/packages/acp-adapter/test/session-new.test.ts index 61220b462b..2a32bf467a 100644 --- a/packages/acp-adapter/test/session-new.test.ts +++ b/packages/acp-adapter/test/session-new.test.ts @@ -266,7 +266,9 @@ describe('AcpServer session/new', () => { const response = await client.newSession({ cwd: '/tmp/work', mcpServers: [] }); - expect(fakeSession.getStatus).toHaveBeenCalledOnce(); + // Called twice: once by resolveCurrentThinkingEffort, once by + // pushContextUsage (best-effort usage_update after session creation). + expect(fakeSession.getStatus).toHaveBeenCalledTimes(2); const thinking = response.configOptions?.find((option) => option.id === 'thinking'); if (thinking?.type !== 'select') throw new Error('thinking option must be a select'); expect(thinking.currentValue).toBe(expected); diff --git a/packages/acp-adapter/test/session-slash.test.ts b/packages/acp-adapter/test/session-slash.test.ts index f13dea0b1d..15c8c2ae2e 100644 --- a/packages/acp-adapter/test/session-slash.test.ts +++ b/packages/acp-adapter/test/session-slash.test.ts @@ -140,6 +140,10 @@ async function waitForAvailableCommands( throw new Error('available_commands_update never arrived'); } +async function flushNdjson(): Promise { + await new Promise((resolve) => setTimeout(resolve, 25)); +} + describe('AcpSession slash routing', () => { it('routes `/skill:foo bar` to Session.activateSkill (not Session.prompt)', async () => { const sessionId = 'sess-slash-A'; @@ -371,4 +375,79 @@ describe('AcpSession slash routing', () => { expect(text).toContain('Model: mock-model'); expect(text).toContain('Context: 1,234 / 200,000 (0.6%)'); }); + + it('pushes usage_update notification after /compact completes', async () => { + const sessionId = 'sess-slash-compact-usage'; + // We need a fake session with controlled listeners so compact() can + // emit compaction lifecycle events through the adapter's onEvent + // subscription (which is registered before compact() is called). + const listeners = new Set<(event: Event) => void>(); + const session = { + id: sessionId, + prompt: async () => undefined, + cancel: async () => undefined, + activateSkill: async () => undefined, + listSkills: async () => [] as const, + onEvent: (fn: (event: Event) => void) => { + listeners.add(fn); + return () => { listeners.delete(fn); }; + }, + compact: async () => { + // The adapter subscribes to onEvent (line 926) before calling + // compact() (line 960), so listeners already contains the + // adapter's handler. Emit the compaction lifecycle. + // Use setTimeout to yield so the adapter's subscription is active. + await new Promise((resolve) => { + setTimeout(() => { + for (const fn of listeners) { + fn({ type: 'compaction.started', sessionId, agentId: 'main' } as Event); + fn({ + type: 'compaction.completed', + sessionId, + agentId: 'main', + result: { tokensBefore: 50000, tokensAfter: 5000, compactedCount: 45, keptUserMessageCount: 1, droppedCount: 0 }, + } as Event); + } + resolve(); + }, 0); + }); + }, + getStatus: async () => ({ + model: 'mock-model', + thinkingEffort: 'low', + permission: 'ask', + planMode: false, + contextTokens: 5000, + maxContextTokens: 200000, + contextUsage: 0.025, + }), + } as unknown as Session; + const harness = { + auth: { status: async () => AUTHED_STATUS }, + createSession: async () => session, + } as unknown as KimiHarness; + + const { agentStream, clientStream } = makeInMemoryStreamPair(); + new AgentSideConnection((c) => new AcpServer(harness, c), agentStream); + const collecting = new CollectingClient(); + const client = new ClientSideConnection(() => collecting, clientStream); + + await client.newSession({ cwd: '/tmp/x', mcpServers: [] }); + await waitForAvailableCommands(collecting); + const beforeCount = collecting.updates.length; + + await client.prompt({ sessionId, prompt: [textBlock('/compact')] }); + await flushNdjson(); + + const usageUpdates = collecting.updates.slice(beforeCount).filter( + (n) => + (n.update as { sessionUpdate: string }).sessionUpdate === 'usage_update', + ); + expect(usageUpdates).toHaveLength(1); + expect(usageUpdates[0]?.update).toEqual({ + sessionUpdate: 'usage_update', + size: 200000, + used: 5000, + }); + }); });