Преглед изворни кода

feat(acp): add session usage (#12299)

Co-authored-by: Aiden Cline <[email protected]>
Steffen Deusch пре 1 недеља
родитељ
комит
8ddef975b7
3 измењених фајлова са 121 додато и 14 уклоњено
  1. 2 2
      bun.lock
  2. 1 1
      packages/opencode/package.json
  3. 118 11
      packages/opencode/src/acp/agent.ts

+ 2 - 2
bun.lock

@@ -265,7 +265,7 @@
       "dependencies": {
         "@actions/core": "1.11.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/anthropic": "2.0.58",
         "@ai-sdk/azure": "2.0.91",
@@ -559,7 +559,7 @@
 
     "@adobe/css-tools": ["@adobe/[email protected]", "", {}, "sha512-Elp+iwUx5rN5+Y8xLt5/GRoG20WGoDCQ/1Fb+1LiGtvwbDavuSk0jhD/eZdckHAuzcDzccnkv+rEjyWfRx18gg=="],
 
-    "@agentclientprotocol/sdk": ["@agentclientprotocol/[email protected]3.0", "", { "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" } }, "sha512-Z6/Fp4cXLbYdMXr5AK752JM5qG2VKb6ShM0Ql6FimBSckMmLyK54OA20UhPYoH4C37FSFwUTARuwQOwQUToYrw=="],
+    "@agentclientprotocol/sdk": ["@agentclientprotocol/[email protected]4.1", "", { "peerDependencies": { "zod": "^3.25.0 || ^4.0.0" } }, "sha512-b6r3PS3Nly+Wyw9U+0nOr47bV8tfS476EgyEMhoKvJCZLbgqoDFN7DJwkxL88RR0aiOqOYV1ZnESHqb+RmdH8w=="],
 
     "@ai-sdk/amazon-bedrock": ["@ai-sdk/[email protected]", "", { "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=="],
 

+ 1 - 1
packages/opencode/package.json

@@ -50,7 +50,7 @@
   "dependencies": {
     "@actions/core": "1.11.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/anthropic": "2.0.58",
     "@ai-sdk/azure": "2.0.91",

+ 118 - 11
packages/opencode/src/acp/agent.ts

@@ -25,6 +25,7 @@ import {
   type SetSessionModeResponse,
   type ToolCallContent,
   type ToolKind,
+  type Usage,
 } from "@agentclientprotocol/sdk"
 
 import { Log } from "../util/log"
@@ -38,7 +39,7 @@ import { Config } from "@/config/config"
 import { Todo } from "@/session/todo"
 import { z } from "zod"
 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"
 
 type ModeOption = { id: string; name: string; description?: string }
@@ -49,6 +50,74 @@ const DEFAULT_VARIANT_VALUE = "default"
 export namespace ACP {
   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 }) {
     return {
       create: (connection: AgentSideConnection, fullConfig: ACPConfig) => {
@@ -546,6 +615,8 @@ export namespace ACP {
           await this.processMessage(msg)
         }
 
+        await sendUsageUpdate(this.connection, this.sdk, sessionId, directory)
+
         return result
       } catch (e) {
         const error = MessageV2.fromError(e, {
@@ -654,6 +725,8 @@ export namespace ACP {
           await this.processMessage(msg)
         }
 
+        await sendUsageUpdate(this.connection, this.sdk, sessionId, directory)
+
         return mode
       } catch (e) {
         const error = MessageV2.fromError(e, {
@@ -677,11 +750,15 @@ export namespace ACP {
 
         log.info("resume_session", { sessionId, mcpServers: mcpServers.length })
 
-        return this.loadSessionMode({
+        const result = await this.loadSessionMode({
           cwd: directory,
           mcpServers,
           sessionId,
         })
+
+        await sendUsageUpdate(this.connection, this.sdk, sessionId, directory)
+
+        return result
       } catch (e) {
         const error = MessageV2.fromError(e, {
           providerID: this.config.defaultModel?.providerID ?? "unknown",
@@ -1239,13 +1316,22 @@ export namespace ACP {
         return { name, args: rest.join(" ").trim() }
       })()
 
-      const done = {
-        stopReason: "end_turn" as const,
-        _meta: {},
-      }
+      const buildUsage = (msg: AssistantMessage): Usage => ({
+        totalTokens:
+          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) {
-        await this.sdk.session.prompt({
+        const response = await this.sdk.session.prompt({
           sessionID,
           model: {
             providerID: model.providerID,
@@ -1256,14 +1342,22 @@ export namespace ACP {
           agent,
           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
         .list({ directory }, { throwOnError: true })
         .then((x) => x.data!.find((c) => c.name === cmd.name))
       if (command) {
-        await this.sdk.session.command({
+        const response = await this.sdk.session.command({
           sessionID,
           command: command.name,
           arguments: cmd.args,
@@ -1271,7 +1365,15 @@ export namespace ACP {
           agent,
           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) {
@@ -1288,7 +1390,12 @@ export namespace ACP {
           break
       }
 
-      return done
+      await sendUsageUpdate(this.connection, this.sdk, sessionID, directory)
+
+      return {
+        stopReason: "end_turn" as const,
+        _meta: {},
+      }
     }
 
     async cancel(params: CancelNotification) {