feat(agent): report per-run token usage to TaskRun.state

The agent-server now accumulates each settled turn's ACP usage (input/output/cache read/cache write/thought tokens) into run-level totals and PATCHes the snapshot into TaskRun.state.token_usage after every turn. Best-effort and adapter-agnostic: both the claude and codex adapters resolve prompt() with usage. This gives cloud runs a persisted, queryable token expenditure record even when a run never completes cleanly.

Generated-By: PostHog Code
Task-Id: 2d6eed7f-9209-4d1f-be54-ae3e22bf2990

Author: tatoalo

packages/agent/src/server/agent-server.test.ts

@@ -685,6 +685,78 @@ describe("AgentServer HTTP Mode", () => {
       expect(testServer.posthogAPI.updateTaskRun).not.toHaveBeenCalled();
     });
 
+    function createUsageTestServer() {
+      const testServer = new AgentServer({
+        port,
+        jwtPublicKey: TEST_PUBLIC_KEY,
+        repositoryPath: repo.path,
+        apiUrl: "http://localhost:8000",
+        apiKey: "test-api-key",
+        projectId: 1,
+        mode: "interactive",
+        taskId: "test-task-id",
+        runId: "test-run-id",
+      }) as unknown as {
+        session: { payload: JwtPayload } | null;
+        posthogAPI: { updateTaskRun: ReturnType<typeof vi.fn> };
+        recordTurnUsage(usage: unknown): void;
+      };
+      testServer.posthogAPI = { updateTaskRun: vi.fn(async () => ({})) };
+      testServer.session = {
+        payload: {
+          run_id: "run-1",
+          task_id: "task-1",
+          team_id: 1,
+          user_id: 1,
+          distinct_id: "distinct-id",
+          mode: "interactive",
+        },
+      };
+      return testServer;
+    }
+
+    it("reports cumulative run token usage into TaskRun.state after each settled turn", () => {
+      const testServer = createUsageTestServer();
+      const turnUsage = {
+        inputTokens: 100,
+        outputTokens: 50,
+        cachedReadTokens: 10,
+        cachedWriteTokens: 5,
+        totalTokens: 165,
+      };
+
+      testServer.recordTurnUsage(turnUsage);
+      testServer.recordTurnUsage(turnUsage);
+
+      expect(testServer.posthogAPI.updateTaskRun).toHaveBeenCalledTimes(2);
+      // The second report carries run-cumulative totals, not per-turn figures.
+      expect(testServer.posthogAPI.updateTaskRun).toHaveBeenLastCalledWith(
+        "task-1",
+        "run-1",
+        {
+          state: {
+            token_usage: {
+              input_tokens: 200,
+              output_tokens: 100,
+              cache_read_tokens: 20,
+              cache_write_tokens: 10,
+              thought_tokens: 0,
+              total_tokens: 330,
+              turns: 2,
+            },
+          },
+        },
+      );
+    });
+
+    it("does not report anything when a turn settles without usage", () => {
+      const testServer = createUsageTestServer();
+
+      testServer.recordTurnUsage(undefined);
+
+      expect(testServer.posthogAPI.updateTaskRun).not.toHaveBeenCalled();
+    });
+
     function createFailureTestServer() {
       const appendRawLine = vi.fn();
       const testServer = new AgentServer({

packages/agent/src/server/agent-server.ts

@@ -77,6 +77,7 @@ import {
 } from "./cloud-prompt";
 import { TaskRunEventStreamSender } from "./event-stream-sender";
 import { type JwtPayload, JwtValidationError, validateJwt } from "./jwt";
+import { RunUsageAccumulator } from "./run-usage";
 import {
   handoffLocalGitStateSchema,
   jsonRpcRequestSchema,
@@ -308,6 +309,7 @@ export class AgentServer {
   private eventStreamSender: TaskRunEventStreamSender | null = null;
   private questionRelayedToSlack = false;
   private adapterEmittedTurnComplete = false;
+  private readonly runUsage = new RunUsageAccumulator();
   private detectedPrUrl: string | null = null;
   // Reset per session. `evaluatedPrUrls` dedupes per URL; `prAttributionChain` serializes
   // attributions so the most recently created PR in a run wins.
@@ -947,6 +949,7 @@ export class AgentServer {
           void this.syncCloudBranchMetadata(this.session.payload);
         }
 
+        this.recordTurnUsage(result.usage);
         this.broadcastTurnComplete(result.stopReason);
 
         if (result.stopReason === "end_turn") {
@@ -1659,6 +1662,7 @@ export class AgentServer {
         void this.syncCloudBranchMetadata(payload);
       }
 
+      this.recordTurnUsage(result.usage);
       this.broadcastTurnComplete(result.stopReason);
 
       if (result.stopReason === "end_turn") {
@@ -1810,6 +1814,7 @@ export class AgentServer {
         void this.syncCloudBranchMetadata(payload);
       }
 
+      this.recordTurnUsage(result.usage);
       this.broadcastTurnComplete(result.stopReason);
 
       if (result.stopReason === "end_turn") {
@@ -3611,6 +3616,24 @@ ${signedCommitInstructions}${prLinkInstructions}${shellEfficiencyInstructions}
     return result.success ? result.data : null;
   }
 
+  /**
+   * Accumulates a settled turn's token usage into the run total and reports it
+   * to the backend, merged into `TaskRun.state.token_usage`. Best-effort: a
+   * reporting failure must never affect the turn outcome.
+   */
+  private recordTurnUsage(usage: PromptResponse["usage"]): void {
+    if (!this.runUsage.add(usage)) return;
+    const payload = this.session?.payload;
+    if (!payload) return;
+    void this.posthogAPI
+      .updateTaskRun(payload.task_id, payload.run_id, {
+        state: { token_usage: this.runUsage.snapshot() },
+      })
+      .catch((error) => {
+        this.logger.warn("Failed to report run token usage", error);
+      });
+  }
+
   private broadcastTurnComplete(stopReason: string): void {
     if (!this.session) return;
     if (this.adapterEmittedTurnComplete) {

packages/agent/src/server/run-usage.test.ts

@@ -0,0 +1,58 @@
+import { describe, expect, test } from "vitest";
+import { RunUsageAccumulator } from "./run-usage";
+
+describe("RunUsageAccumulator", () => {
+  test("accumulates across turns and defaults nullable ACP fields to 0", () => {
+    const acc = new RunUsageAccumulator();
+
+    // Claude-shaped turn: cache components present, no thought tokens.
+    expect(
+      acc.add({
+        inputTokens: 100,
+        outputTokens: 50,
+        cachedReadTokens: 10,
+        cachedWriteTokens: 5,
+        totalTokens: 165,
+      }),
+    ).toBe(true);
+    // Codex-shaped turn: null cache writes, reasoning as thought tokens.
+    expect(
+      acc.add({
+        inputTokens: 200,
+        outputTokens: 80,
+        cachedReadTokens: null,
+        cachedWriteTokens: null,
+        thoughtTokens: 40,
+        totalTokens: 320,
+      }),
+    ).toBe(true);
+
+    expect(acc.snapshot()).toEqual({
+      input_tokens: 300,
+      output_tokens: 130,
+      cache_read_tokens: 10,
+      cache_write_tokens: 5,
+      thought_tokens: 40,
+      total_tokens: 485,
+      turns: 2,
+    });
+  });
+
+  test.each([[null], [undefined]])(
+    "ignores a turn that settles with %s usage",
+    (usage) => {
+      const acc = new RunUsageAccumulator();
+      expect(acc.add(usage)).toBe(false);
+      expect(acc.snapshot().turns).toBe(0);
+    },
+  );
+
+  test("snapshot is a copy — later turns don't mutate earlier snapshots", () => {
+    const acc = new RunUsageAccumulator();
+    acc.add({ inputTokens: 1, outputTokens: 1, totalTokens: 2 });
+    const first = acc.snapshot();
+    acc.add({ inputTokens: 1, outputTokens: 1, totalTokens: 2 });
+    expect(first.turns).toBe(1);
+    expect(acc.snapshot().turns).toBe(2);
+  });
+});

packages/agent/src/server/run-usage.ts

@@ -0,0 +1,51 @@
+import type { Usage } from "@agentclientprotocol/sdk";
+
+/**
+ * Cumulative token usage for a task run, shaped for `TaskRun.state.token_usage`
+ * (snake_case, matching the backend's state conventions). `turns` counts the
+ * settled turns that contributed usage, giving consumers a per-turn denominator.
+ */
+export type RunTokenUsage = {
+  input_tokens: number;
+  output_tokens: number;
+  cache_read_tokens: number;
+  cache_write_tokens: number;
+  thought_tokens: number;
+  total_tokens: number;
+  turns: number;
+};
+
+/**
+ * Accumulates per-turn ACP `Usage` into run-level totals. The ACP usage fields
+ * are optional and nullable, so every component defaults to 0 to keep the sums
+ * numeric across adapters (codex reports no cache writes, claude no thought
+ * tokens on some models).
+ */
+export class RunUsageAccumulator {
+  private totals: RunTokenUsage = {
+    input_tokens: 0,
+    output_tokens: 0,
+    cache_read_tokens: 0,
+    cache_write_tokens: 0,
+    thought_tokens: 0,
+    total_tokens: 0,
+    turns: 0,
+  };
+
+  /** Adds a settled turn's usage. Returns false when there was nothing to add. */
+  add(usage: Usage | null | undefined): boolean {
+    if (!usage) return false;
+    this.totals.input_tokens += usage.inputTokens ?? 0;
+    this.totals.output_tokens += usage.outputTokens ?? 0;
+    this.totals.cache_read_tokens += usage.cachedReadTokens ?? 0;
+    this.totals.cache_write_tokens += usage.cachedWriteTokens ?? 0;
+    this.totals.thought_tokens += usage.thoughtTokens ?? 0;
+    this.totals.total_tokens += usage.totalTokens ?? 0;
+    this.totals.turns += 1;
+    return true;
+  }
+
+  snapshot(): RunTokenUsage {
+    return { ...this.totals };
+  }
+}