diff --git a/packages/core/src/sessions/contextUsage.test.ts b/packages/core/src/sessions/contextUsage.test.ts index b753f257a0..4a7421e92f 100644 --- a/packages/core/src/sessions/contextUsage.test.ts +++ b/packages/core/src/sessions/contextUsage.test.ts @@ -17,6 +17,31 @@ function usageUpdateEvent(used: number, size: number): AcpMessage { }; } +function costUsageUpdateEvent( + used: number, + size: number, + amount: number, + currency = "USD", +): AcpMessage { + return { + type: "acp_message", + ts: 1, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: "s1", + update: { + sessionUpdate: "usage_update", + used, + size, + cost: { amount, currency }, + }, + }, + }, + }; +} + function sizelessUsageUpdateEvent(used: number): AcpMessage { return { type: "acp_message", @@ -119,6 +144,30 @@ describe("extractContextUsage", () => { expect(result?.breakdown?.conversation).toBe(45_500); }); + it("reports null cost when no update carries a cost", () => { + const result = extractContextUsage([usageUpdateEvent(50_000, 200_000)]); + expect(result?.cost).toBeNull(); + }); + + it("surfaces the cost from a single turn", () => { + const result = extractContextUsage([ + costUsageUpdateEvent(50_000, 200_000, 0.42), + ]); + expect(result?.cost).toEqual({ amount: 0.42, currency: "USD" }); + }); + + it("sums cost across turns since each result reports only its own spend", () => { + const result = extractContextUsage([ + costUsageUpdateEvent(40_000, 200_000, 0.4), + costUsageUpdateEvent(90_000, 200_000, 0.35), + costUsageUpdateEvent(120_000, 200_000, 0.25), + ]); + // Context occupancy tracks the newest turn; cost accrues across all of them. + expect(result?.used).toBe(120_000); + expect(result?.cost?.amount).toBeCloseTo(1.0, 10); + expect(result?.cost?.currency).toBe("USD"); + }); + it("tolerates the double-underscore method prefix from extNotification", () => { const result = extractContextUsage([ usageUpdateEvent(50_000, 200_000), @@ -180,6 +229,28 @@ describe("createContextUsageTracker", () => { expect(tracker.update([earlier])?.used).toBe(50_000); }); + it("accumulates cost only over newly appended turns", () => { + const tracker = createContextUsageTracker(); + const first = costUsageUpdateEvent(40_000, 200_000, 0.4); + + expect(tracker.update([first])?.cost?.amount).toBeCloseTo(0.4, 10); + + const result = tracker.update([ + first, + costUsageUpdateEvent(90_000, 200_000, 0.35), + ]); + expect(result?.cost?.amount).toBeCloseTo(0.75, 10); + }); + + it("matches the batch extractor for a cost-bearing log", () => { + const tracker = createContextUsageTracker(); + const events = [ + costUsageUpdateEvent(40_000, 200_000, 0.4), + costUsageUpdateEvent(90_000, 200_000, 0.35), + ]; + expect(tracker.update(events)).toEqual(extractContextUsage(events)); + }); + it("rebuilds when the tail changes at the same length", () => { const tracker = createContextUsageTracker(); const first = usageUpdateEvent(50_000, 200_000); diff --git a/packages/core/src/sessions/contextUsage.ts b/packages/core/src/sessions/contextUsage.ts index 6d37bc156c..df3f4001e3 100644 --- a/packages/core/src/sessions/contextUsage.ts +++ b/packages/core/src/sessions/contextUsage.ts @@ -15,18 +15,28 @@ export interface ContextUsage { used: number; size: number; percentage: number; + /** Cumulative estimated session cost, summed across turns; `null` if none reported (e.g. codex). */ cost: { amount: number; currency: string } | null; breakdown: ContextBreakdown | null; } -type ContextUsageAggregate = Omit; +type ContextUsageAggregate = Omit; export function extractContextUsage(events: AcpMessage[]): ContextUsage | null { let aggregate: ContextUsageAggregate | null = null; let breakdown: ContextBreakdown | null = null; + let costAmount: number | null = null; + let costCurrency = "USD"; + // Cost sums over every turn, so this can't early-break once the newest + // aggregate/breakdown is found — it walks the full log. for (let i = events.length - 1; i >= 0; i--) { const msg = events[i].message; + const cost = extractCost(msg); + if (cost) { + costAmount = (costAmount ?? 0) + cost.amount; + costCurrency = cost.currency; + } if (!aggregate) { aggregate = extractAggregate(msg); } else if (aggregate.size <= 0) { @@ -37,36 +47,55 @@ export function extractContextUsage(events: AcpMessage[]): ContextUsage | null { if (!breakdown) { breakdown = extractBreakdown(msg); } - if (aggregate && aggregate.size > 0 && breakdown) break; } if (!aggregate) return null; - return { ...aggregate, breakdown }; + return { ...aggregate, cost: toCost(costAmount, costCurrency), breakdown }; } interface ContextUsageState { aggregate: ContextUsageAggregate | null; + costAmount: number | null; + costCurrency: string; breakdown: ContextBreakdown | null; } export function createContextUsageTracker() { return createAppendOnlyTracker({ - init: () => ({ aggregate: null, breakdown: null }), + init: () => ({ + aggregate: null, + costAmount: null, + costCurrency: "USD", + breakdown: null, + }), processEvent: (state, event) => { const msg = event.message; const next = extractAggregate(msg); if (next) { state.aggregate = withCarriedSize(next, state.aggregate); } + const cost = extractCost(msg); + if (cost) { + state.costAmount = (state.costAmount ?? 0) + cost.amount; + state.costCurrency = cost.currency; + } state.breakdown = extractBreakdown(msg) ?? state.breakdown; }, getResult: (state) => state.aggregate - ? { ...state.aggregate, breakdown: state.breakdown } + ? { + ...state.aggregate, + cost: toCost(state.costAmount, state.costCurrency), + breakdown: state.breakdown, + } : null, }); } +function toCost(amount: number | null, currency: string): ContextUsage["cost"] { + return amount != null ? { amount, currency } : null; +} + /** * An update that omits `size` must not wipe a previously known context window * (codex reports `modelContextWindow` intermittently), so keep the last known @@ -116,11 +145,38 @@ function extractAggregate( const size = typeof update.size === "number" ? update.size : 0; const percentage = size > 0 ? Math.min(100, Math.round((update.used / size) * 100)) : 0; + return { used: update.used, size, percentage }; + } + } + return null; +} + +function extractCost( + msg: AcpMessage["message"], +): { amount: number; currency: string } | null { + if ( + "method" in msg && + msg.method === "session/update" && + !("id" in msg) && + "params" in msg + ) { + const params = msg.params as + | { + update?: { + sessionUpdate?: string; + cost?: { amount: number; currency: string } | null; + }; + } + | undefined; + const update = params?.update; + if ( + update?.sessionUpdate === "usage_update" && + update.cost && + typeof update.cost.amount === "number" + ) { return { - used: update.used, - size, - percentage, - cost: update.cost ?? null, + amount: update.cost.amount, + currency: update.cost.currency ?? "USD", }; } } diff --git a/packages/shared/src/flags.ts b/packages/shared/src/flags.ts index d692da1fda..8600cec7d0 100644 --- a/packages/shared/src/flags.ts +++ b/packages/shared/src/flags.ts @@ -21,3 +21,5 @@ export const GLM_MODEL_FLAG = "posthog-code-glm-model"; export const SPOKEN_NARRATION_FLAG = "posthog-code-spoken-narration"; // Gates importing and relaying local MCP servers into cloud task runs. export const LOCAL_MCP_IMPORT_FLAG = "posthog-code-local-mcp-import"; +/** Per-task estimated cost readout in the context usage indicator. */ +export const TASK_COST_FLAG = "posthog-code-task-cost"; diff --git a/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx b/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx index 316ca05aa0..4269e1504f 100644 --- a/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx +++ b/packages/ui/src/features/sessions/components/ContextBreakdownPopover.tsx @@ -1,5 +1,6 @@ import { CONTEXT_CATEGORIES, + formatCostUsd, formatTokensCompact, getOverallUsageColor, } from "@posthog/ui/features/sessions/contextColors"; @@ -8,12 +9,14 @@ import { Flex, Text } from "@radix-ui/themes"; interface ContextBreakdownPopoverProps { usage: ContextUsage; + showCost?: boolean; } export function ContextBreakdownPopover({ usage, + showCost = false, }: ContextBreakdownPopoverProps) { - const { used, size, percentage, breakdown } = usage; + const { used, size, percentage, cost, breakdown } = usage; const fillColor = getOverallUsageColor(percentage); // The context window can be unknown (size 0) — show just the token count // rather than a misleading "~X / 0 tokens · 0% full". @@ -71,6 +74,19 @@ export function ContextBreakdownPopover({ Detailed breakdown available after the first response. )} + + {showCost && cost && ( + + Estimated cost + + {formatCostUsd(cost.amount)} + + + )} ); } diff --git a/packages/ui/src/features/sessions/components/ContextUsageIndicator.test.tsx b/packages/ui/src/features/sessions/components/ContextUsageIndicator.test.tsx index 49945a0f7d..08013c5181 100644 --- a/packages/ui/src/features/sessions/components/ContextUsageIndicator.test.tsx +++ b/packages/ui/src/features/sessions/components/ContextUsageIndicator.test.tsx @@ -1,9 +1,18 @@ import type { ContextUsage } from "@posthog/ui/features/sessions/hooks/useContextUsage"; import { Theme } from "@radix-ui/themes"; import { render, screen } from "@testing-library/react"; -import { describe, expect, it } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { ContextUsageIndicator } from "./ContextUsageIndicator"; +const flagState = vi.hoisted(() => ({ enabled: false })); +vi.mock("@posthog/ui/features/feature-flags/useFeatureFlag", () => ({ + useFeatureFlag: () => flagState.enabled, +})); + +beforeEach(() => { + flagState.enabled = false; +}); + function usage(overrides?: Partial): ContextUsage { return { used: 50_000, @@ -53,6 +62,18 @@ describe("ContextUsageIndicator", () => { ).toBeInTheDocument(); }); + it("appends the estimated cost to the label when the flag is enabled", () => { + flagState.enabled = true; + render( + + + , + ); + expect(screen.getByText(/50K\/200K · 25% · \$0\.42/)).toBeInTheDocument(); + }); + it("renders a finite stroke offset at 0% (no NaN/Infinity)", () => { const { container } = render( diff --git a/packages/ui/src/features/sessions/components/ContextUsageIndicator.tsx b/packages/ui/src/features/sessions/components/ContextUsageIndicator.tsx index 1ae3ac46a1..c088a620be 100644 --- a/packages/ui/src/features/sessions/components/ContextUsageIndicator.tsx +++ b/packages/ui/src/features/sessions/components/ContextUsageIndicator.tsx @@ -1,4 +1,7 @@ +import { TASK_COST_FLAG } from "@posthog/shared"; +import { useFeatureFlag } from "@posthog/ui/features/feature-flags/useFeatureFlag"; import { + formatCostUsd, formatTokensCompact, getOverallUsageColor, } from "@posthog/ui/features/sessions/contextColors"; @@ -16,14 +19,23 @@ interface ContextUsageIndicatorProps { } export function ContextUsageIndicator({ usage }: ContextUsageIndicatorProps) { + const costEnabled = useFeatureFlag(TASK_COST_FLAG) || import.meta.env.DEV; + if (!usage) return null; - const { used, size, percentage } = usage; + const { used, size, percentage, cost } = usage; // The context window can be unknown (size 0) — show just the token count // rather than a misleading "X/0 · 0%". const hasSize = size > 0; const strokeDashoffset = CIRCUMFERENCE - (percentage / 100) * CIRCUMFERENCE; const color = getOverallUsageColor(percentage); + const showCost = costEnabled && cost !== null; + const tokenLabel = hasSize + ? `${formatTokensCompact(used)}/${formatTokensCompact(size)} · ${percentage}%` + : formatTokensCompact(used); + const label = showCost + ? `${tokenLabel} · ${formatCostUsd(cost.amount)}` + : tokenLabel; return ( @@ -66,15 +78,13 @@ export function ContextUsageIndicator({ usage }: ContextUsageIndicatorProps) { /> - {hasSize - ? `${formatTokensCompact(used)}/${formatTokensCompact(size)} · ${percentage}%` - : formatTokensCompact(used)} + {label} - + ); diff --git a/packages/ui/src/features/sessions/contextColors.ts b/packages/ui/src/features/sessions/contextColors.ts index 38c03621f5..a3e5896cfd 100644 --- a/packages/ui/src/features/sessions/contextColors.ts +++ b/packages/ui/src/features/sessions/contextColors.ts @@ -28,3 +28,14 @@ export function formatTokensCompact(tokens: number): string { if (tokens >= 1000) return `${Math.round(tokens / 1000)}K`; return tokens.toString(); } + +/** + * Formats a USD cost estimate for display. Sub-cent amounts collapse to + * `<$0.01` so a non-zero spend never reads as free; everything else shows two + * decimals ($0.42, $12.34). + */ +export function formatCostUsd(amount: number): string { + if (amount <= 0) return "$0.00"; + if (amount < 0.01) return "<$0.01"; + return `$${amount.toFixed(2)}`; +}