diff --git a/packages/cli/src/acp/event.ts b/packages/cli/src/acp/event.ts index a3fa95d5d2b..6785fc75d74 100644 --- a/packages/cli/src/acp/event.ts +++ b/packages/cli/src/acp/event.ts @@ -5,7 +5,9 @@ import type { SessionMessageAssistant, SessionMessageInfo, SessionStructuredError, + TokenUsageInfo, } from "@opencode/client/promise" +import { TokenUsage } from "@opencode/schema/token-usage" import type { ACPConnection } from "./connection" import { partsToContentChunks, type ReplayPart } from "./content" import { ACPError } from "./error" @@ -90,7 +92,7 @@ export async function streamTurn(input: { readonly childSessionUpdate?: (update: ChildSessionUpdate) => Promise readonly connectionSignal?: AbortSignal readonly sessionSignal?: AbortSignal -}): Promise { +}): Promise<{ readonly response: PromptResponse; readonly contextTokens: number | undefined }> { const streamController = new AbortController() const connectionAbort = () => streamController.abort() input.connectionSignal?.addEventListener("abort", connectionAbort, { once: true }) @@ -100,15 +102,31 @@ export async function streamTurn(input: { const control = input.control let started = false - let assistantMessageID: string | undefined let finish: SessionMessageAssistant["finish"] let executionError: { readonly type: string; readonly message: string } | undefined + let stepError: SessionStructuredError | undefined + let usage: { readonly turn: TokenUsageInfo; readonly last: TokenUsageInfo } | undefined const tools = new Map() const retries = new Map() const children = new Map() const openChildren = new Set() let handedOff = false + const recordStep = (tokens: TokenUsageInfo) => { + const turn = usage?.turn + usage = { + turn: turn + ? { + input: turn.input + tokens.input, + output: turn.output + tokens.output, + reasoning: turn.reasoning + tokens.reasoning, + cache: { read: turn.cache.read + tokens.cache.read, write: turn.cache.write + tokens.cache.write }, + } + : tokens, + last: tokens, + } + } + const notifyChild = async (child: ChildSession, value: ChildSessionEvent) => { if (!input.childSessionUpdate) return await input @@ -195,7 +213,7 @@ export async function streamTurn(input: { } if (event.type === "session.step.started") { - if (!child) assistantMessageID = event.data.assistantMessageID + if (!child) stepError = undefined if (retries.delete(eventSessionID)) await send({ sessionUpdate: "session_info_update", _meta: { [RetryMeta]: null } }) continue @@ -211,7 +229,6 @@ export async function streamTurn(input: { continue } if (event.type === "session.text.delta") { - if (!child) assistantMessageID = event.data.assistantMessageID await send({ sessionUpdate: "agent_message_chunk", messageId: event.data.assistantMessageID, @@ -220,7 +237,6 @@ export async function streamTurn(input: { continue } if (event.type === "session.reasoning.delta") { - if (!child) assistantMessageID = event.data.assistantMessageID await send({ sessionUpdate: "agent_thought_chunk", messageId: `${event.data.assistantMessageID}:reasoning:${event.data.ordinal}`, @@ -229,7 +245,6 @@ export async function streamTurn(input: { continue } if (event.type === "session.tool.input.started") { - if (!child) assistantMessageID = event.data.assistantMessageID tools.set(toolKey(event.data.sessionID, event.data.id), { name: event.data.name, input: {}, @@ -248,7 +263,6 @@ export async function streamTurn(input: { continue } if (event.type === "session.tool.called") { - if (!child) assistantMessageID = event.data.assistantMessageID const key = toolKey(event.data.sessionID, event.data.id) const current = tools.get(key) ?? emptyToolState() current.input = event.data.input @@ -315,8 +329,15 @@ export async function streamTurn(input: { } if (event.type === "session.step.ended") { if (!child) { - assistantMessageID = event.data.assistantMessageID finish = event.data.finish + recordStep(event.data.tokens) + } + continue + } + if (event.type === "session.step.failed") { + if (!child) { + stepError = event.data.error + if (event.data.tokens) recordStep(event.data.tokens) } continue } @@ -362,14 +383,17 @@ export async function streamTurn(input: { if (input.action) { streamController.abort() await completed.catch(() => {}) - return response(undefined, undefined, "succeeded", control.cancelled, undefined) + return { + response: response(undefined, undefined, "succeeded", control.cancelled, undefined), + contextTokens: undefined, + } } if (control.cancelled) { await input.client.session.interrupt({ sessionID: input.sessionID }).catch(() => {}) if (!started) { streamController.abort() await completed.catch(() => {}) - return response(undefined, undefined, "interrupted", true, undefined) + return { response: response(undefined, undefined, "interrupted", true, undefined), contextTokens: undefined } } } const terminal = await completed @@ -380,19 +404,17 @@ export async function streamTurn(input: { .catch(() => {}) .finally(closeStream) } - const assistant = assistantMessageID - ? await input.client.session - .message.get({ sessionID: input.sessionID, messageID: assistantMessageID }) - .catch(() => undefined) - : undefined - return response( - assistant?.type === "assistant" ? assistant : undefined, - executionError, - terminal, - control.cancelled, - finish, - retries.get(input.sessionID), - ) + return { + response: response( + usage?.turn, + stepError ?? executionError, + terminal, + control.cancelled, + finish, + retries.get(input.sessionID), + ), + contextTokens: usage ? TokenUsage.total(usage.last) : undefined, + } } catch (error) { streamController.abort() await completed.catch(() => {}) @@ -566,14 +588,13 @@ function matchesStart(event: EventSubscribeOutput, start: TurnStart) { } function response( - assistant: SessionMessageAssistant | undefined, - executionError: { readonly type: string; readonly message: string } | undefined, + tokens: TokenUsageInfo | undefined, + error: { readonly type: string; readonly message: string } | undefined, terminal: "succeeded" | "failed" | "interrupted", cancelled: boolean, finish: SessionMessageAssistant["finish"], retry?: RetryStatus, ): PromptResponse { - const error = assistant?.error ?? executionError if (error?.type === "provider.auth") throw new ACPError.AuthRequiredError() if (error && error.type !== "aborted" && error.type !== "provider.content-filter") { throw new ACPError.ServiceFailureError({ @@ -582,12 +603,11 @@ function response( errorName: error.type, }) } - const tokens = assistant?.tokens const usage = tokens ? { inputTokens: tokens.input, outputTokens: tokens.output, - totalTokens: tokens.input + tokens.output + tokens.reasoning + tokens.cache.read + tokens.cache.write, + totalTokens: TokenUsage.total(tokens), ...(tokens.reasoning > 0 ? { thoughtTokens: tokens.reasoning } : {}), ...(tokens.cache.read > 0 ? { cachedReadTokens: tokens.cache.read } : {}), ...(tokens.cache.write > 0 ? { cachedWriteTokens: tokens.cache.write } : {}), diff --git a/packages/cli/src/acp/service.ts b/packages/cli/src/acp/service.ts index 66c9f634128..6de273d6962 100644 --- a/packages/cli/src/acp/service.ts +++ b/packages/cli/src/acp/service.ts @@ -357,9 +357,9 @@ export function make(input: { submit: (signal) => submitPrompt(input.client, state, prepared, signal), ...(childSessionUpdate ? { childSessionUpdate } : {}), }) - .then(async (response) => { - await sendUsageUpdate(input.client, input.connection, state, response.usage?.totalTokens).catch(() => {}) - return response + .then(async (result) => { + await sendUsageUpdate(input.client, input.connection, state, result.contextTokens).catch(() => {}) + return result.response }) .finally(() => { signal?.removeEventListener("abort", cancel) diff --git a/packages/cli/test/acp/prompt.test.ts b/packages/cli/test/acp/prompt.test.ts index aefa3a88756..26ae313ed28 100644 --- a/packages/cli/test/acp/prompt.test.ts +++ b/packages/cli/test/acp/prompt.test.ts @@ -2,7 +2,7 @@ import { describe, expect, test } from "bun:test" import type { StopReason } from "@agentclientprotocol/sdk" import type { OpenCodeEvent } from "@opencode/client/promise" import { - assistantMessage, + childCreated, delivered, durableEvent, failed, @@ -32,7 +32,6 @@ describe("acp prompt turns over the wire", () => { onPrompt: ({ sessionID, id }) => turn(sessionID, id, textDelta(sessionID, "msg_assistant", "hello"), stepEnded(sessionID, "msg_assistant")), }) - acp.server.messages.set(acp.sessionId, [assistantMessage("msg_assistant")]) const response = await acp.prompt(acp.sessionId, "hi") @@ -114,7 +113,6 @@ describe("acp prompt turns over the wire", () => { value: "test/second-model", }) acp.server.sessions.set(acp.sessionId, makeSession(acp.sessionId, { cost: 3.5 })) - acp.server.messages.set(acp.sessionId, [assistantMessage("msg_assistant", { tokens: assistantTokens })]) const response = await acp.prompt(acp.sessionId, "hello") @@ -145,7 +143,6 @@ describe("acp prompt turns over the wire", () => { return turn(sessionID, id, stepEnded(sessionID, "msg_assistant")) }, }) - acp.server.messages.set(acp.sessionId, [assistantMessage("msg_assistant")]) expect((await acp.prompt(acp.sessionId, "hello")).stopReason).toBe("end_turn") }) @@ -191,14 +188,20 @@ describe("acp prompt turns over the wire", () => { expect(await rpcError(acp.prompt(acp.sessionId, "hello"))).toMatchObject({ code: -32000 }) }) - test("maps an assistant message auth error to auth required", async () => { + test("maps an assistant step auth error to auth required", async () => { await using acp = await startSession({ onPrompt: ({ sessionID, id }) => - turn(sessionID, id, textDelta(sessionID, "msg_auth", "partial"), stepEnded(sessionID, "msg_auth")), + turn( + sessionID, + id, + textDelta(sessionID, "msg_auth", "partial"), + durableEvent("session.step.failed", { + sessionID, + assistantMessageID: "msg_auth", + error: { type: "provider.auth", message: "expired" }, + }), + ), }) - acp.server.messages.set(acp.sessionId, [ - assistantMessage("msg_auth", { error: { type: "provider.auth", message: "expired" } }), - ]) expect(await rpcError(acp.prompt(acp.sessionId, "hello"))).toMatchObject({ code: -32000 }) }) @@ -353,27 +356,110 @@ describe("acp prompt turns over the wire", () => { expect((await first).stopReason).toBe("cancelled") }) - test.todo( - "reports usage summed across every step of the turn (https://github.com/anomalyco/opencode/issues/41660)", - async () => { - await using acp = await startSession({ - onPrompt: ({ sessionID, id }) => - turn( + test("reports usage summed across every step of the turn", async () => { + await using acp = await startSession({ + onPrompt: ({ sessionID, id }) => + turn( + sessionID, + id, + stepEnded(sessionID, "msg_step_1", { finish: "tool-calls", tokens: { ...tokens(), input: 10, output: 5 } }), + stepEnded(sessionID, "msg_step_2", { tokens: { ...tokens(), input: 20, output: 7 } }), + ), + }) + + const response = await acp.prompt(acp.sessionId, "hello") + + expect(response.usage).toEqual({ inputTokens: 30, outputTokens: 12, totalTokens: 42 }) + }) + + test("publishes the last step's context usage rather than the turn sum", async () => { + await using acp = await startSession({ + onPrompt: ({ sessionID, id }) => + turn( + sessionID, + id, + stepEnded(sessionID, "msg_step_1", { + finish: "tool-calls", + tokens: { input: 100, output: 10, reasoning: 0, cache: { read: 0, write: 50 } }, + }), + stepEnded(sessionID, "msg_step_2", { + tokens: { input: 20, output: 5, reasoning: 3, cache: { read: 150, write: 0 } }, + }), + ), + }) + + const response = await acp.prompt(acp.sessionId, "hello") + + expect(response.usage).toEqual({ + inputTokens: 120, + outputTokens: 15, + thoughtTokens: 3, + cachedReadTokens: 150, + cachedWriteTokens: 50, + totalTokens: 338, + }) + expect(await acp.waitForUpdate((item) => item.update.sessionUpdate === "usage_update")).toEqual({ + sessionId: acp.sessionId, + update: { sessionUpdate: "usage_update", used: 178, size: 100_000, cost: { amount: 0, currency: "USD" } }, + }) + }) + + test("counts a failed step's tokens and clears its error when the next step starts", async () => { + await using acp = await startSession({ + onPrompt: ({ sessionID, id }) => + turn( + sessionID, + id, + durableEvent("session.step.failed", { sessionID, - id, - stepEnded(sessionID, "msg_step_1", { finish: "tool-calls", tokens: { ...tokens(), input: 10, output: 5 } }), - stepEnded(sessionID, "msg_step_2", { tokens: { ...tokens(), input: 20, output: 7 } }), - ), - }) - acp.server.messages.set(acp.sessionId, [ - assistantMessage("msg_step_2", { tokens: { ...tokens(), input: 20, output: 7 } }), - ]) + assistantMessageID: "msg_1", + error: { type: "provider.stream", message: "stream interrupted" }, + cost: 0, + tokens: { ...tokens(), input: 40, output: 4 }, + }), + durableEvent("session.step.started", { + sessionID, + assistantMessageID: "msg_2", + agent: "build", + model: { providerID: "test", id: "test-model" }, + started: 0, + }), + stepEnded(sessionID, "msg_2", { tokens: { ...tokens(), input: 20, output: 7 } }), + ), + }) - const response = await acp.prompt(acp.sessionId, "hello") + const response = await acp.prompt(acp.sessionId, "hello") - expect(response.usage).toEqual({ inputTokens: 30, outputTokens: 12, totalTokens: 42 }) - }, - ) + expect(response).toEqual({ + stopReason: "end_turn", + usage: { inputTokens: 60, outputTokens: 11, totalTokens: 71 }, + _meta: {}, + }) + expect(await acp.waitForUpdate((item) => item.update.sessionUpdate === "usage_update")).toMatchObject({ + update: { used: 27 }, + }) + }) + + test("excludes child session steps from the turn usage", async () => { + await using acp = await startSession({ + onPrompt: ({ sessionID, id }) => + turn( + sessionID, + id, + childCreated("ses_child", sessionID, "Explore"), + stepEnded("ses_child", "msg_child", { tokens: { ...tokens(), input: 500, output: 50 } }), + succeeded("ses_child"), + stepEnded(sessionID, "msg_root", { tokens: { ...tokens(), input: 20, output: 7 } }), + ), + }) + + const response = await acp.prompt(acp.sessionId, "hello") + + expect(response.usage).toEqual({ inputTokens: 20, outputTokens: 7, totalTokens: 27 }) + expect(await acp.waitForUpdate((item) => item.update.sessionUpdate === "usage_update")).toMatchObject({ + update: { used: 27 }, + }) + }) }) // The server answered admission before streaming the chunk, and this request round-trips through the server after it,