feat(acp-adapter): Support displaying context length usage in Zed ACP like Codex

.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.

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,
+    },
+  };
+}

packages/acp-adapter/src/index.ts

@@ -19,6 +19,7 @@ export {
 export {
   acpToolCallId,
   assistantDeltaToSessionUpdate,
+  contextUsageToUsageUpdate,
   inferToolKind,
   stringifyArgs,
   thinkingDeltaToSessionUpdate,

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 };
   }
 

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<void> {
+    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**.

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);
+  });
+});

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);