diff --git a/apps/server/src/provider/Layers/CursorAdapter.ts b/apps/server/src/provider/Layers/CursorAdapter.ts index 80475a5c269..12f13009dde 100644 --- a/apps/server/src/provider/Layers/CursorAdapter.ts +++ b/apps/server/src/provider/Layers/CursorAdapter.ts @@ -867,6 +867,11 @@ export function makeCursorAdapter( }), ); return; + case "TokenUsageUpdated": + // Parsed from ACP usage_update; Grok wires this to the context + // meter. Cursor can emit the same makeAcpTokenUsageUpdatedEvent + // path once verified against a real Cursor ACP install. + return; } }), ), diff --git a/apps/server/src/provider/Layers/GrokAdapter.ts b/apps/server/src/provider/Layers/GrokAdapter.ts index 977cc8caadd..a606fe07456 100644 --- a/apps/server/src/provider/Layers/GrokAdapter.ts +++ b/apps/server/src/provider/Layers/GrokAdapter.ts @@ -49,15 +49,21 @@ import { makeAcpPlanUpdatedEvent, makeAcpRequestOpenedEvent, makeAcpRequestResolvedEvent, + makeAcpTokenUsageUpdatedEvent, makeAcpToolCallEvent, } from "../acp/AcpCoreRuntimeEvents.ts"; import { parsePermissionRequest } from "../acp/AcpRuntimeModel.ts"; import { makeAcpNativeLoggerFactory } from "../acp/AcpNativeLogging.ts"; import { applyGrokAcpModelSelection, + contextWindowForModelId, + contextWindowsFromSessionModels, currentGrokModelIdFromSessionSetup, + enrichGrokTokenUsage, makeGrokAcpRuntime, resolveGrokAcpBaseModelId, + resolveInitialGrokContextWindow, + tokenUsageFromGrokPromptMeta, } from "../acp/GrokAcpSupport.ts"; import { extractXAiAskUserQuestions, @@ -117,6 +123,10 @@ interface GrokSessionContext { * continues it, and only the last remaining prompt settles the turn. */ promptsInFlight: number; currentModelId: string | undefined; + /** modelId → totalContextTokens from Grok session model `_meta`. */ + contextWindowsByModelId: ReadonlyMap; + /** Last resolved window size so fill % still works if model id lookup misses. */ + lastKnownMaxTokens: number | undefined; stopped: boolean; } @@ -745,6 +755,15 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte mapError: (cause) => mapAcpToAdapterError(PROVIDER, input.threadId, "session/set_model", cause), }); + const setupModelId = currentGrokModelIdFromSessionSetup(started.sessionSetupResult); + const contextWindowsByModelId = contextWindowsFromSessionModels( + started.sessionSetupResult.models, + ); + const lastKnownMaxTokens = resolveInitialGrokContextWindow({ + windows: contextWindowsByModelId, + boundModelId, + setupModelId, + }); const now = yield* nowIso; const session: ProviderSession = { @@ -777,7 +796,9 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte activeTurnId: undefined, interruptedTurnIds: new Set(), promptsInFlight: 0, - currentModelId: boundModelId, + currentModelId: boundModelId ?? setupModelId, + contextWindowsByModelId, + lastKnownMaxTokens, stopped: false, }; @@ -791,7 +812,8 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte if ( event._tag === "PlanUpdated" || event._tag === "ToolCallUpdated" || - event._tag === "ContentDelta" + event._tag === "ContentDelta" || + event._tag === "TokenUsageUpdated" ) { yield* logNative(ctx.threadId, "session/update", event.rawPayload); } @@ -800,6 +822,33 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte return; } + // Session-level usage can arrive after activeTurnId is cleared; still + // surface it so the composer context meter can update between turns. + if (event._tag === "TokenUsageUpdated") { + const usageTurnId = resolveNotificationTurnId(ctx); + if (usageTurnId !== undefined && ctx.interruptedTurnIds.has(usageTurnId)) { + return; + } + const maxTokens = + contextWindowForModelId(ctx.contextWindowsByModelId, ctx.currentModelId) ?? + contextWindowForModelId(ctx.contextWindowsByModelId, ctx.session.model) ?? + ctx.lastKnownMaxTokens; + if (maxTokens !== undefined) { + ctx.lastKnownMaxTokens = maxTokens; + } + yield* offerRuntimeEvent( + makeAcpTokenUsageUpdatedEvent({ + stamp: yield* makeEventStamp(), + provider: PROVIDER, + threadId: ctx.threadId, + turnId: usageTurnId, + usage: enrichGrokTokenUsage(event.usage, maxTokens), + rawPayload: event.rawPayload, + }), + ); + return; + } + const notificationTurnId = resolveNotificationTurnId(ctx); if ( notificationTurnId === undefined || @@ -998,6 +1047,13 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte } ctx.currentModelId = currentModelId; + const turnMaxTokens = contextWindowForModelId( + ctx.contextWindowsByModelId, + currentModelId, + ); + if (turnMaxTokens !== undefined) { + ctx.lastKnownMaxTokens = turnMaxTokens; + } const displayModel = currentModelId ? resolveGrokAcpBaseModelId(currentModelId) : undefined; @@ -1150,6 +1206,28 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte updatedAt: yield* nowIso, ...(prepared.displayModel ? { model: prepared.displayModel } : {}), }; + // Grok streams fill via session/update _meta; prompt response + // carries the final totals as a reliable end-of-turn snapshot. + const promptMaxTokens = + contextWindowForModelId(ctx.contextWindowsByModelId, ctx.currentModelId) ?? + contextWindowForModelId(ctx.contextWindowsByModelId, ctx.session.model) ?? + ctx.lastKnownMaxTokens; + if (promptMaxTokens !== undefined) { + ctx.lastKnownMaxTokens = promptMaxTokens; + } + const promptUsage = tokenUsageFromGrokPromptMeta(result._meta, promptMaxTokens); + if (promptUsage) { + yield* offerRuntimeEvent( + makeAcpTokenUsageUpdatedEvent({ + stamp: yield* makeEventStamp(), + provider: PROVIDER, + threadId: ctx.threadId, + turnId: prepared.turnId, + usage: promptUsage, + rawPayload: result, + }), + ); + } const remainingPrompts = Math.max(0, ctx.promptsInFlight - 1); ctx.promptsInFlight = remainingPrompts; diff --git a/apps/server/src/provider/acp/AcpCoreRuntimeEvents.test.ts b/apps/server/src/provider/acp/AcpCoreRuntimeEvents.test.ts index 7fe25699bbc..79cd865f1f9 100644 --- a/apps/server/src/provider/acp/AcpCoreRuntimeEvents.test.ts +++ b/apps/server/src/provider/acp/AcpCoreRuntimeEvents.test.ts @@ -7,6 +7,7 @@ import { makeAcpPlanUpdatedEvent, makeAcpRequestOpenedEvent, makeAcpRequestResolvedEvent, + makeAcpTokenUsageUpdatedEvent, makeAcpToolCallEvent, } from "./AcpCoreRuntimeEvents.ts"; @@ -151,5 +152,35 @@ describe("AcpCoreRuntimeEvents", () => { status: "inProgress", }, }); + + expect( + makeAcpTokenUsageUpdatedEvent({ + stamp, + provider: ProviderDriverKind.make("grok"), + threadId: "thread-1" as never, + turnId, + usage: { + usedTokens: 12_345, + maxTokens: 200_000, + }, + rawPayload: { + sessionId: "session-1", + update: { sessionUpdate: "usage_update", used: 12_345, size: 200_000 }, + }, + }), + ).toMatchObject({ + type: "thread.token-usage.updated", + provider: "grok", + payload: { + usage: { + usedTokens: 12_345, + maxTokens: 200_000, + }, + }, + raw: { + source: "acp.jsonrpc", + method: "session/update", + }, + }); }); }); diff --git a/apps/server/src/provider/acp/AcpCoreRuntimeEvents.ts b/apps/server/src/provider/acp/AcpCoreRuntimeEvents.ts index c93e61dc37b..012db752162 100644 --- a/apps/server/src/provider/acp/AcpCoreRuntimeEvents.ts +++ b/apps/server/src/provider/acp/AcpCoreRuntimeEvents.ts @@ -8,6 +8,7 @@ import { type ProviderRuntimeEvent, type RuntimeRequestId, type ThreadId, + type ThreadTokenUsageSnapshot, type ToolLifecycleItemType, type TurnId, } from "@t3tools/contracts"; @@ -240,3 +241,28 @@ export function makeAcpContentDeltaEvent(input: { }, }; } + +export function makeAcpTokenUsageUpdatedEvent(input: { + readonly stamp: AcpEventStamp; + readonly provider: ProviderDriverKind; + readonly threadId: ThreadId; + readonly turnId: TurnId | undefined; + readonly usage: ThreadTokenUsageSnapshot; + readonly rawPayload: unknown; +}): ProviderRuntimeEvent { + return { + type: "thread.token-usage.updated", + ...input.stamp, + provider: input.provider, + threadId: input.threadId, + turnId: input.turnId, + payload: { + usage: input.usage, + }, + raw: { + source: "acp.jsonrpc", + method: "session/update", + payload: input.rawPayload, + }, + }; +} diff --git a/apps/server/src/provider/acp/AcpRuntimeModel.test.ts b/apps/server/src/provider/acp/AcpRuntimeModel.test.ts index 7682c5f5f9c..21c0499835e 100644 --- a/apps/server/src/provider/acp/AcpRuntimeModel.test.ts +++ b/apps/server/src/provider/acp/AcpRuntimeModel.test.ts @@ -336,6 +336,120 @@ describe("AcpRuntimeModel", () => { ]); }); + it("projects ACP usage_update into TokenUsageUpdated snapshots", () => { + const raw = { + sessionId: "session-1", + update: { + sessionUpdate: "usage_update" as const, + used: 12_345, + size: 200_000, + }, + } satisfies EffectAcpSchema.SessionNotification; + + expect(parseSessionUpdateEvent(raw).events).toEqual([ + { + _tag: "TokenUsageUpdated", + usage: { + usedTokens: 12_345, + maxTokens: 200_000, + }, + rawPayload: raw, + }, + ]); + + expect( + parseSessionUpdateEvent({ + sessionId: "session-1", + update: { + sessionUpdate: "usage_update", + used: 0, + size: 200_000, + }, + } satisfies EffectAcpSchema.SessionNotification).events, + ).toEqual([]); + + expect( + parseSessionUpdateEvent({ + sessionId: "session-1", + update: { + sessionUpdate: "usage_update", + used: 99, + size: 0, + }, + } satisfies EffectAcpSchema.SessionNotification).events, + ).toEqual([ + { + _tag: "TokenUsageUpdated", + usage: { + usedTokens: 99, + }, + rawPayload: { + sessionId: "session-1", + update: { + sessionUpdate: "usage_update", + used: 99, + size: 0, + }, + }, + }, + ]); + }); + + it("projects Grok-style session/update _meta.totalTokens into TokenUsageUpdated", () => { + const raw = { + sessionId: "session-1", + update: { + sessionUpdate: "agent_message_chunk", + content: { + type: "text", + text: "hello", + }, + }, + _meta: { + totalTokens: 19_267, + }, + } satisfies EffectAcpSchema.SessionNotification; + + expect(parseSessionUpdateEvent(raw).events).toEqual([ + { + _tag: "ContentDelta", + text: "hello", + rawPayload: raw, + }, + { + _tag: "TokenUsageUpdated", + usage: { + usedTokens: 19_267, + }, + rawPayload: raw, + }, + ]); + + // Explicit usage_update wins over notification _meta when both exist. + const withBoth = { + sessionId: "session-1", + update: { + sessionUpdate: "usage_update" as const, + used: 1_000, + size: 50_000, + }, + _meta: { + totalTokens: 9_999, + }, + } satisfies EffectAcpSchema.SessionNotification; + + expect(parseSessionUpdateEvent(withBoth).events).toEqual([ + { + _tag: "TokenUsageUpdated", + usage: { + usedTokens: 1_000, + maxTokens: 50_000, + }, + rawPayload: withBoth, + }, + ]); + }); + it("keeps permission request parsing compatible with loose extension payloads", () => { const request = parsePermissionRequest({ sessionId: "session-1", diff --git a/apps/server/src/provider/acp/AcpRuntimeModel.ts b/apps/server/src/provider/acp/AcpRuntimeModel.ts index e6bfc127e6e..a577d86759b 100644 --- a/apps/server/src/provider/acp/AcpRuntimeModel.ts +++ b/apps/server/src/provider/acp/AcpRuntimeModel.ts @@ -5,7 +5,7 @@ import * as Option from "effect/Option"; import * as Ref from "effect/Ref"; import type * as EffectAcpSchema from "effect-acp/schema"; import { deriveToolActivityPresentation } from "@t3tools/shared/toolActivity"; -import type { ToolLifecycleItemType } from "@t3tools/contracts"; +import type { ThreadTokenUsageSnapshot, ToolLifecycleItemType } from "@t3tools/contracts"; function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); @@ -108,6 +108,11 @@ export type AcpParsedSessionEvent = readonly itemId?: string; readonly text: string; readonly rawPayload: unknown; + } + | { + readonly _tag: "TokenUsageUpdated"; + readonly usage: ThreadTokenUsageSnapshot; + readonly rawPayload: unknown; }; type AcpSessionSetupResponse = @@ -574,9 +579,84 @@ export function parseSessionUpdateEvent(params: EffectAcpSchema.SessionNotificat } break; } + case "usage_update": { + // ACP Session Context Size and Cost: used = tokens in context, size = window. + // Ingestion only materializes context-window activities when usedTokens > 0. + const usage = tokenUsageFromAcpUsageUpdate(upd); + if (usage) { + events.push({ + _tag: "TokenUsageUpdated", + usage, + rawPayload: params, + }); + } + break; + } default: break; } + // Prefer explicit usage_update above. Grok (and similar agents) also put the + // current context fill on SessionNotification._meta.totalTokens while streaming. + if (!events.some((event) => event._tag === "TokenUsageUpdated")) { + const usageFromMeta = tokenUsageFromSessionNotificationMeta(params._meta); + if (usageFromMeta) { + events.push({ + _tag: "TokenUsageUpdated", + usage: usageFromMeta, + rawPayload: params, + }); + } + } + return { ...(modeId !== undefined ? { modeId } : {}), events }; } + +function finiteNonNegativeInt(value: unknown): number | undefined { + if (typeof value !== "number" || !Number.isFinite(value)) { + return undefined; + } + const truncated = Math.trunc(value); + return truncated >= 0 ? truncated : undefined; +} + +function finitePositiveInt(value: unknown): number | undefined { + const valueInt = finiteNonNegativeInt(value); + return valueInt !== undefined && valueInt > 0 ? valueInt : undefined; +} + +function tokenUsageFromAcpUsageUpdate(update: { + readonly used: number; + readonly size: number; +}): ThreadTokenUsageSnapshot | undefined { + const usedTokens = finitePositiveInt(update.used); + if (usedTokens === undefined) { + return undefined; + } + const maxTokens = finitePositiveInt(update.size); + return { + usedTokens, + ...(maxTokens !== undefined ? { maxTokens } : {}), + }; +} + +/** Read context fill from ACP notification `_meta` (Grok streams this today). */ +export function tokenUsageFromSessionNotificationMeta( + meta: unknown, +): ThreadTokenUsageSnapshot | undefined { + if (!isRecord(meta)) { + return undefined; + } + const usedTokens = finitePositiveInt(meta.totalTokens); + if (usedTokens === undefined) { + return undefined; + } + const maxTokens = + finitePositiveInt(meta.totalContextTokens) ?? + finitePositiveInt(meta.contextWindow) ?? + finitePositiveInt(meta.size); + return { + usedTokens, + ...(maxTokens !== undefined ? { maxTokens } : {}), + }; +} diff --git a/apps/server/src/provider/acp/GrokAcpSupport.test.ts b/apps/server/src/provider/acp/GrokAcpSupport.test.ts index 02d60976b24..35085f5409b 100644 --- a/apps/server/src/provider/acp/GrokAcpSupport.test.ts +++ b/apps/server/src/provider/acp/GrokAcpSupport.test.ts @@ -5,7 +5,13 @@ import * as EffectAcpErrors from "effect-acp/errors"; import { applyGrokAcpModelSelection, buildGrokAcpSpawnInput, + contextWindowForModelId, + contextWindowsFromSessionModels, + enrichGrokTokenUsage, resolveGrokAcpBaseModelId, + resolveInitialGrokContextWindow, + tokenUsageFromGrokPromptMeta, + totalContextTokensFromModelMeta, } from "./GrokAcpSupport.ts"; describe("resolveGrokAcpBaseModelId", () => { @@ -107,3 +113,163 @@ describe("applyGrokAcpModelSelection", () => { }), ); }); + +describe("Grok context window helpers", () => { + it("reads totalContextTokens from model meta and session model state", () => { + expect(totalContextTokensFromModelMeta({ totalContextTokens: 500_000 })).toBe(500_000); + expect(totalContextTokensFromModelMeta({ totalContextTokens: 0 })).toBeUndefined(); + expect(totalContextTokensFromModelMeta(null)).toBeUndefined(); + + const windows = contextWindowsFromSessionModels({ + currentModelId: "model-a", + availableModels: [ + { + modelId: "model-a", + name: "Model A", + _meta: { totalContextTokens: 500_000 }, + }, + { + modelId: "model-b", + name: "Model B", + _meta: { totalContextTokens: 128_000 }, + }, + { + modelId: "model-c", + name: "Model C", + }, + ], + }); + + expect(contextWindowForModelId(windows, "model-a")).toBe(500_000); + expect(contextWindowForModelId(windows, "model-b")).toBe(128_000); + expect(contextWindowForModelId(windows, "model-c")).toBeUndefined(); + }); + + it("maps prompt _meta usage into a ThreadTokenUsageSnapshot with model window", () => { + expect( + tokenUsageFromGrokPromptMeta( + { + totalTokens: 19_267, + inputTokens: 19_237, + outputTokens: 29, + cachedReadTokens: 2_560, + reasoningTokens: 18, + usage: { + inputTokens: 19_237, + outputTokens: 29, + totalTokens: 19_266, + cachedReadTokens: 2_560, + reasoningTokens: 18, + }, + }, + 500_000, + ), + ).toEqual({ + usedTokens: 19_267, + maxTokens: 500_000, + inputTokens: 19_237, + lastInputTokens: 19_237, + outputTokens: 29, + lastOutputTokens: 29, + cachedInputTokens: 2_560, + lastCachedInputTokens: 2_560, + reasoningOutputTokens: 18, + lastReasoningOutputTokens: 18, + lastUsedTokens: 19_267, + compactsAutomatically: true, + }); + + expect(tokenUsageFromGrokPromptMeta({}, 500_000)).toBeUndefined(); + expect(enrichGrokTokenUsage({ usedTokens: 100 }, 200_000)).toEqual({ + usedTokens: 100, + maxTokens: 200_000, + compactsAutomatically: true, + }); + // Always stamps max + auto-compact so the meter can show used/max. + expect(enrichGrokTokenUsage({ usedTokens: 16_000 }, undefined)).toEqual({ + usedTokens: 16_000, + compactsAutomatically: true, + }); + }); + + it("resolves an initial window even when the bound model id is missing", () => { + const windows = contextWindowsFromSessionModels({ + currentModelId: "model-a", + availableModels: [ + { + modelId: "model-a", + name: "Model A", + _meta: { totalContextTokens: 500_000 }, + }, + { + modelId: "model-b", + name: "Model B", + _meta: { totalContextTokens: 200_000 }, + }, + ], + }); + + expect( + resolveInitialGrokContextWindow({ + windows, + boundModelId: undefined, + setupModelId: "model-a", + }), + ).toBe(500_000); + + // Multiple known windows + no matching model id → do not guess. + expect( + resolveInitialGrokContextWindow({ + windows, + boundModelId: "missing-model", + setupModelId: "also-missing", + }), + ).toBeUndefined(); + + const singleWindow = contextWindowsFromSessionModels({ + currentModelId: "model-a", + availableModels: [ + { + modelId: "model-a", + name: "Model A", + _meta: { totalContextTokens: 500_000 }, + }, + ], + }); + expect( + resolveInitialGrokContextWindow({ + windows: singleWindow, + boundModelId: "missing-model", + setupModelId: "also-missing", + }), + ).toBe(500_000); + }); + + it("preserves zero-valued token breakdowns from prompt meta", () => { + expect( + tokenUsageFromGrokPromptMeta( + { + totalTokens: 100, + inputTokens: 100, + outputTokens: 0, + cachedReadTokens: 0, + reasoningTokens: 0, + }, + 50_000, + ), + ).toEqual({ + usedTokens: 100, + maxTokens: 50_000, + inputTokens: 100, + lastInputTokens: 100, + outputTokens: 0, + lastOutputTokens: 0, + cachedInputTokens: 0, + lastCachedInputTokens: 0, + reasoningOutputTokens: 0, + lastReasoningOutputTokens: 0, + lastUsedTokens: 100, + compactsAutomatically: true, + }); + }); +}); diff --git a/apps/server/src/provider/acp/GrokAcpSupport.ts b/apps/server/src/provider/acp/GrokAcpSupport.ts index c928b3ed80e..26bb18791de 100644 --- a/apps/server/src/provider/acp/GrokAcpSupport.ts +++ b/apps/server/src/provider/acp/GrokAcpSupport.ts @@ -1,4 +1,8 @@ -import { type GrokSettings, ProviderDriverKind } from "@t3tools/contracts"; +import { + type GrokSettings, + ProviderDriverKind, + type ThreadTokenUsageSnapshot, +} from "@t3tools/contracts"; import * as Crypto from "effect/Crypto"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; @@ -106,3 +110,150 @@ export function applyGrokAcpModelSelection(input: { .setSessionModel(input.requestedModelId) .pipe(Effect.mapError(input.mapError), Effect.as(input.requestedModelId)); } + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function finiteNonNegativeInt(value: unknown): number | undefined { + if (typeof value !== "number" || !Number.isFinite(value)) { + return undefined; + } + const truncated = Math.trunc(value); + return truncated >= 0 ? truncated : undefined; +} + +function finitePositiveInt(value: unknown): number | undefined { + const valueInt = finiteNonNegativeInt(value); + return valueInt !== undefined && valueInt > 0 ? valueInt : undefined; +} + +/** Model context window from Grok ACP `availableModels[]._meta.totalContextTokens`. */ +export function totalContextTokensFromModelMeta(meta: unknown): number | undefined { + if (!isRecord(meta)) { + return undefined; + } + return finitePositiveInt(meta.totalContextTokens); +} + +/** Build modelId → context-window map from a Grok session model state. */ +export function contextWindowsFromSessionModels( + models: EffectAcpSchema.SessionModelState | null | undefined, +): ReadonlyMap { + const windows = new Map(); + for (const model of models?.availableModels ?? []) { + const size = totalContextTokensFromModelMeta(model._meta); + if (size === undefined) { + continue; + } + windows.set(model.modelId, size); + // Index the normalized slug too so UI slugs and ACP ids both resolve. + const baseId = resolveGrokAcpBaseModelId(model.modelId); + if (baseId && baseId !== model.modelId) { + windows.set(baseId, size); + } + } + return windows; +} + +export function contextWindowForModelId( + windows: ReadonlyMap | undefined, + modelId: string | undefined, +): number | undefined { + if (!windows || windows.size === 0) { + return undefined; + } + if (modelId) { + const direct = windows.get(modelId) ?? windows.get(resolveGrokAcpBaseModelId(modelId)); + if (direct !== undefined) { + return direct; + } + } + // Prefer the sole known window when the active model id does not match. + if (windows.size === 1) { + return windows.values().next().value; + } + return undefined; +} + +/** Best-effort window for a new session: bound model, then setup current, then sole entry. */ +export function resolveInitialGrokContextWindow(input: { + readonly windows: ReadonlyMap; + readonly boundModelId: string | undefined; + readonly setupModelId: string | undefined; +}): number | undefined { + return ( + contextWindowForModelId(input.windows, input.boundModelId) ?? + contextWindowForModelId(input.windows, input.setupModelId) ?? + // Only fall back when a single window is known — never pick an arbitrary model. + (input.windows.size === 1 ? input.windows.values().next().value : undefined) + ); +} + +/** + * Enrich a token-usage snapshot with the current model context window and + * Grok's auto-compact behavior so the composer meter can show fill %. + */ +export function enrichGrokTokenUsage( + usage: ThreadTokenUsageSnapshot, + maxTokens: number | undefined, +): ThreadTokenUsageSnapshot { + const resolvedMax = usage.maxTokens ?? maxTokens; + return { + ...usage, + ...(resolvedMax !== undefined ? { maxTokens: resolvedMax } : {}), + compactsAutomatically: usage.compactsAutomatically ?? true, + }; +} + +/** + * Grok does not yet emit ACP `usage_update`, but prompt responses carry + * `_meta.totalTokens` / `_meta.usage` for the session context fill. + */ +export function tokenUsageFromGrokPromptMeta( + meta: unknown, + maxTokens: number | undefined, +): ThreadTokenUsageSnapshot | undefined { + if (!isRecord(meta)) { + return undefined; + } + + const usageRecord = isRecord(meta.usage) ? meta.usage : undefined; + const usedTokens = + finitePositiveInt(meta.totalTokens) ?? finitePositiveInt(usageRecord?.totalTokens); + if (usedTokens === undefined) { + return undefined; + } + + // Breakdown fields may legitimately be 0; usedTokens/window stay strictly positive. + const inputTokens = + finiteNonNegativeInt(meta.inputTokens) ?? finiteNonNegativeInt(usageRecord?.inputTokens); + const outputTokens = + finiteNonNegativeInt(meta.outputTokens) ?? finiteNonNegativeInt(usageRecord?.outputTokens); + const cachedInputTokens = + finiteNonNegativeInt(meta.cachedReadTokens) ?? + finiteNonNegativeInt(usageRecord?.cachedReadTokens); + const reasoningOutputTokens = + finiteNonNegativeInt(meta.reasoningTokens) ?? + finiteNonNegativeInt(usageRecord?.reasoningTokens); + const window = + maxTokens ?? + finitePositiveInt(meta.totalContextTokens) ?? + finitePositiveInt(usageRecord?.totalContextTokens); + + return enrichGrokTokenUsage( + { + usedTokens, + ...(inputTokens !== undefined ? { inputTokens, lastInputTokens: inputTokens } : {}), + ...(outputTokens !== undefined ? { outputTokens, lastOutputTokens: outputTokens } : {}), + ...(cachedInputTokens !== undefined + ? { cachedInputTokens, lastCachedInputTokens: cachedInputTokens } + : {}), + ...(reasoningOutputTokens !== undefined + ? { reasoningOutputTokens, lastReasoningOutputTokens: reasoningOutputTokens } + : {}), + lastUsedTokens: usedTokens, + }, + window, + ); +} diff --git a/apps/web/src/lib/contextWindow.ts b/apps/web/src/lib/contextWindow.ts index 80f7d31cf2f..424d895b0dd 100644 --- a/apps/web/src/lib/contextWindow.ts +++ b/apps/web/src/lib/contextWindow.ts @@ -36,6 +36,8 @@ export function formatProviderDisplayName(provider: string | null | undefined): return "Codex"; case "cursor": return "Cursor"; + case "grok": + return "Grok"; case "opencode": return "OpenCode"; default: {