diff --git a/packages/plugin/src/server/acp-internal/connection.ts b/packages/plugin/src/server/acp-internal/connection.ts index f853a6006a..6c8ac31b71 100644 --- a/packages/plugin/src/server/acp-internal/connection.ts +++ b/packages/plugin/src/server/acp-internal/connection.ts @@ -47,6 +47,7 @@ import { type ProviderSessionConfig, type ProviderSetting, type ProviderTimelineItem, + type ProviderUsage, } from "../provider.js"; const ADAPTER_CAPABILITIES = [ @@ -395,6 +396,7 @@ class AcpRuntime { private notificationLane: Promise = Promise.resolve(); private promptLane: Promise = Promise.resolve(); private activePrompt: { turnId: string; settled: Promise } | null = null; + private lastUsage: ProviderUsage = {}; static async start(options: StartRuntimeOptions): Promise { let child: ChildProcessWithoutNullStreams | null = null; @@ -623,6 +625,20 @@ class AcpRuntime { ).then( async (response): Promise => { await this.drainNotifications(); + if (response.usage) { + this.lastUsage = { + ...this.lastUsage, + inputTokens: response.usage.inputTokens, + outputTokens: response.usage.outputTokens, + cachedInputTokens: response.usage.cachedReadTokens ?? undefined, + }; + this.emit({ + type: "session.usage", + sessionId: this.options.boundarySessionId, + turnId, + usage: this.lastUsage, + }); + } const state = response.stopReason === "cancelled" ? "canceled" : "completed"; this.fallbackChunkIds.clear(); this.terminalizeTransientItems(state); @@ -1066,14 +1082,16 @@ class AcpRuntime { return; } if (update.sessionUpdate === "usage_update") { + this.lastUsage = { + ...this.lastUsage, + contextWindowUsedTokens: update.used, + contextWindowMaxTokens: update.size, + totalCostUsd: update.cost?.currency === "USD" ? update.cost.amount : undefined, + }; this.emit({ type: "session.usage", sessionId: this.options.boundarySessionId, - usage: { - contextWindowUsedTokens: update.used, - contextWindowMaxTokens: update.size, - totalCostUsd: update.cost?.currency === "USD" ? update.cost.amount : undefined, - }, + usage: this.lastUsage, }); return; } diff --git a/packages/plugin/src/server/acp.test.ts b/packages/plugin/src/server/acp.test.ts index 266d0b0109..b022e0a5ef 100644 --- a/packages/plugin/src/server/acp.test.ts +++ b/packages/plugin/src/server/acp.test.ts @@ -792,6 +792,104 @@ lines.on("line", (line) => { await connection.close(); }); + it.each([ + { counts: null, context: false, expected: [] }, + { + counts: { totalTokens: 0, inputTokens: 0, outputTokens: 0 }, + context: false, + expected: [ + { + turnId: "acp:first", + usage: { inputTokens: 0, outputTokens: 0, cachedInputTokens: undefined }, + }, + ], + }, + { + counts: { totalTokens: 18971, inputTokens: 12152, outputTokens: 78, cachedReadTokens: 6741 }, + context: true, + expected: [ + { + turnId: undefined, + usage: { contextWindowUsedTokens: 42, contextWindowMaxTokens: 100, totalCostUsd: 0.25 }, + }, + { + turnId: "acp:first", + usage: { + inputTokens: 12152, + outputTokens: 78, + cachedInputTokens: 6741, + contextWindowUsedTokens: 42, + contextWindowMaxTokens: 100, + totalCostUsd: 0.25, + }, + }, + ], + }, + ])( + "forwards prompt-result usage $counts while preserving context updates: $context", + async ({ counts, context, expected }) => { + const harness = connectorHarness({ + handleRequest(instance, message) { + if (message.method !== "session/prompt") return false; + if (context) + instance.notify("session/update", { + sessionId: "connector-session", + update: { + sessionUpdate: "usage_update", + used: 42, + size: 100, + cost: { amount: 0.25, currency: "USD" }, + }, + }); + instance.respond(message, { + stopReason: "end_turn", + ...(counts ? { usage: counts } : {}), + }); + return true; + }, + }); + const registration = runAcpProvider({ + id: "usage-acp", + label: "Usage ACP", + connector: harness.connector, + }); + const connection = await registration.connect({ + versions: [1], + capabilities: ["prompt.message"], + }); + const events: ProviderEvent[] = []; + connection.onEvent((event) => events.push(event)); + try { + await connection.send(openInput()); + await waitForEvent(events, (event) => event.type === "session.ready"); + await connection.send({ + type: "session.prompt", + sessionId: "session-1", + prompt: { + clientMessageId: "first", + delivery: "auto", + input: { type: "message", content: [{ type: "text", text: "hello" }] }, + }, + }); + await waitForEvent( + events, + (event) => event.type === "session.turn" && event.state === "completed", + ); + const usage = events.flatMap((event) => + event.type === "session.usage" ? [{ turnId: event.turnId, usage: event.usage }] : [], + ); + expect(usage).toEqual(expected); + const lastUsage = events.findLastIndex((event) => event.type === "session.usage"); + const completed = events.findIndex( + (event) => event.type === "session.turn" && event.state === "completed", + ); + expect(lastUsage).toBeLessThan(completed); + } finally { + await connection.close(); + } + }, + ); + it("serializes configuration mutations", async () => { let releaseMutation: (() => void) | null = null; const configOptions = [