diff --git a/plugin/statusline.sh b/plugin/statusline.sh index 0e9db38..cef3b79 100755 --- a/plugin/statusline.sh +++ b/plugin/statusline.sh @@ -249,7 +249,8 @@ const local = localCost(); const runUsage = getRunUsage(); const orchestratorUsage = runUsage ? getOrchestratorUsage(runUsage) : undefined; const workerTokens = runUsage - ? runUsage.inputTokens + runUsage.outputTokens + runUsage.cachedTokens + ? nonNegativeNumber(runUsage.totalTokens) ?? + runUsage.inputTokens + runUsage.outputTokens + runUsage.cachedTokens : undefined; const tokenField = () => { const total = workerTokens ?? (runUsage ? undefined : localTokens()); diff --git a/src/core/pricing.ts b/src/core/pricing.ts index 8539d5a..2f9899c 100644 --- a/src/core/pricing.ts +++ b/src/core/pricing.ts @@ -179,3 +179,13 @@ export function computeSessionCost({ export function cachedInInputFor(agent: string | undefined | null): boolean { return agent === "codex"; } + +export function totalTokensFor( + agent: string | undefined | null, + usage?: SessionCostUsage, +): number { + const inputTokens = usage?.inputTokens ?? 0; + const outputTokens = usage?.outputTokens ?? 0; + const cachedTokens = usage?.cachedTokens ?? 0; + return inputTokens + outputTokens + (cachedInInputFor(agent) ? 0 : cachedTokens); +} diff --git a/src/core/run-usage.ts b/src/core/run-usage.ts index cbfc0b5..b0c59b2 100644 --- a/src/core/run-usage.ts +++ b/src/core/run-usage.ts @@ -1,4 +1,4 @@ -import { cachedInInputFor, computeSessionCost } from "./pricing.js"; +import { cachedInInputFor, computeSessionCost, totalTokensFor } from "./pricing.js"; import { isActiveStatus, type Session } from "./session.js"; export interface RunAttribution { @@ -17,6 +17,7 @@ export interface RunUsageSummary { inputTokens: number; outputTokens: number; cachedTokens: number; + totalTokens: number; costUsd: number; sessionCount: number; activeSessionCount: number; @@ -28,6 +29,7 @@ export interface RunUsageSummary { inputTokens: number; outputTokens: number; cachedTokens: number; + totalTokens: number; sources: Array<{ nativeId: string; costUsd: number }>; }; total: { costUsd: number }; @@ -56,6 +58,7 @@ export function aggregateRunUsage( inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, costUsd: 0, sessionCount: 0, activeSessionCount: 0, @@ -67,6 +70,7 @@ export function aggregateRunUsage( inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, sources: [], }, total: { costUsd: 0 }, @@ -79,6 +83,7 @@ export function aggregateRunUsage( target.inputTokens += usage?.inputTokens ?? 0; target.outputTokens += usage?.outputTokens ?? 0; target.cachedTokens += usage?.cachedTokens ?? 0; + target.totalTokens += totalTokensFor(session.agent, usage); const costUsd = computeSessionCost({ model: session.model, diff --git a/src/store/sessions.ts b/src/store/sessions.ts index ec35cd4..9de419d 100644 --- a/src/store/sessions.ts +++ b/src/store/sessions.ts @@ -1,7 +1,7 @@ import type { DatabaseSync } from "node:sqlite"; import { type Session, type SessionStatus, type AgentId, isActiveStatus } from "../core/session.js"; import type { FailureInfo } from "../core/errors.js"; -import { cachedInInputFor, computeSessionCost } from "../core/pricing.js"; +import { cachedInInputFor, computeSessionCost, totalTokensFor } from "../core/pricing.js"; import type { UsageQueryParams, UsageQueryResult, UsageMetricBucket, UsageTotals } from "../daemon/protocol.js"; export interface SessionRow { @@ -443,7 +443,7 @@ export class SessionStore { const inputTokens = row.usage_input_tokens ?? 0; const outputTokens = row.usage_output_tokens ?? 0; const cachedTokens = row.usage_cached_tokens ?? 0; - const totalTokens = inputTokens + outputTokens + cachedTokens; + const totalTokens = totalTokensFor(row.agent, { inputTokens, outputTokens, cachedTokens }); const calculatedCost = computeSessionCost({ model: row.model, @@ -522,7 +522,7 @@ export class SessionStore { const inputTokens = row.input_tokens; const outputTokens = row.output_tokens; const cachedTokens = row.cached_tokens; - const totalTokens = inputTokens + outputTokens + cachedTokens; + const totalTokens = totalTokensFor("claude", { inputTokens, outputTokens, cachedTokens }); const cost = row.cost; totals.sessionCount++; diff --git a/tests/pricing.test.ts b/tests/pricing.test.ts index 787ba36..cb16f26 100644 --- a/tests/pricing.test.ts +++ b/tests/pricing.test.ts @@ -1,5 +1,11 @@ import { describe, expect, it } from "vitest"; -import { cachedInInputFor, computeSessionCost, MODEL_PRICES, resolveModelPrice } from "../src/core/pricing.js"; +import { + cachedInInputFor, + computeSessionCost, + MODEL_PRICES, + resolveModelPrice, + totalTokensFor, +} from "../src/core/pricing.js"; describe("computeSessionCost", () => { it("uses a valid reported cost, including zero, without recalculating", () => { @@ -242,3 +248,12 @@ describe("computeSessionCost", () => { expect(resolveModelPrice("CLAUDE-SONNET-4-6")).toEqual(MODEL_PRICES["claude-sonnet-4-6"]); }); }); + +describe("totalTokensFor", () => { + it("counts cached tokens once for Codex and separately for Claude", () => { + const usage = { inputTokens: 1_000, outputTokens: 100, cachedTokens: 800 }; + + expect(totalTokensFor("codex", usage)).toBe(1_100); + expect(totalTokensFor("claude", usage)).toBe(1_900); + }); +}); diff --git a/tests/statusline.test.ts b/tests/statusline.test.ts index 4b653b1..bc21008 100644 --- a/tests/statusline.test.ts +++ b/tests/statusline.test.ts @@ -265,6 +265,49 @@ describe("Claude statusline", () => { expect(stripAnsi(small.output)).toContain("980 tok"); }); + it("uses a valid run totalTokens value for the tok field", async () => { + const result = await render({ + payload: payload(0), + runId: "run-cached-tokens", + usage: { + runId: "run-cached-tokens", + inputTokens: 1_000, + outputTokens: 100, + cachedTokens: 800, + totalTokens: 1_100, + costUsd: 0, + sessionCount: 1, + activeSessionCount: 1, + costComplete: true, + sessionsWithoutCost: 0, + }, + }); + + expect(stripAnsi(result.output)).toContain("1.1k tok"); + expect(stripAnsi(result.output)).not.toContain("1.9k tok"); + }); + + it("falls back to the token sum when totalTokens is invalid", async () => { + const result = await render({ + payload: payload(0), + runId: "run-invalid-token-total", + usage: { + runId: "run-invalid-token-total", + inputTokens: 1_000, + outputTokens: 100, + cachedTokens: 800, + totalTokens: -1, + costUsd: 0, + sessionCount: 1, + activeSessionCount: 1, + costComplete: true, + sessionsWithoutCost: 0, + }, + }); + + expect(stripAnsi(result.output)).toContain("1.9k tok"); + }); + it("marks a partial aggregate even when the local cost is zero", async () => { const result = await render({ payload: payload(0), diff --git a/tests/usage-cli.test.ts b/tests/usage-cli.test.ts index 528fcc1..94f3599 100644 --- a/tests/usage-cli.test.ts +++ b/tests/usage-cli.test.ts @@ -18,6 +18,7 @@ const runSummary = { inputTokens: 100, outputTokens: 20, cachedTokens: 5, + totalTokens: 125, costUsd: 0.25, sessionCount: 1, activeSessionCount: 0, @@ -29,6 +30,7 @@ const runSummary = { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, sources: [], }, total: { costUsd: 0.25 }, diff --git a/tests/usage-daemon.test.ts b/tests/usage-daemon.test.ts index b35b5c8..581cd3f 100644 --- a/tests/usage-daemon.test.ts +++ b/tests/usage-daemon.test.ts @@ -93,6 +93,7 @@ describe("usage daemon methods", () => { inputTokens: 1_300, outputTokens: 850, cachedTokens: 320, + totalTokens: 2_170, costUsd: 0.5, sessionCount: 2, activeSessionCount: 1, @@ -104,6 +105,7 @@ describe("usage daemon methods", () => { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, sources: [], }, total: { costUsd: 0.5 }, diff --git a/tests/usage-query.test.ts b/tests/usage-query.test.ts index 54e4abe..eb614de 100644 --- a/tests/usage-query.test.ts +++ b/tests/usage-query.test.ts @@ -125,6 +125,40 @@ describe("SessionStore.queryUsage", () => { expect(worker).toMatchObject({ sessionCount: 2, costUsd: 1.25 }); }); + it("uses harness-aware totals in every usage bucket", () => { + store.create( + makeSession("codex", { + agent: "codex", + model: "gpt-5.6-luna", + repository: "/dev/codedeck", + origin: null, + usage: { inputTokens: 1_000, outputTokens: 100, cachedTokens: 800 }, + }), + ); + store.create( + makeSession("claude", { + agent: "claude", + model: "claude-sonnet-4-6", + repository: "/dev/codedeck", + origin: "open", + usage: { inputTokens: 1_000, outputTokens: 100, cachedTokens: 800 }, + }), + ); + + const result = store.queryUsage({ period: "all" }); + + expect(result.totals.totalTokens).toBe(3_000); + expect(result.byDay[0]?.totalTokens).toBe(3_000); + expect(result.byRepository.find((bucket) => bucket.key === "codedeck")?.totalTokens).toBe(3_000); + expect(result.byModel.find((bucket) => bucket.key === "gpt-5.6-luna")?.totalTokens).toBe(1_100); + expect(result.byModel.find((bucket) => bucket.key === "claude-sonnet-4-6")?.totalTokens).toBe(1_900); + expect(result.byAgent.find((bucket) => bucket.key === "codex")?.totalTokens).toBe(1_100); + expect(result.byAgent.find((bucket) => bucket.key === "claude")?.totalTokens).toBe(1_900); + expect(result.byRun.find((bucket) => bucket.key === "run-1")?.totalTokens).toBe(3_000); + expect(result.byOrigin.find((bucket) => bucket.key === "worker")?.totalTokens).toBe(1_100); + expect(result.byOrigin.find((bucket) => bucket.key === "orchestrator")?.totalTokens).toBe(1_900); + }); + it("merges in-range legacy usage into analytics without listing it as a session", () => { const endedAt = new Date(2026, 8, 7, 12).toISOString(); const outsideRange = new Date(2026, 7, 31, 12).toISOString(); diff --git a/tests/usage.test.ts b/tests/usage.test.ts index 76d1b51..67546cb 100644 --- a/tests/usage.test.ts +++ b/tests/usage.test.ts @@ -31,6 +31,7 @@ describe("aggregateRunUsage", () => { inputTokens: 1_300, outputTokens: 850, cachedTokens: 320, + totalTokens: 2_150, costUsd: 0.5, sessionCount: 2, activeSessionCount: 0, @@ -42,6 +43,7 @@ describe("aggregateRunUsage", () => { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, sources: [], }, total: { costUsd: 0.5 }, @@ -65,6 +67,7 @@ describe("aggregateRunUsage", () => { inputTokens: 300, outputTokens: 150, cachedTokens: 30, + totalTokens: 450, costUsd: 0.42, sessionCount: 2, activeSessionCount: 0, @@ -76,6 +79,7 @@ describe("aggregateRunUsage", () => { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, sources: [], }, total: { costUsd: 0.42 }, @@ -142,6 +146,7 @@ describe("aggregateRunUsage", () => { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, costUsd: 0, sessionCount: 0, activeSessionCount: 0, @@ -153,6 +158,7 @@ describe("aggregateRunUsage", () => { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, sources: [], }, total: { costUsd: 0 }, @@ -184,6 +190,7 @@ describe("aggregateRunUsage", () => { inputTokens: 0, outputTokens: 0, cachedTokens: 0, + totalTokens: 0, costUsd: 0.5, sessionCount: 2, activeSessionCount: 0, @@ -195,6 +202,7 @@ describe("aggregateRunUsage", () => { inputTokens: 10, outputTokens: 20, cachedTokens: 3, + totalTokens: 33, sources: [ { nativeId: "X", costUsd: 3 }, { nativeId: "Y", costUsd: 0.8 }, @@ -250,4 +258,21 @@ describe("aggregateRunUsage", () => { expect(summary.total.costUsd).toBe(1); expect(summary.costComplete).toBe(true); }); + + it("counts cached tokens once for Codex workers and separately for Claude orchestrators", () => { + const summary = aggregateRunUsage("run-tokens", [ + makeSession("codex", { + agent: "codex", + usage: { inputTokens: 1_000, outputTokens: 100, cachedTokens: 800, cost: 0 }, + }), + makeSession("claude", { + agent: "claude", + origin: "open", + usage: { inputTokens: 1_000, outputTokens: 100, cachedTokens: 800, cost: 0 }, + }), + ]); + + expect(summary.totalTokens).toBe(1_100); + expect(summary.orchestrator.totalTokens).toBe(1_900); + }); });