diff --git a/.github/workflows/claude.yml b/.github/workflows/claude.yml index 0459b4c..5e49c40 100644 --- a/.github/workflows/claude.yml +++ b/.github/workflows/claude.yml @@ -57,6 +57,7 @@ jobs: bun-version: "1.3.12" - name: Run Claude + id: claude if: steps.prompt.outputs.skip != 'true' uses: anthropics/claude-code-action@v1 with: @@ -78,6 +79,17 @@ jobs: } } + - name: Report token usage + if: steps.claude.outcome != 'skipped' && steps.claude.outputs.execution_file != '' + continue-on-error: true + env: + GITHUB_TOKEN: ${{ secrets.AGENT_TOKEN }} + EXEC_FILE: ${{ steps.claude.outputs.execution_file }} + ENTITY_TYPE: issue + ENTITY_NUMBER: ${{ github.event.issue.number }} + REPO: ${{ github.repository }} + run: bash scripts/report-token-usage.sh + # --------------------------------------------------------------------------- # Comment on a plain issue → update the related PR # --------------------------------------------------------------------------- @@ -117,6 +129,7 @@ jobs: bun-version: "1.3.12" - name: Run Claude + id: claude if: steps.prompt.outputs.skip != 'true' uses: anthropics/claude-code-action@v1 with: @@ -138,6 +151,17 @@ jobs: } } + - name: Report token usage + if: steps.claude.outcome != 'skipped' && steps.claude.outputs.execution_file != '' + continue-on-error: true + env: + GITHUB_TOKEN: ${{ secrets.AGENT_TOKEN }} + EXEC_FILE: ${{ steps.claude.outputs.execution_file }} + ENTITY_TYPE: issue + ENTITY_NUMBER: ${{ github.event.issue.number }} + REPO: ${{ github.repository }} + run: bash scripts/report-token-usage.sh + # --------------------------------------------------------------------------- # PR opened/updated (not Claude's) → post a review # --------------------------------------------------------------------------- @@ -264,6 +288,17 @@ jobs: GH_TOKEN: ${{ secrets.AGENT_TOKEN }} PR_NUMBER: ${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.event.workflow_run.pull_requests[0].number }} + - name: Report token usage + if: steps.skip-check.outputs.skip != 'true' && steps.claude.outcome != 'skipped' && steps.claude.outputs.execution_file != '' + continue-on-error: true + env: + GITHUB_TOKEN: ${{ secrets.AGENT_TOKEN }} + EXEC_FILE: ${{ steps.claude.outputs.execution_file }} + ENTITY_TYPE: pr + ENTITY_NUMBER: ${{ github.event_name == 'pull_request' && github.event.pull_request.number || github.event.workflow_run.pull_requests[0].number }} + REPO: ${{ github.repository }} + run: bash scripts/report-token-usage.sh + # --------------------------------------------------------------------------- # PR synchronize (new commits pushed) → clear status, check prior concerns # --------------------------------------------------------------------------- @@ -314,6 +349,7 @@ jobs: bun-version: "1.3.12" - name: Run Claude + id: claude if: steps.prompt.outputs.skip != 'true' uses: anthropics/claude-code-action@v1 with: @@ -328,6 +364,17 @@ jobs: } } + - name: Report token usage + if: steps.claude.outcome != 'skipped' && steps.claude.outputs.execution_file != '' + continue-on-error: true + env: + GITHUB_TOKEN: ${{ secrets.AGENT_TOKEN }} + EXEC_FILE: ${{ steps.claude.outputs.execution_file }} + ENTITY_TYPE: pr + ENTITY_NUMBER: ${{ github.event.pull_request.number }} + REPO: ${{ github.repository }} + run: bash scripts/report-token-usage.sh + # --------------------------------------------------------------------------- # Comment on Claude's own PR → address the feedback # Handles both timeline comments (issue_comment) and inline diff comments @@ -380,6 +427,7 @@ jobs: bun-version: "1.3.12" - name: Run Claude + id: claude if: steps.prompt.outputs.skip != 'true' uses: anthropics/claude-code-action@v1 with: @@ -400,3 +448,14 @@ jobs: ] } } + + - name: Report token usage + if: steps.claude.outcome != 'skipped' && steps.claude.outputs.execution_file != '' + continue-on-error: true + env: + GITHUB_TOKEN: ${{ secrets.AGENT_TOKEN }} + EXEC_FILE: ${{ steps.claude.outputs.execution_file }} + ENTITY_TYPE: issue + ENTITY_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }} + REPO: ${{ github.repository }} + run: bash scripts/report-token-usage.sh diff --git a/packages/agent/src/__tests__/agent.test.ts b/packages/agent/src/__tests__/agent.test.ts index 55c5b95..949ba73 100644 --- a/packages/agent/src/__tests__/agent.test.ts +++ b/packages/agent/src/__tests__/agent.test.ts @@ -24,7 +24,7 @@ describe("runAgent", () => { })() ); - const result = await runAgent("find the Q4 report", "You are Squirrel."); + const { result } = await runAgent("find the Q4 report", "You are Squirrel."); expect(result).toBe("Here is the document you requested."); }); @@ -35,7 +35,7 @@ describe("runAgent", () => { })() ); - const result = await runAgent("find something", "system"); + const { result } = await runAgent("find something", "system"); expect(result).toBe("I was unable to generate a response."); }); @@ -47,10 +47,51 @@ describe("runAgent", () => { })() ); - const result = await runAgent("prompt", "system"); + const { result } = await runAgent("prompt", "system"); expect(result).toBe("second"); }); + it("captures token usage and cost from the result message", async () => { + vi.mocked(query).mockReturnValue( + (async function* () { + yield { + result: "answer", + total_cost_usd: 0.0042, + // real SDKResultMessage.usage uses the API's snake_case keys + usage: { + input_tokens: 800, + output_tokens: 300, + cache_read_input_tokens: 150, + cache_creation_input_tokens: 50, + }, + }; + })() + ); + + const agentResult = await runAgent("prompt", "system"); + expect(agentResult.result).toBe("answer"); + expect(agentResult.costUsd).toBeCloseTo(0.0042); + expect(agentResult.inputTokens).toBe(800); + expect(agentResult.outputTokens).toBe(300); + expect(agentResult.cacheReadTokens).toBe(150); + expect(agentResult.cacheCreationTokens).toBe(50); + }); + + it("defaults token counts and cost to zero when usage is absent", async () => { + vi.mocked(query).mockReturnValue( + (async function* () { + yield { result: "ok" }; + })() + ); + + const agentResult = await runAgent("prompt", "system"); + expect(agentResult.costUsd).toBe(0); + expect(agentResult.inputTokens).toBe(0); + expect(agentResult.outputTokens).toBe(0); + expect(agentResult.cacheReadTokens).toBe(0); + expect(agentResult.cacheCreationTokens).toBe(0); + }); + it("passes systemPrompt and model to query options", async () => { vi.mocked(query).mockReturnValue( (async function* () { @@ -126,7 +167,7 @@ describe("runAgent", () => { const resultPromise = runAgent("prompt", "system"); await vi.runAllTimersAsync(); - const result = await resultPromise; + const { result } = await resultPromise; expect(result).toBe("recovered"); expect(vi.mocked(query)).toHaveBeenCalledTimes(2); vi.useRealTimers(); diff --git a/packages/agent/src/__tests__/index.test.ts b/packages/agent/src/__tests__/index.test.ts index b2562f4..8a647ce 100644 --- a/packages/agent/src/__tests__/index.test.ts +++ b/packages/agent/src/__tests__/index.test.ts @@ -45,8 +45,12 @@ vi.mock("../slack.js", () => ({ postEphemeral: vi.fn().mockResolvedValue(undefined), })); +const mockAgentResult = (result: string, costUsd = 0) => ({ + result, inputTokens: 100, outputTokens: 50, cacheReadTokens: 10, cacheCreationTokens: 5, costUsd, +}); + vi.mock("../agent.js", () => ({ - runAgent: vi.fn().mockResolvedValue("Agent response"), + runAgent: vi.fn().mockResolvedValue(mockAgentResult("Agent response")), })); vi.mock("../prompt.js", () => ({ @@ -103,7 +107,7 @@ describe("processEvent", () => { it("fetches thread history, runs agent, and posts response", async () => { const history = [{ user: "U1", text: "earlier message" }]; vi.mocked(fetchThreadHistory).mockResolvedValueOnce(history); - vi.mocked(runAgent).mockResolvedValueOnce("Here is your answer."); + vi.mocked(runAgent).mockResolvedValueOnce(mockAgentResult("Here is your answer.")); await processEvent(baseEvent); @@ -120,7 +124,7 @@ describe("processEvent", () => { }); it("does not post or audit when agent returns the no-reply sentinel", async () => { - vi.mocked(runAgent).mockResolvedValueOnce("__NO_REPLY__"); + vi.mocked(runAgent).mockResolvedValueOnce(mockAgentResult("__NO_REPLY__")); await processEvent({ ...baseEvent, requires_discretion: true }); expect(postMessage).not.toHaveBeenCalled(); expect(mockAuditLog).not.toHaveBeenCalled(); @@ -136,7 +140,7 @@ describe("processEvent", () => { }); it("passes the agent response to postMessage", async () => { - vi.mocked(runAgent).mockResolvedValueOnce("Custom agent answer"); + vi.mocked(runAgent).mockResolvedValueOnce(mockAgentResult("Custom agent answer")); await processEvent(baseEvent); expect(postMessage).toHaveBeenCalledWith("C_CHAN", "1.0", "Custom agent answer"); }); @@ -176,9 +180,9 @@ describe("processEvent", () => { await expect(processEvent(baseEvent)).rejects.toThrow("Agent failed"); }); - it("writes audit record with channel, user, model, and duration after posting response", async () => { + it("writes audit record with channel, user, model, duration, and token usage after posting response", async () => { vi.mocked(mockConfigStore.get).mockResolvedValueOnce("claude-opus-4-6"); - vi.mocked(runAgent).mockResolvedValueOnce("The answer."); + vi.mocked(runAgent).mockResolvedValueOnce(mockAgentResult("The answer.", 0.0042)); await processEvent(baseEvent); @@ -189,6 +193,11 @@ describe("processEvent", () => { user_id: "U1", response: "The answer.", model: "claude-opus-4-6", + input_tokens: 100, + output_tokens: 50, + cache_read_tokens: 10, + cache_creation_tokens: 5, + cost_usd: "0.004200", }) ); expect(typeof mockAuditLog.mock.calls[0][0].duration_ms).toBe("number"); diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index d42773b..b187bbe 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -3,13 +3,22 @@ import * as path from "path"; import { randomUUID } from "crypto"; import type { ImageContent } from "./slack.js"; +export interface AgentResult { + result: string; + inputTokens: number; + outputTokens: number; + cacheReadTokens: number; + cacheCreationTokens: number; + costUsd: number; +} + export async function runAgent( prompt: string, systemPrompt: string, model?: string, maxTokens?: number, images?: ImageContent[], -): Promise { +): Promise { const maxRetries = parseInt(process.env.MAX_MCP_RETRIES ?? "2", 10); let lastError: Error | undefined; @@ -82,7 +91,7 @@ async function runAgentOnce( model?: string, maxTokens?: number, images?: ImageContent[], -): Promise { +): Promise { const urlFetcherPath = path.resolve( __dirname, "../../mcp-url-fetcher/dist/index.js" @@ -105,6 +114,11 @@ async function runAgentOnce( : prompt; let result = ""; + let inputTokens = 0; + let outputTokens = 0; + let cacheReadTokens = 0; + let cacheCreationTokens = 0; + let costUsd = 0; for await (const message of query({ prompt: sdkPrompt, @@ -174,8 +188,25 @@ async function runAgentOnce( })) { if ("result" in message) { result = message.result; + const msg = message as Record; + costUsd = typeof msg.total_cost_usd === "number" ? msg.total_cost_usd : 0; + // SDKResultMessage.usage is the API Usage shape — snake_case keys + const usage = msg.usage as Record | undefined; + if (usage) { + inputTokens = usage.input_tokens ?? 0; + outputTokens = usage.output_tokens ?? 0; + cacheReadTokens = usage.cache_read_input_tokens ?? 0; + cacheCreationTokens = usage.cache_creation_input_tokens ?? 0; + } } } - return result || "I was unable to generate a response."; + return { + result: result || "I was unable to generate a response.", + inputTokens, + outputTokens, + cacheReadTokens, + cacheCreationTokens, + costUsd, + }; } diff --git a/packages/agent/src/index.ts b/packages/agent/src/index.ts index 21803eb..c54b836 100644 --- a/packages/agent/src/index.ts +++ b/packages/agent/src/index.ts @@ -76,23 +76,28 @@ export async function processEvent(event: SlackEvent): Promise { : []; const start = Date.now(); - const response = await runAgent(prompt, systemPrompt, model, maxTokens, images); + const agentResult = await runAgent(prompt, systemPrompt, model, maxTokens, images); const duration_ms = Date.now() - start; - if (response.trim() === NO_REPLY_SENTINEL) { + if (agentResult.result.trim() === NO_REPLY_SENTINEL) { console.log(`[discretion] agent elected not to reply in channel ${event.channel}`); return; } - await postMessage(event.channel, event.thread_ts, response); + await postMessage(event.channel, event.thread_ts, agentResult.result); await auditLogger.log({ channel: event.channel, thread_ts: event.thread_ts, user_id: event.user, prompt, - response, + response: agentResult.result, model: model ?? null, duration_ms, + input_tokens: agentResult.inputTokens, + output_tokens: agentResult.outputTokens, + cache_read_tokens: agentResult.cacheReadTokens, + cache_creation_tokens: agentResult.cacheCreationTokens, + cost_usd: agentResult.costUsd.toFixed(6), }); } diff --git a/packages/db/drizzle/0001_audit_token_cost.sql b/packages/db/drizzle/0001_audit_token_cost.sql new file mode 100644 index 0000000..8004fdf --- /dev/null +++ b/packages/db/drizzle/0001_audit_token_cost.sql @@ -0,0 +1,9 @@ +ALTER TABLE "audit_log" ADD COLUMN "input_tokens" integer NOT NULL DEFAULT 0; +--> statement-breakpoint +ALTER TABLE "audit_log" ADD COLUMN "output_tokens" integer NOT NULL DEFAULT 0; +--> statement-breakpoint +ALTER TABLE "audit_log" ADD COLUMN "cache_read_tokens" integer NOT NULL DEFAULT 0; +--> statement-breakpoint +ALTER TABLE "audit_log" ADD COLUMN "cache_creation_tokens" integer NOT NULL DEFAULT 0; +--> statement-breakpoint +ALTER TABLE "audit_log" ADD COLUMN "cost_usd" text NOT NULL DEFAULT '0'; diff --git a/packages/db/drizzle/meta/_journal.json b/packages/db/drizzle/meta/_journal.json index 8473706..a9d46d8 100644 --- a/packages/db/drizzle/meta/_journal.json +++ b/packages/db/drizzle/meta/_journal.json @@ -8,6 +8,13 @@ "when": 1742688000000, "tag": "0000_initial", "breakpoints": true + }, + { + "idx": 1, + "version": "7", + "when": 1751760000000, + "tag": "0001_audit_token_cost", + "breakpoints": true } ] } diff --git a/packages/db/src/audit-store.ts b/packages/db/src/audit-store.ts index fddd2c4..c998121 100644 --- a/packages/db/src/audit-store.ts +++ b/packages/db/src/audit-store.ts @@ -22,6 +22,11 @@ export class PostgresAuditStore implements AuditStore { response: record.response, model: record.model ?? null, duration_ms: record.duration_ms, + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cache_read_tokens: record.cache_read_tokens, + cache_creation_tokens: record.cache_creation_tokens, + cost_usd: record.cost_usd, }); } @@ -54,6 +59,11 @@ function toRecord(row: typeof auditLog.$inferSelect): AuditRecord { response: row.response, model: row.model, duration_ms: row.duration_ms, + input_tokens: row.input_tokens, + output_tokens: row.output_tokens, + cache_read_tokens: row.cache_read_tokens, + cache_creation_tokens: row.cache_creation_tokens, + cost_usd: row.cost_usd, created_at: row.created_at.toISOString(), }; } diff --git a/packages/db/src/schema-sqlite.ts b/packages/db/src/schema-sqlite.ts index 87a7dad..d91afdb 100644 --- a/packages/db/src/schema-sqlite.ts +++ b/packages/db/src/schema-sqlite.ts @@ -22,5 +22,10 @@ export const auditLog = sqliteTable("audit_log", { response: text("response").notNull(), model: text("model"), duration_ms: integer("duration_ms").notNull(), + input_tokens: integer("input_tokens").notNull().default(0), + output_tokens: integer("output_tokens").notNull().default(0), + cache_read_tokens: integer("cache_read_tokens").notNull().default(0), + cache_creation_tokens: integer("cache_creation_tokens").notNull().default(0), + cost_usd: text("cost_usd").notNull().default("0"), created_at: text("created_at").notNull().$default(() => new Date().toISOString()), }); diff --git a/packages/db/src/schema.ts b/packages/db/src/schema.ts index 5e98a2d..0a9b377 100644 --- a/packages/db/src/schema.ts +++ b/packages/db/src/schema.ts @@ -22,5 +22,10 @@ export const auditLog = pgTable("audit_log", { response: text("response").notNull(), model: text("model"), duration_ms: integer("duration_ms").notNull(), + input_tokens: integer("input_tokens").notNull().default(0), + output_tokens: integer("output_tokens").notNull().default(0), + cache_read_tokens: integer("cache_read_tokens").notNull().default(0), + cache_creation_tokens: integer("cache_creation_tokens").notNull().default(0), + cost_usd: text("cost_usd").notNull().default("0"), created_at: timestamp("created_at", { withTimezone: true }).notNull().defaultNow(), }); diff --git a/packages/db/src/sqlite-stores.ts b/packages/db/src/sqlite-stores.ts index 32b7fff..9739c19 100644 --- a/packages/db/src/sqlite-stores.ts +++ b/packages/db/src/sqlite-stores.ts @@ -37,6 +37,11 @@ function openDb(filePath: string): SqliteDb { "response" text NOT NULL, "model" text, "duration_ms" integer NOT NULL, + "input_tokens" integer NOT NULL DEFAULT 0, + "output_tokens" integer NOT NULL DEFAULT 0, + "cache_read_tokens" integer NOT NULL DEFAULT 0, + "cache_creation_tokens" integer NOT NULL DEFAULT 0, + "cost_usd" text NOT NULL DEFAULT '0', "created_at" text NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')) ); `); @@ -127,6 +132,11 @@ export class SqliteAuditStore implements AuditStore { response: record.response, model: record.model ?? null, duration_ms: record.duration_ms, + input_tokens: record.input_tokens, + output_tokens: record.output_tokens, + cache_read_tokens: record.cache_read_tokens, + cache_creation_tokens: record.cache_creation_tokens, + cost_usd: record.cost_usd, }) .run(); } @@ -190,6 +200,11 @@ function toAuditRecord(row: typeof schema.auditLog.$inferSelect): AuditRecord { response: row.response, model: row.model, duration_ms: row.duration_ms, + input_tokens: row.input_tokens, + output_tokens: row.output_tokens, + cache_read_tokens: row.cache_read_tokens, + cache_creation_tokens: row.cache_creation_tokens, + cost_usd: row.cost_usd, created_at: row.created_at, }; } diff --git a/packages/db/src/store.ts b/packages/db/src/store.ts index ef5475b..64834c4 100644 --- a/packages/db/src/store.ts +++ b/packages/db/src/store.ts @@ -38,6 +38,11 @@ export interface AuditRecord { response: string; model: string | null; duration_ms: number; + input_tokens: number; + output_tokens: number; + cache_read_tokens: number; + cache_creation_tokens: number; + cost_usd: string; created_at: string; // ISO 8601 } diff --git a/scripts/report-token-usage.sh b/scripts/report-token-usage.sh new file mode 100755 index 0000000..d54656b --- /dev/null +++ b/scripts/report-token-usage.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +# Posts a token-usage/cost comment from a claude-code-action execution file. +# Env: EXEC_FILE, ENTITY_TYPE (issue|pr), ENTITY_NUMBER, REPO, GITHUB_TOKEN. +# Shared by every job in claude.yml — fix jq extraction here, once. +set -euo pipefail + +# The result entry's cost field is total_cost_usd (older action versions wrote cost_usd). +cost=$(jq -r '[.[] | select(.type == "result") | (.total_cost_usd // .cost_usd // 0)] | last // 0' "$EXEC_FILE" 2>/dev/null || echo "0") +in_tok=$(jq '[.[] | select(.type == "assistant") | (.message.usage.input_tokens // 0)] | add // 0' "$EXEC_FILE" 2>/dev/null || echo "0") +out_tok=$(jq '[.[] | select(.type == "assistant") | (.message.usage.output_tokens // 0)] | add // 0' "$EXEC_FILE" 2>/dev/null || echo "0") +cache_r=$(jq '[.[] | select(.type == "assistant") | (.message.usage.cache_read_input_tokens // 0)] | add // 0' "$EXEC_FILE" 2>/dev/null || echo "0") +cache_w=$(jq '[.[] | select(.type == "assistant") | (.message.usage.cache_creation_input_tokens // 0)] | add // 0' "$EXEC_FILE" 2>/dev/null || echo "0") +total_in=$(( ${in_tok:-0} + ${cache_r:-0} + ${cache_w:-0} )) +body="**Claude token usage:** ${total_in} in (${in_tok} direct + ${cache_w} cache write + ${cache_r} cache read) • ${out_tok} out • cost: \$${cost}" +gh "$ENTITY_TYPE" comment "$ENTITY_NUMBER" --repo "$REPO" --body "$body"