Merge branch 'main' into posthog-code/autoresearch-task-switch-latency

Author: MattPua

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);

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<ContextUsage, "breakdown">;
+type ContextUsageAggregate = Omit<ContextUsage, "breakdown" | "cost">;
 
 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<ContextUsageState, ContextUsage | null>({
-    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",
       };
     }
   }

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";

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.
         </Text>
       )}
+
+      {showCost && cost && (
+        <Flex
+          align="center"
+          justify="between"
+          className="border-(--gray-4) border-t pt-2 text-[13px]"
+        >
+          <Text className="text-(--gray-11)">Estimated cost</Text>
+          <Text className="font-medium text-(--gray-12) tabular-nums">
+            {formatCostUsd(cost.amount)}
+          </Text>
+        </Flex>
+      )}
     </Flex>
   );
 }

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>): 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(
+      <Theme>
+        <ContextUsageIndicator
+          usage={usage({ cost: { amount: 0.42, currency: "USD" } })}
+        />
+      </Theme>,
+    );
+    expect(screen.getByText(/50K\/200K · 25% · \$0\.42/)).toBeInTheDocument();
+  });
+
   it("renders a finite stroke offset at 0% (no NaN/Infinity)", () => {
     const { container } = render(
       <Theme>

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 (
     <Popover.Root>
@@ -66,15 +78,13 @@ export function ContextUsageIndicator({ usage }: ContextUsageIndicatorProps) {
               />
             </svg>
             <Text className="text-[13px] text-muted-foreground tabular-nums">
-              {hasSize
-                ? `${formatTokensCompact(used)}/${formatTokensCompact(size)} · ${percentage}%`
-                : formatTokensCompact(used)}
+              {label}
             </Text>
           </Flex>
         </button>
       </Popover.Trigger>
       <Popover.Content size="2" side="top" align="end" sideOffset={6}>
-        <ContextBreakdownPopover usage={usage} />
+        <ContextBreakdownPopover usage={usage} showCost={showCost} />
       </Popover.Content>
     </Popover.Root>
   );

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)}`;
+}