mirrors/opencode
1
0
Fork 0
mirror of https://github.com/anomalyco/opencode.git synced 2026-11-15 12:41:37 +03:00
Code Issues Packages Projects Releases Wiki Activity

fix(acp): report usage for the whole turn (#52467)

Browse Source
This commit is contained in:
Shoubhit Dash
2026-10-01 16:21:26 +05:30
committed by GitHub
Unverified
parent 966e718b31
commit 20cf09752d
3 changed files with 164 additions and 58 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/cli/src/acp/event.ts
+48 -28
View File
@@ -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<void>
readonly connectionSignal?: AbortSignal
readonly sessionSignal?: AbortSignal
}): Promise<PromptResponse> {
}): 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<string, ToolState>()
const retries = new Map<string, RetryStatus>()
const children = new Map<string, ChildSession>()
const openChildren = new Set<string>()
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 } : {}),
packages/cli/src/acp/service.ts
+3 -3
View File
@@ -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)
packages/cli/test/acp/prompt.test.ts
+113 -27
View File
@@ -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,