Skip to content

Commit ca67b12

Browse files
committed
feat(sessions): surface per-task cost estimate
The Claude Agent SDK already reports a per-turn `total_cost_usd`, which flowed into `ContextUsage.cost` but was never rendered. Sum it across turns (each result reports only its own spend, matching how per-turn token usage is already accumulated) and show the running total in the context usage indicator chip and breakdown popover. Codex sessions emit tokens without a cost, so `cost` stays null and the UI hides it. Generated-By: PostHog Code Task-Id: f6d673c3-965e-4598-a4bd-663f2651a5a5
1 parent fc33dfb commit ca67b12

5 files changed

Lines changed: 176 additions & 14 deletions

File tree

packages/core/src/sessions/contextUsage.test.ts

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,31 @@ function usageUpdateEvent(used: number, size: number): AcpMessage {
1717
};
1818
}
1919

20+
function costUsageUpdateEvent(
21+
used: number,
22+
size: number,
23+
amount: number,
24+
currency = "USD",
25+
): AcpMessage {
26+
return {
27+
type: "acp_message",
28+
ts: 1,
29+
message: {
30+
jsonrpc: "2.0",
31+
method: "session/update",
32+
params: {
33+
sessionId: "s1",
34+
update: {
35+
sessionUpdate: "usage_update",
36+
used,
37+
size,
38+
cost: { amount, currency },
39+
},
40+
},
41+
},
42+
};
43+
}
44+
2045
function sizelessUsageUpdateEvent(used: number): AcpMessage {
2146
return {
2247
type: "acp_message",
@@ -119,6 +144,30 @@ describe("extractContextUsage", () => {
119144
expect(result?.breakdown?.conversation).toBe(45_500);
120145
});
121146

147+
it("reports null cost when no update carries a cost", () => {
148+
const result = extractContextUsage([usageUpdateEvent(50_000, 200_000)]);
149+
expect(result?.cost).toBeNull();
150+
});
151+
152+
it("surfaces the cost from a single turn", () => {
153+
const result = extractContextUsage([
154+
costUsageUpdateEvent(50_000, 200_000, 0.42),
155+
]);
156+
expect(result?.cost).toEqual({ amount: 0.42, currency: "USD" });
157+
});
158+
159+
it("sums cost across turns since each result reports only its own spend", () => {
160+
const result = extractContextUsage([
161+
costUsageUpdateEvent(40_000, 200_000, 0.4),
162+
costUsageUpdateEvent(90_000, 200_000, 0.35),
163+
costUsageUpdateEvent(120_000, 200_000, 0.25),
164+
]);
165+
// Context occupancy tracks the newest turn; cost accrues across all of them.
166+
expect(result?.used).toBe(120_000);
167+
expect(result?.cost?.amount).toBeCloseTo(1.0, 10);
168+
expect(result?.cost?.currency).toBe("USD");
169+
});
170+
122171
it("tolerates the double-underscore method prefix from extNotification", () => {
123172
const result = extractContextUsage([
124173
usageUpdateEvent(50_000, 200_000),
@@ -180,6 +229,28 @@ describe("createContextUsageTracker", () => {
180229
expect(tracker.update([earlier])?.used).toBe(50_000);
181230
});
182231

232+
it("accumulates cost only over newly appended turns", () => {
233+
const tracker = createContextUsageTracker();
234+
const first = costUsageUpdateEvent(40_000, 200_000, 0.4);
235+
236+
expect(tracker.update([first])?.cost?.amount).toBeCloseTo(0.4, 10);
237+
238+
const result = tracker.update([
239+
first,
240+
costUsageUpdateEvent(90_000, 200_000, 0.35),
241+
]);
242+
expect(result?.cost?.amount).toBeCloseTo(0.75, 10);
243+
});
244+
245+
it("matches the batch extractor for a cost-bearing log", () => {
246+
const tracker = createContextUsageTracker();
247+
const events = [
248+
costUsageUpdateEvent(40_000, 200_000, 0.4),
249+
costUsageUpdateEvent(90_000, 200_000, 0.35),
250+
];
251+
expect(tracker.update(events)).toEqual(extractContextUsage(events));
252+
});
253+
183254
it("rebuilds when the tail changes at the same length", () => {
184255
const tracker = createContextUsageTracker();
185256
const first = usageUpdateEvent(50_000, 200_000);

packages/core/src/sessions/contextUsage.ts

Lines changed: 70 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,33 @@ export interface ContextUsage {
1515
used: number;
1616
size: number;
1717
percentage: number;
18+
/**
19+
* Cumulative estimated cost of the whole session, summed across turns. `used`
20+
* is a point-in-time snapshot of resident context, but each turn's cost is a
21+
* spend that accrues, so cost is totalled rather than replaced. `null` when
22+
* no turn reported a cost (e.g. codex, which emits tokens without a cost).
23+
*/
1824
cost: { amount: number; currency: string } | null;
1925
breakdown: ContextBreakdown | null;
2026
}
2127

22-
type ContextUsageAggregate = Omit<ContextUsage, "breakdown">;
28+
type ContextUsageAggregate = Omit<ContextUsage, "breakdown" | "cost">;
2329

2430
export function extractContextUsage(events: AcpMessage[]): ContextUsage | null {
2531
let aggregate: ContextUsageAggregate | null = null;
2632
let breakdown: ContextBreakdown | null = null;
33+
let costAmount: number | null = null;
34+
let costCurrency = "USD";
2735

36+
// Cost sums over every turn, so this can't early-break once the newest
37+
// aggregate/breakdown is found — it walks the full log.
2838
for (let i = events.length - 1; i >= 0; i--) {
2939
const msg = events[i].message;
40+
const cost = extractCost(msg);
41+
if (cost) {
42+
costAmount = (costAmount ?? 0) + cost.amount;
43+
costCurrency = cost.currency;
44+
}
3045
if (!aggregate) {
3146
aggregate = extractAggregate(msg);
3247
} else if (aggregate.size <= 0) {
@@ -37,36 +52,55 @@ export function extractContextUsage(events: AcpMessage[]): ContextUsage | null {
3752
if (!breakdown) {
3853
breakdown = extractBreakdown(msg);
3954
}
40-
if (aggregate && aggregate.size > 0 && breakdown) break;
4155
}
4256

4357
if (!aggregate) return null;
44-
return { ...aggregate, breakdown };
58+
return { ...aggregate, cost: toCost(costAmount, costCurrency), breakdown };
4559
}
4660

4761
interface ContextUsageState {
4862
aggregate: ContextUsageAggregate | null;
63+
costAmount: number | null;
64+
costCurrency: string;
4965
breakdown: ContextBreakdown | null;
5066
}
5167

5268
export function createContextUsageTracker() {
5369
return createAppendOnlyTracker<ContextUsageState, ContextUsage | null>({
54-
init: () => ({ aggregate: null, breakdown: null }),
70+
init: () => ({
71+
aggregate: null,
72+
costAmount: null,
73+
costCurrency: "USD",
74+
breakdown: null,
75+
}),
5576
processEvent: (state, event) => {
5677
const msg = event.message;
5778
const next = extractAggregate(msg);
5879
if (next) {
5980
state.aggregate = withCarriedSize(next, state.aggregate);
6081
}
82+
const cost = extractCost(msg);
83+
if (cost) {
84+
state.costAmount = (state.costAmount ?? 0) + cost.amount;
85+
state.costCurrency = cost.currency;
86+
}
6187
state.breakdown = extractBreakdown(msg) ?? state.breakdown;
6288
},
6389
getResult: (state) =>
6490
state.aggregate
65-
? { ...state.aggregate, breakdown: state.breakdown }
91+
? {
92+
...state.aggregate,
93+
cost: toCost(state.costAmount, state.costCurrency),
94+
breakdown: state.breakdown,
95+
}
6696
: null,
6797
});
6898
}
6999

100+
function toCost(amount: number | null, currency: string): ContextUsage["cost"] {
101+
return amount != null ? { amount, currency } : null;
102+
}
103+
70104
/**
71105
* An update that omits `size` must not wipe a previously known context window
72106
* (codex reports `modelContextWindow` intermittently), so keep the last known
@@ -116,11 +150,38 @@ function extractAggregate(
116150
const size = typeof update.size === "number" ? update.size : 0;
117151
const percentage =
118152
size > 0 ? Math.min(100, Math.round((update.used / size) * 100)) : 0;
153+
return { used: update.used, size, percentage };
154+
}
155+
}
156+
return null;
157+
}
158+
159+
function extractCost(
160+
msg: AcpMessage["message"],
161+
): { amount: number; currency: string } | null {
162+
if (
163+
"method" in msg &&
164+
msg.method === "session/update" &&
165+
!("id" in msg) &&
166+
"params" in msg
167+
) {
168+
const params = msg.params as
169+
| {
170+
update?: {
171+
sessionUpdate?: string;
172+
cost?: { amount: number; currency: string } | null;
173+
};
174+
}
175+
| undefined;
176+
const update = params?.update;
177+
if (
178+
update?.sessionUpdate === "usage_update" &&
179+
update.cost &&
180+
typeof update.cost.amount === "number"
181+
) {
119182
return {
120-
used: update.used,
121-
size,
122-
percentage,
123-
cost: update.cost ?? null,
183+
amount: update.cost.amount,
184+
currency: update.cost.currency ?? "USD",
124185
};
125186
}
126187
}

packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import {
22
CONTEXT_CATEGORIES,
3+
formatCostUsd,
34
formatTokensCompact,
45
getOverallUsageColor,
56
} from "@posthog/ui/features/sessions/contextColors";
@@ -13,7 +14,7 @@ interface ContextBreakdownPopoverProps {
1314
export function ContextBreakdownPopover({
1415
usage,
1516
}: ContextBreakdownPopoverProps) {
16-
const { used, size, percentage, breakdown } = usage;
17+
const { used, size, percentage, cost, breakdown } = usage;
1718
const fillColor = getOverallUsageColor(percentage);
1819
// The context window can be unknown (size 0) — show just the token count
1920
// rather than a misleading "~X / 0 tokens · 0% full".
@@ -71,6 +72,19 @@ export function ContextBreakdownPopover({
7172
Detailed breakdown available after the first response.
7273
</Text>
7374
)}
75+
76+
{cost && (
77+
<Flex
78+
align="center"
79+
justify="between"
80+
className="border-(--gray-4) border-t pt-2 text-[13px]"
81+
>
82+
<Text className="text-(--gray-11)">Estimated cost</Text>
83+
<Text className="font-medium text-(--gray-12) tabular-nums">
84+
{formatCostUsd(cost.amount)}
85+
</Text>
86+
</Flex>
87+
)}
7488
</Flex>
7589
);
7690
}

packages/ui/src/features/sessions/components/ContextUsageIndicator.tsx

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import {
2+
formatCostUsd,
23
formatTokensCompact,
34
getOverallUsageColor,
45
} from "@posthog/ui/features/sessions/contextColors";
@@ -18,12 +19,18 @@ interface ContextUsageIndicatorProps {
1819
export function ContextUsageIndicator({ usage }: ContextUsageIndicatorProps) {
1920
if (!usage) return null;
2021

21-
const { used, size, percentage } = usage;
22+
const { used, size, percentage, cost } = usage;
2223
// The context window can be unknown (size 0) — show just the token count
2324
// rather than a misleading "X/0 · 0%".
2425
const hasSize = size > 0;
2526
const strokeDashoffset = CIRCUMFERENCE - (percentage / 100) * CIRCUMFERENCE;
2627
const color = getOverallUsageColor(percentage);
28+
const tokenLabel = hasSize
29+
? `${formatTokensCompact(used)}/${formatTokensCompact(size)} · ${percentage}%`
30+
: formatTokensCompact(used);
31+
const label = cost
32+
? `${tokenLabel} · ${formatCostUsd(cost.amount)}`
33+
: tokenLabel;
2734

2835
return (
2936
<Popover.Root>
@@ -66,9 +73,7 @@ export function ContextUsageIndicator({ usage }: ContextUsageIndicatorProps) {
6673
/>
6774
</svg>
6875
<Text className="text-[13px] text-muted-foreground tabular-nums">
69-
{hasSize
70-
? `${formatTokensCompact(used)}/${formatTokensCompact(size)} · ${percentage}%`
71-
: formatTokensCompact(used)}
76+
{label}
7277
</Text>
7378
</Flex>
7479
</button>

packages/ui/src/features/sessions/contextColors.ts

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,3 +28,14 @@ export function formatTokensCompact(tokens: number): string {
2828
if (tokens >= 1000) return `${Math.round(tokens / 1000)}K`;
2929
return tokens.toString();
3030
}
31+
32+
/**
33+
* Formats a USD cost estimate for display. Sub-cent amounts collapse to
34+
* `<$0.01` so a non-zero spend never reads as free; everything else shows two
35+
* decimals ($0.42, $12.34).
36+
*/
37+
export function formatCostUsd(amount: number): string {
38+
if (amount <= 0) return "$0.00";
39+
if (amount < 0.01) return "<$0.01";
40+
return `$${amount.toFixed(2)}`;
41+
}

0 commit comments

Comments
 (0)