Implement session usage updates (#92)

* Implement session usage updates

* Reduce changes in package.json

* Add token-usage-session-update-multiple.txt to git

* Fix test data

* Fix test data

Author: egor-baranov

package-lock.json

@@ -50,10 +50,10 @@
     "vitest": "^4.0.10"
   },
   "dependencies": {
-    "@agentclientprotocol/sdk": "^0.12.0",
+    "@agentclientprotocol/sdk": "^0.16.1",
     "vscode-jsonrpc": "^8.2.1",
     "open": "^11.0.0",
     "diff": "^8.0.3",
-    "@openai/codex": "^0.114.0"
+    "@openai/codex": "^0.118.0"
   }
 }

package.json

@@ -24,6 +24,7 @@ import type {RateLimitsMap} from "./RateLimitsMap";
 import {ModelId} from "./ModelId";
 import {AgentMode} from "./AgentMode";
 import type {TokenCount} from "./TokenCount";
+import {toPromptUsage} from "./TokenCount";
 import {CodexCommands} from "./CodexCommands";
 import type {QuotaMeta} from "./QuotaMeta";
 import {logger} from "./Logger";
@@ -714,6 +715,7 @@ export class CodexAcpServer implements acp.Agent {
                 logger.log("Prompt handled by a command");
                 return {
                     stopReason: "end_turn",
+                    usage: this.buildPromptUsage(sessionState.lastTokenUsage),
                     _meta: this.buildQuotaMeta(sessionState),
                 };
             }
@@ -751,6 +753,7 @@ export class CodexAcpServer implements acp.Agent {
                 });
                 return {
                     stopReason: "cancelled",
+                    usage: this.buildPromptUsage(sessionState.lastTokenUsage),
                     _meta: this.buildQuotaMeta(sessionState),
                 };
             }
@@ -763,6 +766,7 @@ export class CodexAcpServer implements acp.Agent {
 
             return {
                 stopReason: "end_turn",
+                usage: this.buildPromptUsage(sessionState.lastTokenUsage),
                 _meta: this.buildQuotaMeta(sessionState),
             };
         } catch (err) {
@@ -793,6 +797,13 @@ export class CodexAcpServer implements acp.Agent {
         };
     }
 
+    private buildPromptUsage(lastTokenUsage: TokenCount | null): acp.Usage | null {
+        if (lastTokenUsage == null) {
+            return null;
+        }
+        return toPromptUsage(lastTokenUsage);
+    }
+
     private async runWithProcessCheck<T>(operation: () => Promise<T>): Promise<T> {
         try {
             return await operation();

src/CodexAcpServer.ts

@@ -87,8 +87,7 @@ export class CodexEventHandler {
                 this.sessionState.currentTurnId = null;
                 return null;
             case "thread/tokenUsage/updated":
-                this.handleTokenUsageUpdated(notification.params);
-                return null;
+                return this.createUsageUpdate(notification.params);
             case "item/commandExecution/outputDelta":
                 return this.createCommandOutputDeltaEvent(notification.params);
             case "command/exec/outputDelta":
@@ -383,6 +382,22 @@ export class CodexEventHandler {
         this.sessionState.modelContextWindow = params.tokenUsage.modelContextWindow;
     }
 
+    private createUsageUpdate(params: ThreadTokenUsageUpdatedNotification): UpdateSessionEvent | null {
+        this.handleTokenUsageUpdated(params);
+
+        const used = this.sessionState.totalTokenUsage?.totalTokens;
+        const size = this.sessionState.modelContextWindow;
+        if (used == null || size == null || size <= 0) {
+            return null;
+        }
+
+        return {
+            sessionUpdate: "usage_update",
+            used,
+            size,
+        };
+    }
+
     private handleRateLimitsUpdated(params: AccountRateLimitsUpdatedNotification): void {
         if (!this.sessionState.rateLimits) {
             this.sessionState.rateLimits = new Map();

src/CodexEventHandler.ts

@@ -1,3 +1,4 @@
+import type {Usage} from "@agentclientprotocol/sdk";
 import type {TokenUsageBreakdown} from "./app-server/v2";
 
 /**
@@ -33,3 +34,18 @@ export function toTokenCount(usage: TokenUsageBreakdown): TokenCount {
         reasoningOutputTokens: usage.reasoningOutputTokens,
     };
 }
+
+/**
+ * Maps our per-turn token breakdown to ACP PromptResponse usage fields.
+ * Cached input tokens are reported as ACP cache reads, and reasoning output
+ * tokens are exposed through ACP's thoughtTokens field.
+ */
+export function toPromptUsage(tokenCount: TokenCount): Usage {
+    return {
+        totalTokens: tokenCount.totalTokens,
+        inputTokens: tokenCount.inputTokens,
+        cachedReadTokens: tokenCount.cachedInputTokens,
+        outputTokens: tokenCount.outputTokens,
+        thoughtTokens: tokenCount.reasoningOutputTokens,
+    };
+}

src/TokenCount.ts

@@ -1,5 +1,12 @@
 {
   "stopReason": "cancelled",
+  "usage": {
+    "totalTokens": 1500,
+    "inputTokens": 1200,
+    "cachedReadTokens": 0,
+    "outputTokens": 300,
+    "thoughtTokens": 0
+  },
   "_meta": {
     "quota": {
       "token_count": {
@@ -23,4 +30,4 @@
       ]
     }
   }
-}
+}

src/__tests__/CodexACPAgent/data/token-usage-cancelled.json

@@ -1,5 +1,12 @@
 {
   "stopReason": "end_turn",
+  "usage": {
+    "totalTokens": 2500,
+    "inputTokens": 1500,
+    "cachedReadTokens": 500,
+    "outputTokens": 450,
+    "thoughtTokens": 50
+  },
   "_meta": {
     "quota": {
       "token_count": {
@@ -23,4 +30,4 @@
       ]
     }
   }
-}
+}

src/__tests__/CodexACPAgent/data/token-usage-end-turn.json

@@ -1,5 +1,12 @@
 {
   "stopReason": "end_turn",
+  "usage": {
+    "totalTokens": 1500,
+    "inputTokens": 700,
+    "cachedReadTokens": 500,
+    "outputTokens": 200,
+    "thoughtTokens": 100
+  },
   "_meta": {
     "quota": {
       "token_count": {
@@ -23,4 +30,4 @@
       ]
     }
   }
-}
+}

src/__tests__/CodexACPAgent/data/token-usage-multiple-updates.json

@@ -1,9 +1,10 @@
 {
   "stopReason": "end_turn",
+  "usage": null,
   "_meta": {
     "quota": {
       "token_count": null,
       "model_usage": []
     }
   }
-}
+}

src/__tests__/CodexACPAgent/data/token-usage-null.json

@@ -0,0 +1,41 @@
+[
+  {
+    "method": "sessionUpdate",
+    "args": [
+      {
+        "sessionId": "test-session-id",
+        "update": {
+          "sessionUpdate": "usage_update",
+          "used": 1000,
+          "size": 128000
+        }
+      }
+    ]
+  },
+  {
+    "method": "sessionUpdate",
+    "args": [
+      {
+        "sessionId": "test-session-id",
+        "update": {
+          "sessionUpdate": "usage_update",
+          "used": 2000,
+          "size": 128000
+        }
+      }
+    ]
+  },
+  {
+    "method": "sessionUpdate",
+    "args": [
+      {
+        "sessionId": "test-session-id",
+        "update": {
+          "sessionUpdate": "usage_update",
+          "used": 3500,
+          "size": 128000
+        }
+      }
+    ]
+  }
+]