fix(grok): show context-window usage in the composer

apps/server/scripts/acp-mock-agent.ts

@@ -46,6 +46,7 @@ const emitOverlappingXAiPromptCompleteOutOfOrder =
 const failPrompt = process.env.T3_ACP_FAIL_PROMPT === "1";
 const failSetConfigOption = process.env.T3_ACP_FAIL_SET_CONFIG_OPTION === "1";
 const exitOnSetConfigOption = process.env.T3_ACP_EXIT_ON_SET_CONFIG_OPTION === "1";
+const emitContextUsage = process.env.T3_ACP_EMIT_CONTEXT_USAGE === "1";
 const promptResponseText = process.env.T3_ACP_PROMPT_RESPONSE_TEXT;
 const initialGrokReasoningEffort =
   process.env.T3_ACP_INITIAL_GROK_REASONING_EFFORT?.trim() || undefined;
@@ -1061,6 +1062,7 @@ const program = Effect.gen(function* () {
 
       yield* agent.client.sessionUpdate({
         sessionId: requestedSessionId,
+        ...(emitContextUsage ? { _meta: { totalTokens: 1615 } } : {}),
         update: {
           sessionUpdate: "agent_message_chunk",
           content: { type: "text", text: promptResponseText ?? "hello from mock" },

apps/server/src/provider/Layers/CursorAdapter.ts

@@ -57,6 +57,7 @@ import {
   makeAcpPlanUpdatedEvent,
   makeAcpRequestOpenedEvent,
   makeAcpRequestResolvedEvent,
+  makeAcpTokenUsageEvent,
   makeAcpToolCallEvent,
 } from "../acp/AcpCoreRuntimeEvents.ts";
 import {
@@ -867,6 +868,27 @@ export function makeCursorAdapter(
                       }),
                     );
                     return;
+                  case "UsageUpdated": {
+                    yield* logNative(
+                      ctx.threadId,
+                      "session/update",
+                      event.rawPayload,
+                      "acp.jsonrpc",
+                    );
+                    const usageEvent = makeAcpTokenUsageEvent({
+                      stamp: yield* makeEventStamp(),
+                      provider: PROVIDER,
+                      threadId: ctx.threadId,
+                      turnId: ctx.activeTurnId,
+                      used: event.used,
+                      size: event.size,
+                      rawPayload: event.rawPayload,
+                    });
+                    if (usageEvent) {
+                      yield* offerRuntimeEvent(usageEvent);
+                    }
+                    return;
+                  }
                 }
               }),
             ),

apps/server/src/provider/Layers/GrokAdapter.test.ts

@@ -27,10 +27,13 @@ import {
 
 import { ServerConfig } from "../../config.ts";
 import {
+  GROK_DEFAULT_CONTEXT_WINDOW,
+  grokContextWindowAfterModelChange,
   grokPromptSettlementBelongsToContext,
   isGrokEnterPlanModeToolCall,
   makeGrokAdapter,
   nextGrokPlanModeActive,
+  resolveGrokContextUsage,
   selectGrokPermissionOptionId,
 } from "./GrokAdapter.ts";
 const decodeGrokSettings = Schema.decodeSync(GrokSettings);
@@ -179,6 +182,83 @@ it("prefers allow_always when Grok offers it", () => {
   assert.equal(selectGrokPermissionOptionId(request, "accept"), "allow-once");
 });
 
+it("keeps occupancy snapshots when only the context window changes", () => {
+  assert.deepStrictEqual(
+    resolveGrokContextUsage({
+      used: 1615,
+      eventSize: 0,
+      lastUsed: 1615,
+      lastKnownMaxTokens: 500_000,
+      lastEmittedMaxTokens: 500_000,
+    }),
+    undefined,
+  );
+  assert.deepStrictEqual(
+    resolveGrokContextUsage({
+      used: 1615,
+      eventSize: 256_000,
+      lastUsed: 1615,
+      lastKnownMaxTokens: 500_000,
+      lastEmittedMaxTokens: 500_000,
+    }),
+    { used: 1615, size: 256_000 },
+  );
+  assert.deepStrictEqual(
+    resolveGrokContextUsage({
+      used: 1615,
+      eventSize: 0,
+      lastUsed: 1615,
+      lastKnownMaxTokens: 256_000,
+      lastEmittedMaxTokens: 500_000,
+    }),
+    { used: 1615, size: 256_000 },
+  );
+  assert.equal(
+    resolveGrokContextUsage({
+      used: 1_200,
+      eventSize: 0,
+      lastUsed: undefined,
+      lastKnownMaxTokens: undefined,
+      lastEmittedMaxTokens: undefined,
+    })?.size,
+    GROK_DEFAULT_CONTEXT_WINDOW,
+  );
+});
+
+it("refreshes the cached window when the Grok model changes", () => {
+  const windows = new Map([
+    ["grok-4.6", 500_000],
+    ["grok-code", 256_000],
+  ]);
+  assert.equal(
+    grokContextWindowAfterModelChange({
+      previousModelId: "grok-4.6",
+      nextModelId: "grok-4.6",
+      windows,
+      lastKnownMaxTokens: 500_000,
+    }),
+    500_000,
+  );
+  assert.equal(
+    grokContextWindowAfterModelChange({
+      previousModelId: "grok-4.6",
+      nextModelId: "grok-code",
+      windows,
+      lastKnownMaxTokens: 500_000,
+    }),
+    256_000,
+  );
+  assert.equal(
+    grokContextWindowAfterModelChange({
+      previousModelId: "grok-4.6",
+      nextModelId: "grok-unknown",
+      windows,
+      lastKnownMaxTokens: 500_000,
+    }),
+    undefined,
+  );
+});
+
 it("requires a settlement to match the live Grok turn", () => {
   const staleTurnId = TurnId.make("stale-turn");
   const replacementTurnId = TurnId.make("replacement-turn");
@@ -278,6 +358,56 @@ it.layer(grokAdapterTestLayer)("GrokAdapterLive", (it) => {
     }),
   );
 
+  it.effect("emits context occupancy from session-update _meta.totalTokens", () =>
+    Effect.gen(function* () {
+      const threadId = ThreadId.make("grok-context-window-thread");
+      const wrapperPath = yield* Effect.promise(() =>
+        makeMockGrokWrapper({ T3_ACP_EMIT_CONTEXT_USAGE: "1" }),
+      );
+      const adapter = yield* makeTestAdapter(wrapperPath);
+
+      const runtimeEvents: ProviderRuntimeEvent[] = [];
+      const turnCompleted = yield* Deferred.make<void>();
+      const runtimeEventsFiber = yield* Stream.runForEach(adapter.streamEvents, (event) =>
+        Effect.sync(() => {
+          runtimeEvents.push(event);
+        }).pipe(
+          Effect.andThen(
+            event.type === "turn.completed"
+              ? Deferred.succeed(turnCompleted, undefined)
+              : Effect.void,
+          ),
+        ),
+      ).pipe(Effect.forkChild);
+
+      yield* adapter.startSession({
+        threadId,
+        provider: ProviderDriverKind.make("grok"),
+        cwd: process.cwd(),
+        runtimeMode: "full-access",
+        modelSelection: { instanceId: ProviderInstanceId.make("grok"), model: "grok-4.6" },
+      });
+
+      yield* adapter.sendTurn({
+        threadId,
+        input: "hello grok",
+        attachments: [],
+      });
+
+      yield* Deferred.await(turnCompleted);
+      yield* Fiber.interrupt(runtimeEventsFiber);
+
+      const usageEvent = runtimeEvents.find((event) => event.type === "thread.token-usage.updated");
+      assert.equal(usageEvent?.type, "thread.token-usage.updated");
+      if (usageEvent?.type === "thread.token-usage.updated") {
+        assert.equal(usageEvent.payload.usage.usedTokens, 1615);
+        assert.equal(usageEvent.payload.usage.maxTokens, 500_000);
+      }
+
+      yield* adapter.stopSession(threadId);
+    }),
+  );
+
   it.effect("closes the ACP child process when a session stops", () =>
     Effect.gen(function* () {
       const threadId = ThreadId.make("grok-stop-session-close");

apps/server/src/provider/Layers/GrokAdapter.ts

@@ -55,12 +55,16 @@ import {
   makeAcpPlanUpdatedEvent,
   makeAcpRequestOpenedEvent,
   makeAcpRequestResolvedEvent,
+  makeAcpTokenUsageEvent,
   makeAcpToolCallEvent,
 } from "../acp/AcpCoreRuntimeEvents.ts";
 import { parsePermissionRequest } from "../acp/AcpRuntimeModel.ts";
 import { makeAcpNativeLoggerFactory } from "../acp/AcpNativeLogging.ts";
 import {
   applyGrokAcpModelSelection,
+  contextWindowForGrokModelId,
+  contextWindowFromGrokSessionSetup,
+  contextWindowsFromGrokSessionModels,
   currentGrokModelIdFromSessionSetup,
   currentGrokReasoningEffortFromSessionSetup,
   makeGrokAcpRuntime,
@@ -169,6 +173,10 @@ interface GrokSessionContext {
   promptResponsesReady: number;
   currentModelId: string | undefined;
   currentReasoningEffort: string | undefined;
+  contextWindows: ReadonlyMap<string, number>;
+  lastContextTokens: number | undefined;
+  lastKnownMaxTokens: number | undefined;
+  lastEmittedMaxTokens: number | undefined;
   stopped: boolean;
 }
 
@@ -314,6 +322,37 @@ function selectAutoApprovedPermissionOption(
   );
 }
 
+export const GROK_DEFAULT_CONTEXT_WINDOW = 500_000;
+
+export function resolveGrokContextUsage(input: {
+  readonly used: number;
+  readonly eventSize: number;
+  readonly lastUsed: number | undefined;
+  readonly lastKnownMaxTokens: number | undefined;
+  readonly lastEmittedMaxTokens: number | undefined;
+}): { readonly used: number; readonly size: number } | undefined {
+  const size =
+    input.eventSize > 0
+      ? input.eventSize
+      : (input.lastKnownMaxTokens ?? GROK_DEFAULT_CONTEXT_WINDOW);
+  if (input.lastUsed === input.used && input.lastEmittedMaxTokens === size) {
+    return undefined;
+  }
+  return { used: input.used, size };
+}
+
+export function grokContextWindowAfterModelChange(input: {
+  readonly previousModelId: string | undefined;
+  readonly nextModelId: string | undefined;
+  readonly windows: ReadonlyMap<string, number>;
+  readonly lastKnownMaxTokens: number | undefined;
+}): number | undefined {
+  if (input.nextModelId === input.previousModelId) {
+    return input.lastKnownMaxTokens;
+  }
+  return contextWindowForGrokModelId(input.windows, input.nextModelId);
+}
+
 function completedStopReasonFromPromptResponse(
   response: EffectAcpSchema.PromptResponse | undefined,
 ): EffectAcpSchema.StopReason | null {
@@ -1300,6 +1339,13 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte
               requestedStartReasoningEffort !== undefined
                 ? normalizeGrokReasoningEffort(requestedStartReasoningEffort)
                 : currentStartReasoningEffort,
+            contextWindows: contextWindowsFromGrokSessionModels(started.sessionSetupResult.models),
+            lastContextTokens: undefined,
+            lastKnownMaxTokens: contextWindowFromGrokSessionSetup(
+              started.sessionSetupResult,
+              boundModelId,
+            ),
+            lastEmittedMaxTokens: undefined,
             stopped: false,
           };
 
@@ -1313,7 +1359,8 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte
                 if (
                   event._tag === "PlanUpdated" ||
                   event._tag === "ToolCallUpdated" ||
-                  event._tag === "ContentDelta"
+                  event._tag === "ContentDelta" ||
+                  event._tag === "UsageUpdated"
                 ) {
                   yield* logNative(ctx.threadId, "session/update", event.rawPayload);
                 }
@@ -1322,6 +1369,41 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte
                   return;
                 }
 
+                if (event._tag === "UsageUpdated") {
+                  const usageTurnId = resolveNotificationTurnId(ctx);
+                  if (usageTurnId !== undefined && ctx.interruptedTurnIds.has(usageTurnId)) {
+                    return;
+                  }
+                  const snapshot = resolveGrokContextUsage({
+                    used: event.used,
+                    eventSize: event.size,
+                    lastUsed: ctx.lastContextTokens,
+                    lastKnownMaxTokens: ctx.lastKnownMaxTokens,
+                    lastEmittedMaxTokens: ctx.lastEmittedMaxTokens,
+                  });
+                  if (!snapshot) {
+                    return;
+                  }
+                  ctx.lastContextTokens = snapshot.used;
+                  if (event.size > 0) {
+                    ctx.lastKnownMaxTokens = event.size;
+                  }
+                  ctx.lastEmittedMaxTokens = snapshot.size;
+                  const usageEvent = makeAcpTokenUsageEvent({
+                    stamp: yield* makeEventStamp(),
+                    provider: PROVIDER,
+                    threadId: ctx.threadId,
+                    turnId: usageTurnId,
+                    used: snapshot.used,
+                    size: snapshot.size,
+                    rawPayload: event.rawPayload,
+                  });
+                  if (usageEvent) {
+                    yield* offerRuntimeEvent(usageEvent);
+                  }
+                  return;
+                }
+
                 const notificationTurnId = resolveNotificationTurnId(ctx);
                 if (
                   notificationTurnId === undefined ||
@@ -1574,6 +1656,12 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte
                 mapError: (cause) =>
                   mapAcpToAdapterError(PROVIDER, input.threadId, "session/set_model", cause),
               });
+              ctx.lastKnownMaxTokens = grokContextWindowAfterModelChange({
+                previousModelId: ctx.currentModelId,
+                nextModelId: currentModelId,
+                windows: ctx.contextWindows,
+                lastKnownMaxTokens: ctx.lastKnownMaxTokens,
+              });
               ctx.currentModelId = currentModelId;
               if (requestedTurnReasoningEffort !== undefined) {
                 ctx.currentReasoningEffort = normalizeGrokReasoningEffort(

apps/server/src/provider/acp/AcpCoreRuntimeEvents.test.ts

@@ -7,6 +7,7 @@ import {
   makeAcpPlanUpdatedEvent,
   makeAcpRequestOpenedEvent,
   makeAcpRequestResolvedEvent,
+  makeAcpTokenUsageEvent,
   makeAcpToolCallEvent,
 } from "./AcpCoreRuntimeEvents.ts";
 
@@ -191,5 +192,37 @@ describe("AcpCoreRuntimeEvents", () => {
         status: "inProgress",
       },
     });
+
+    expect(
+      makeAcpTokenUsageEvent({
+        stamp,
+        provider: ProviderDriverKind.make("grok"),
+        threadId: "thread-1" as never,
+        turnId,
+        used: 12_345,
+        size: 200_000,
+        rawPayload: { sessionUpdate: "usage_update" },
+      }),
+    ).toMatchObject({
+      type: "thread.token-usage.updated",
+      payload: {
+        usage: {
+          usedTokens: 12_345,
+          maxTokens: 200_000,
+        },
+      },
+    });
+
+    expect(
+      makeAcpTokenUsageEvent({
+        stamp,
+        provider: ProviderDriverKind.make("grok"),
+        threadId: "thread-1" as never,
+        turnId,
+        used: 0,
+        size: 200_000,
+        rawPayload: { sessionUpdate: "usage_update" },
+      }),
+    ).toBeNull();
   });
 });