feat(acp): add session usage (#12299)
Co-authored-by: Aiden Cline <63023139+rekram1-node@users.noreply.github.com>
This commit is contained in:
4
bun.lock
4
bun.lock
@@ -265,7 +265,7 @@
|
|||||||
"dependencies": {
|
"dependencies": {
|
||||||
"@actions/core": "1.11.1",
|
"@actions/core": "1.11.1",
|
||||||
"@actions/github": "6.0.1",
|
"@actions/github": "6.0.1",
|
||||||
"@agentclientprotocol/sdk": "0.13.0",
|
"@agentclientprotocol/sdk": "0.14.1",
|
||||||
"@ai-sdk/amazon-bedrock": "3.0.74",
|
"@ai-sdk/amazon-bedrock": "3.0.74",
|
||||||
"@ai-sdk/anthropic": "2.0.58",
|
"@ai-sdk/anthropic": "2.0.58",
|
||||||
"@ai-sdk/azure": "2.0.91",
|
"@ai-sdk/azure": "2.0.91",
|
||||||
@@ -559,7 +559,7 @@
|
|||||||
|
|
||||||
"@adobe/css-tools": ["@adobe/css-tools@4.4.4", "", {}, "sha512-Elp+iwUx5rN5+Y8xLt5/GRoG20WGoDCQ/1Fb+1LiGtvwbDavuSk0jhD/eZdckHAuzcDzccnkv+rEjyWfRx18gg=="],
|
"@adobe/css-tools": ["@adobe/css-tools@4.4.4", "", {}, "sha512-Elp+iwUx5rN5+Y8xLt5/GRoG20WGoDCQ/1Fb+1LiGtvwbDavuSk0jhD/eZdckHAuzcDzccnkv+rEjyWfRx18gg=="],
|
||||||
|
|
||||||
"@agentclientprotocol/sdk": ["@agentclientprotocol/sdk@0.13.0", "", { "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" } }, "sha512-Z6/Fp4cXLbYdMXr5AK752JM5qG2VKb6ShM0Ql6FimBSckMmLyK54OA20UhPYoH4C37FSFwUTARuwQOwQUToYrw=="],
|
"@agentclientprotocol/sdk": ["@agentclientprotocol/sdk@0.14.1", "", { "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" } }, "sha512-b6r3PS3Nly+Wyw9U+0nOr47bV8tfS476EgyEMhoKvJCZLbgqoDFN7DJwkxL88RR0aiOqOYV1ZnESHqb+RmdH8w=="],
|
||||||
|
|
||||||
"@ai-sdk/amazon-bedrock": ["@ai-sdk/amazon-bedrock@3.0.74", "", { "dependencies": { "@ai-sdk/anthropic": "2.0.58", "@ai-sdk/provider": "2.0.1", "@ai-sdk/provider-utils": "3.0.20", "@smithy/eventstream-codec": "^4.0.1", "@smithy/util-utf8": "^4.0.0", "aws4fetch": "^1.0.20" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-q83HE3FBb/HPIvjXsehrHOgCuGHPorSMFt6BYnzIYZy8gNnSqV1OWX4oXVsCAuYPPMtYW/KMK35hmoIFV8QKoQ=="],
|
"@ai-sdk/amazon-bedrock": ["@ai-sdk/amazon-bedrock@3.0.74", "", { "dependencies": { "@ai-sdk/anthropic": "2.0.58", "@ai-sdk/provider": "2.0.1", "@ai-sdk/provider-utils": "3.0.20", "@smithy/eventstream-codec": "^4.0.1", "@smithy/util-utf8": "^4.0.0", "aws4fetch": "^1.0.20" }, "peerDependencies": { "zod": "^3.25.76 || ^4.1.8" } }, "sha512-q83HE3FBb/HPIvjXsehrHOgCuGHPorSMFt6BYnzIYZy8gNnSqV1OWX4oXVsCAuYPPMtYW/KMK35hmoIFV8QKoQ=="],
|
||||||
|
|
||||||
|
|||||||
@@ -50,7 +50,7 @@
|
|||||||
"dependencies": {
|
"dependencies": {
|
||||||
"@actions/core": "1.11.1",
|
"@actions/core": "1.11.1",
|
||||||
"@actions/github": "6.0.1",
|
"@actions/github": "6.0.1",
|
||||||
"@agentclientprotocol/sdk": "0.13.0",
|
"@agentclientprotocol/sdk": "0.14.1",
|
||||||
"@ai-sdk/amazon-bedrock": "3.0.74",
|
"@ai-sdk/amazon-bedrock": "3.0.74",
|
||||||
"@ai-sdk/anthropic": "2.0.58",
|
"@ai-sdk/anthropic": "2.0.58",
|
||||||
"@ai-sdk/azure": "2.0.91",
|
"@ai-sdk/azure": "2.0.91",
|
||||||
|
|||||||
@@ -25,6 +25,7 @@ import {
|
|||||||
type SetSessionModeResponse,
|
type SetSessionModeResponse,
|
||||||
type ToolCallContent,
|
type ToolCallContent,
|
||||||
type ToolKind,
|
type ToolKind,
|
||||||
|
type Usage,
|
||||||
} from "@agentclientprotocol/sdk"
|
} from "@agentclientprotocol/sdk"
|
||||||
|
|
||||||
import { Log } from "../util/log"
|
import { Log } from "../util/log"
|
||||||
@@ -38,7 +39,7 @@ import { Config } from "@/config/config"
|
|||||||
import { Todo } from "@/session/todo"
|
import { Todo } from "@/session/todo"
|
||||||
import { z } from "zod"
|
import { z } from "zod"
|
||||||
import { LoadAPIKeyError } from "ai"
|
import { LoadAPIKeyError } from "ai"
|
||||||
import type { Event, OpencodeClient, SessionMessageResponse } from "@opencode-ai/sdk/v2"
|
import type { AssistantMessage, Event, OpencodeClient, SessionMessageResponse } from "@opencode-ai/sdk/v2"
|
||||||
import { applyPatch } from "diff"
|
import { applyPatch } from "diff"
|
||||||
|
|
||||||
type ModeOption = { id: string; name: string; description?: string }
|
type ModeOption = { id: string; name: string; description?: string }
|
||||||
@@ -49,6 +50,74 @@ const DEFAULT_VARIANT_VALUE = "default"
|
|||||||
export namespace ACP {
|
export namespace ACP {
|
||||||
const log = Log.create({ service: "acp-agent" })
|
const log = Log.create({ service: "acp-agent" })
|
||||||
|
|
||||||
|
async function getContextLimit(
|
||||||
|
sdk: OpencodeClient,
|
||||||
|
providerID: string,
|
||||||
|
modelID: string,
|
||||||
|
directory: string,
|
||||||
|
): Promise<number | null> {
|
||||||
|
const providers = await sdk.config
|
||||||
|
.providers({ directory })
|
||||||
|
.then((x) => x.data?.providers ?? [])
|
||||||
|
.catch((error) => {
|
||||||
|
log.error("failed to get providers for context limit", { error })
|
||||||
|
return []
|
||||||
|
})
|
||||||
|
|
||||||
|
const provider = providers.find((p) => p.id === providerID)
|
||||||
|
const model = provider?.models[modelID]
|
||||||
|
return model?.limit.context ?? null
|
||||||
|
}
|
||||||
|
|
||||||
|
async function sendUsageUpdate(
|
||||||
|
connection: AgentSideConnection,
|
||||||
|
sdk: OpencodeClient,
|
||||||
|
sessionID: string,
|
||||||
|
directory: string,
|
||||||
|
): Promise<void> {
|
||||||
|
const messages = await sdk.session
|
||||||
|
.messages({ sessionID, directory }, { throwOnError: true })
|
||||||
|
.then((x) => x.data)
|
||||||
|
.catch((error) => {
|
||||||
|
log.error("failed to fetch messages for usage update", { error })
|
||||||
|
return undefined
|
||||||
|
})
|
||||||
|
|
||||||
|
if (!messages) return
|
||||||
|
|
||||||
|
const assistantMessages = messages.filter(
|
||||||
|
(m): m is { info: AssistantMessage; parts: SessionMessageResponse["parts"] } => m.info.role === "assistant",
|
||||||
|
)
|
||||||
|
|
||||||
|
const lastAssistant = assistantMessages[assistantMessages.length - 1]
|
||||||
|
if (!lastAssistant) return
|
||||||
|
|
||||||
|
const msg = lastAssistant.info
|
||||||
|
const size = await getContextLimit(sdk, msg.providerID, msg.modelID, directory)
|
||||||
|
|
||||||
|
if (!size) {
|
||||||
|
// Cannot calculate usage without known context size
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
const used = msg.tokens.input + (msg.tokens.cache?.read ?? 0)
|
||||||
|
const totalCost = assistantMessages.reduce((sum, m) => sum + m.info.cost, 0)
|
||||||
|
|
||||||
|
await connection
|
||||||
|
.sessionUpdate({
|
||||||
|
sessionId: sessionID,
|
||||||
|
update: {
|
||||||
|
sessionUpdate: "usage_update",
|
||||||
|
used,
|
||||||
|
size,
|
||||||
|
cost: { amount: totalCost, currency: "USD" },
|
||||||
|
},
|
||||||
|
})
|
||||||
|
.catch((error) => {
|
||||||
|
log.error("failed to send usage update", { error })
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
export async function init({ sdk: _sdk }: { sdk: OpencodeClient }) {
|
export async function init({ sdk: _sdk }: { sdk: OpencodeClient }) {
|
||||||
return {
|
return {
|
||||||
create: (connection: AgentSideConnection, fullConfig: ACPConfig) => {
|
create: (connection: AgentSideConnection, fullConfig: ACPConfig) => {
|
||||||
@@ -546,6 +615,8 @@ export namespace ACP {
|
|||||||
await this.processMessage(msg)
|
await this.processMessage(msg)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
await sendUsageUpdate(this.connection, this.sdk, sessionId, directory)
|
||||||
|
|
||||||
return result
|
return result
|
||||||
} catch (e) {
|
} catch (e) {
|
||||||
const error = MessageV2.fromError(e, {
|
const error = MessageV2.fromError(e, {
|
||||||
@@ -654,6 +725,8 @@ export namespace ACP {
|
|||||||
await this.processMessage(msg)
|
await this.processMessage(msg)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
await sendUsageUpdate(this.connection, this.sdk, sessionId, directory)
|
||||||
|
|
||||||
return mode
|
return mode
|
||||||
} catch (e) {
|
} catch (e) {
|
||||||
const error = MessageV2.fromError(e, {
|
const error = MessageV2.fromError(e, {
|
||||||
@@ -677,11 +750,15 @@ export namespace ACP {
|
|||||||
|
|
||||||
log.info("resume_session", { sessionId, mcpServers: mcpServers.length })
|
log.info("resume_session", { sessionId, mcpServers: mcpServers.length })
|
||||||
|
|
||||||
return this.loadSessionMode({
|
const result = await this.loadSessionMode({
|
||||||
cwd: directory,
|
cwd: directory,
|
||||||
mcpServers,
|
mcpServers,
|
||||||
sessionId,
|
sessionId,
|
||||||
})
|
})
|
||||||
|
|
||||||
|
await sendUsageUpdate(this.connection, this.sdk, sessionId, directory)
|
||||||
|
|
||||||
|
return result
|
||||||
} catch (e) {
|
} catch (e) {
|
||||||
const error = MessageV2.fromError(e, {
|
const error = MessageV2.fromError(e, {
|
||||||
providerID: this.config.defaultModel?.providerID ?? "unknown",
|
providerID: this.config.defaultModel?.providerID ?? "unknown",
|
||||||
@@ -1239,13 +1316,22 @@ export namespace ACP {
|
|||||||
return { name, args: rest.join(" ").trim() }
|
return { name, args: rest.join(" ").trim() }
|
||||||
})()
|
})()
|
||||||
|
|
||||||
const done = {
|
const buildUsage = (msg: AssistantMessage): Usage => ({
|
||||||
stopReason: "end_turn" as const,
|
totalTokens:
|
||||||
_meta: {},
|
msg.tokens.input +
|
||||||
}
|
msg.tokens.output +
|
||||||
|
msg.tokens.reasoning +
|
||||||
|
(msg.tokens.cache?.read ?? 0) +
|
||||||
|
(msg.tokens.cache?.write ?? 0),
|
||||||
|
inputTokens: msg.tokens.input,
|
||||||
|
outputTokens: msg.tokens.output,
|
||||||
|
thoughtTokens: msg.tokens.reasoning || undefined,
|
||||||
|
cachedReadTokens: msg.tokens.cache?.read || undefined,
|
||||||
|
cachedWriteTokens: msg.tokens.cache?.write || undefined,
|
||||||
|
})
|
||||||
|
|
||||||
if (!cmd) {
|
if (!cmd) {
|
||||||
await this.sdk.session.prompt({
|
const response = await this.sdk.session.prompt({
|
||||||
sessionID,
|
sessionID,
|
||||||
model: {
|
model: {
|
||||||
providerID: model.providerID,
|
providerID: model.providerID,
|
||||||
@@ -1256,14 +1342,22 @@ export namespace ACP {
|
|||||||
agent,
|
agent,
|
||||||
directory,
|
directory,
|
||||||
})
|
})
|
||||||
return done
|
const msg = response.data?.info
|
||||||
|
|
||||||
|
await sendUsageUpdate(this.connection, this.sdk, sessionID, directory)
|
||||||
|
|
||||||
|
return {
|
||||||
|
stopReason: "end_turn" as const,
|
||||||
|
usage: msg ? buildUsage(msg) : undefined,
|
||||||
|
_meta: {},
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
const command = await this.config.sdk.command
|
const command = await this.config.sdk.command
|
||||||
.list({ directory }, { throwOnError: true })
|
.list({ directory }, { throwOnError: true })
|
||||||
.then((x) => x.data!.find((c) => c.name === cmd.name))
|
.then((x) => x.data!.find((c) => c.name === cmd.name))
|
||||||
if (command) {
|
if (command) {
|
||||||
await this.sdk.session.command({
|
const response = await this.sdk.session.command({
|
||||||
sessionID,
|
sessionID,
|
||||||
command: command.name,
|
command: command.name,
|
||||||
arguments: cmd.args,
|
arguments: cmd.args,
|
||||||
@@ -1271,7 +1365,15 @@ export namespace ACP {
|
|||||||
agent,
|
agent,
|
||||||
directory,
|
directory,
|
||||||
})
|
})
|
||||||
return done
|
const msg = response.data?.info
|
||||||
|
|
||||||
|
await sendUsageUpdate(this.connection, this.sdk, sessionID, directory)
|
||||||
|
|
||||||
|
return {
|
||||||
|
stopReason: "end_turn" as const,
|
||||||
|
usage: msg ? buildUsage(msg) : undefined,
|
||||||
|
_meta: {},
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
switch (cmd.name) {
|
switch (cmd.name) {
|
||||||
@@ -1288,7 +1390,12 @@ export namespace ACP {
|
|||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
|
||||||
return done
|
await sendUsageUpdate(this.connection, this.sdk, sessionID, directory)
|
||||||
|
|
||||||
|
return {
|
||||||
|
stopReason: "end_turn" as const,
|
||||||
|
_meta: {},
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async cancel(params: CancelNotification) {
|
async cancel(params: CancelNotification) {
|
||||||
|
|||||||
Reference in New Issue
Block a user