From 5fdae9e90b6aab3985ae4ac52a795402eb48d825 Mon Sep 17 00:00:00 2001 From: Oleksandr Zabolotnyi Date: Mon, 31 Aug 2026 16:08:52 +0200 Subject: [PATCH] fix(server): show Cursor, Grok, and Jcode token usage in chat ACP agents already send usage_update and prompt usage, but those events were dropped. Map them into thread.token-usage.updated so the chat chip and context meter work for those providers too. Co-authored-by: Cursor --- .../src/provider/Layers/CursorAdapter.ts | 13 +++ .../server/src/provider/Layers/GrokAdapter.ts | 13 +++ .../src/provider/Layers/JcodeAdapter.ts | 13 +++ .../src/provider/acp/AcpCoreRuntimeEvents.ts | 26 ++++++ .../src/provider/acp/AcpRuntimeModel.ts | 16 +++- .../src/provider/acp/AcpSessionRuntime.ts | 5 ++ .../src/provider/acp/acpTokenUsage.test.ts | 80 +++++++++++++++++ apps/server/src/provider/acp/acpTokenUsage.ts | 90 +++++++++++++++++++ 8 files changed, 254 insertions(+), 2 deletions(-) create mode 100644 apps/server/src/provider/acp/acpTokenUsage.test.ts create mode 100644 apps/server/src/provider/acp/acpTokenUsage.ts diff --git a/apps/server/src/provider/Layers/CursorAdapter.ts b/apps/server/src/provider/Layers/CursorAdapter.ts index 818fe567b9b3..2097e18e8f50 100644 --- a/apps/server/src/provider/Layers/CursorAdapter.ts +++ b/apps/server/src/provider/Layers/CursorAdapter.ts @@ -54,6 +54,7 @@ import type * as AcpSessionRuntime from "../acp/AcpSessionRuntime.ts"; import { makeAcpAssistantItemEvent, makeAcpContentDeltaEvent, + makeAcpTokenUsageEvent, makeAcpPlanUpdatedEvent, makeAcpRequestOpenedEvent, makeAcpRequestResolvedEvent, @@ -867,6 +868,18 @@ export function makeCursorAdapter( }), ); return; + case "TokenUsageUpdated": + yield* offerRuntimeEvent( + makeAcpTokenUsageEvent({ + stamp: yield* makeEventStamp(), + provider: PROVIDER, + threadId: ctx.threadId, + turnId: ctx.activeTurnId, + usage: event.usage, + rawPayload: event.rawPayload, + }), + ); + return; } }), ), diff --git a/apps/server/src/provider/Layers/GrokAdapter.ts b/apps/server/src/provider/Layers/GrokAdapter.ts index 4bec298c628d..ed6dc6a27955 100644 --- a/apps/server/src/provider/Layers/GrokAdapter.ts +++ b/apps/server/src/provider/Layers/GrokAdapter.ts @@ -52,6 +52,7 @@ import type * as AcpSessionRuntime from "../acp/AcpSessionRuntime.ts"; import { makeAcpAssistantItemEvent, makeAcpContentDeltaEvent, + makeAcpTokenUsageEvent, makeAcpPlanUpdatedEvent, makeAcpRequestOpenedEvent, makeAcpRequestResolvedEvent, @@ -1413,6 +1414,18 @@ export function makeGrokAdapter(grokSettings: GrokSettings, options?: GrokAdapte }), ); return; + case "TokenUsageUpdated": + yield* offerRuntimeEvent( + makeAcpTokenUsageEvent({ + stamp, + provider: PROVIDER, + threadId: ctx.threadId, + turnId: notificationTurnId, + usage: event.usage, + rawPayload: event.rawPayload, + }), + ); + return; } }), ), diff --git a/apps/server/src/provider/Layers/JcodeAdapter.ts b/apps/server/src/provider/Layers/JcodeAdapter.ts index 0d467e00e15f..e0ba6cf91d1a 100644 --- a/apps/server/src/provider/Layers/JcodeAdapter.ts +++ b/apps/server/src/provider/Layers/JcodeAdapter.ts @@ -48,6 +48,7 @@ import type * as AcpSessionRuntime from "../acp/AcpSessionRuntime.ts"; import { makeAcpAssistantItemEvent, makeAcpContentDeltaEvent, + makeAcpTokenUsageEvent, makeAcpPlanUpdatedEvent, makeAcpRequestOpenedEvent, makeAcpRequestResolvedEvent, @@ -924,6 +925,18 @@ export function makeJcodeAdapter(jcodeSettings: JcodeSettings, options?: JcodeAd }), ); return; + case "TokenUsageUpdated": + yield* offerRuntimeEvent( + makeAcpTokenUsageEvent({ + stamp, + provider: PROVIDER, + threadId: ctx.threadId, + turnId: notificationTurnId, + usage: event.usage, + rawPayload: event.rawPayload, + }), + ); + return; } }), ), diff --git a/apps/server/src/provider/acp/AcpCoreRuntimeEvents.ts b/apps/server/src/provider/acp/AcpCoreRuntimeEvents.ts index bd25e9815aef..7c2720e80242 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 makeAcpTokenUsageEvent(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.ts b/apps/server/src/provider/acp/AcpRuntimeModel.ts index 8fd7af3c5610..afd0eba85622 100644 --- a/apps/server/src/provider/acp/AcpRuntimeModel.ts +++ b/apps/server/src/provider/acp/AcpRuntimeModel.ts @@ -5,7 +5,9 @@ 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"; + +import { tokenUsageEventFromAcpSessionUpdate } from "./acpTokenUsage.ts"; function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); @@ -108,6 +110,11 @@ export type AcpParsedSessionEvent = readonly itemId?: string; readonly text: string; readonly rawPayload: unknown; + } + | { + readonly _tag: "TokenUsageUpdated"; + readonly usage: ThreadTokenUsageSnapshot; + readonly rawPayload: unknown; }; type AcpSessionSetupResponse = @@ -823,8 +830,13 @@ export function parseSessionUpdateEvent(params: EffectAcpSchema.SessionNotificat } break; } - default: + default: { + const tokenUsage = tokenUsageEventFromAcpSessionUpdate(params); + if (tokenUsage) { + events.push(tokenUsage); + } break; + } } return { ...(modeId !== undefined ? { modeId } : {}), events }; diff --git a/apps/server/src/provider/acp/AcpSessionRuntime.ts b/apps/server/src/provider/acp/AcpSessionRuntime.ts index 262b1a2c4642..6f2a58bd3c3b 100644 --- a/apps/server/src/provider/acp/AcpSessionRuntime.ts +++ b/apps/server/src/provider/acp/AcpSessionRuntime.ts @@ -21,6 +21,7 @@ import type * as EffectAcpSchema from "effect-acp/schema"; import type * as EffectAcpProtocol from "effect-acp/protocol"; import { resolveSpawnCommand } from "@t3tools/shared/shell"; +import { tokenUsageEventFromAcpPromptResponse } from "./acpTokenUsage.ts"; import { collectSessionConfigOptionValues, decideToolCallUpdateEmission, @@ -772,6 +773,10 @@ export const make = ( assistantSegmentRef, }), ), + Effect.tap((result) => { + const event = tokenUsageEventFromAcpPromptResponse(result); + return event ? Queue.offer(eventQueue, event) : Effect.void; + }), ); }), ), diff --git a/apps/server/src/provider/acp/acpTokenUsage.test.ts b/apps/server/src/provider/acp/acpTokenUsage.test.ts new file mode 100644 index 000000000000..12595de93f72 --- /dev/null +++ b/apps/server/src/provider/acp/acpTokenUsage.test.ts @@ -0,0 +1,80 @@ +import { describe, expect, it } from "vite-plus/test"; + +import { parseSessionUpdateEvent } from "./AcpRuntimeModel.ts"; +import { + snapshotFromAcpPromptUsage, + snapshotFromAcpUsageUpdate, + tokenUsageEventFromAcpPromptResponse, + tokenUsageEventFromAcpSessionUpdate, +} from "./acpTokenUsage.ts"; + +describe("acpTokenUsage", () => { + it("maps usage_update into a context-window snapshot the chat can render", () => { + expect(snapshotFromAcpUsageUpdate({ used: 12_400, size: 200_000 })).toEqual({ + usedTokens: 12_400, + lastUsedTokens: 12_400, + maxTokens: 200_000, + }); + expect( + tokenUsageEventFromAcpSessionUpdate({ + sessionId: "session-1", + update: { sessionUpdate: "usage_update", used: 12_400, size: 200_000 }, + })?._tag, + ).toBe("TokenUsageUpdated"); + expect( + parseSessionUpdateEvent({ + sessionId: "session-1", + update: { sessionUpdate: "usage_update", used: 12_400, size: 200_000 }, + }).events, + ).toEqual([ + { + _tag: "TokenUsageUpdated", + usage: { usedTokens: 12_400, lastUsedTokens: 12_400, maxTokens: 200_000 }, + rawPayload: { + sessionId: "session-1", + update: { sessionUpdate: "usage_update", used: 12_400, size: 200_000 }, + }, + }, + ]); + }); + + it("maps prompt usage into last-turn token counts", () => { + expect( + snapshotFromAcpPromptUsage({ + inputTokens: 800, + outputTokens: 200, + totalTokens: 1_050, + cachedReadTokens: 50, + thoughtTokens: 0, + }), + ).toEqual({ + usedTokens: 1_050, + lastUsedTokens: 1_050, + inputTokens: 800, + lastInputTokens: 800, + outputTokens: 200, + lastOutputTokens: 200, + cachedInputTokens: 50, + lastCachedInputTokens: 50, + }); + expect( + tokenUsageEventFromAcpPromptResponse({ + stopReason: "end_turn", + usage: { inputTokens: 10, outputTokens: 5, totalTokens: 15 }, + })?.usage.lastUsedTokens, + ).toBe(15); + }); + + it("ignores empty usage payloads", () => { + expect(snapshotFromAcpUsageUpdate({ used: 0, size: 200_000 })).toBeUndefined(); + expect( + snapshotFromAcpPromptUsage({ inputTokens: 0, outputTokens: 0, totalTokens: 0 }), + ).toBeUndefined(); + expect( + tokenUsageEventFromAcpSessionUpdate({ + sessionId: "session-1", + update: { sessionUpdate: "agent_message_chunk", content: { type: "text", text: "hi" } }, + }), + ).toBeUndefined(); + }); +}); diff --git a/apps/server/src/provider/acp/acpTokenUsage.ts b/apps/server/src/provider/acp/acpTokenUsage.ts new file mode 100644 index 000000000000..b1b3899d7606 --- /dev/null +++ b/apps/server/src/provider/acp/acpTokenUsage.ts @@ -0,0 +1,90 @@ +import type { ThreadTokenUsageSnapshot } from "@t3tools/contracts"; +import type * as EffectAcpSchema from "effect-acp/schema"; + +export type AcpTokenUsageSessionEvent = { + readonly _tag: "TokenUsageUpdated"; + readonly usage: ThreadTokenUsageSnapshot; + readonly rawPayload: unknown; +}; + +function asCount(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value >= 0 + ? Math.trunc(value) + : undefined; +} + +export function snapshotFromAcpUsageUpdate(update: { + readonly used: number; + readonly size: number; +}): ThreadTokenUsageSnapshot | undefined { + const usedTokens = asCount(update.used); + if (usedTokens === undefined || usedTokens <= 0) { + return undefined; + } + const maxTokens = asCount(update.size); + return { + usedTokens, + lastUsedTokens: usedTokens, + ...(maxTokens !== undefined && maxTokens > 0 ? { maxTokens } : {}), + }; +} + +export function snapshotFromAcpPromptUsage( + usage: EffectAcpSchema.Usage | null | undefined, +): ThreadTokenUsageSnapshot | undefined { + if (!usage) { + return undefined; + } + const inputTokens = asCount(usage.inputTokens); + const outputTokens = asCount(usage.outputTokens); + const totalTokens = asCount(usage.totalTokens); + const lastUsedTokens = + totalTokens !== undefined && totalTokens > 0 + ? totalTokens + : (inputTokens ?? 0) + (outputTokens ?? 0); + if (lastUsedTokens <= 0) { + return undefined; + } + const cachedInputTokens = asCount(usage.cachedReadTokens ?? undefined); + const reasoningOutputTokens = asCount(usage.thoughtTokens ?? undefined); + return { + usedTokens: lastUsedTokens, + lastUsedTokens, + ...(inputTokens !== undefined && inputTokens > 0 + ? { inputTokens, lastInputTokens: inputTokens } + : {}), + ...(outputTokens !== undefined && outputTokens > 0 + ? { outputTokens, lastOutputTokens: outputTokens } + : {}), + ...(cachedInputTokens !== undefined && cachedInputTokens > 0 + ? { cachedInputTokens, lastCachedInputTokens: cachedInputTokens } + : {}), + ...(reasoningOutputTokens !== undefined && reasoningOutputTokens > 0 + ? { reasoningOutputTokens, lastReasoningOutputTokens: reasoningOutputTokens } + : {}), + }; +} + +export function tokenUsageEventFromAcpSessionUpdate( + params: EffectAcpSchema.SessionNotification, +): AcpTokenUsageSessionEvent | undefined { + const update = params.update; + if (update.sessionUpdate !== "usage_update") { + return undefined; + } + const usage = snapshotFromAcpUsageUpdate(update); + if (!usage) { + return undefined; + } + return { _tag: "TokenUsageUpdated", usage, rawPayload: params }; +} + +export function tokenUsageEventFromAcpPromptResponse( + result: EffectAcpSchema.PromptResponse, +): AcpTokenUsageSessionEvent | undefined { + const usage = snapshotFromAcpPromptUsage(result.usage); + if (!usage) { + return undefined; + } + return { _tag: "TokenUsageUpdated", usage, rawPayload: result }; +}