fix(billing): scope the usage snapshot to the identity; emit on dollar changes; fail hydration open

Review fixes: isSameUsage now compares ai_credits.used_usd/limit_usd so a
subscribed org's dollar meter doesn't freeze all period (nothing else in
its snapshot ever changes); the usage monitor clears its snapshot on org
transitions and the renderer drops the cached query on identity change, so
a new sign-in never sees the previous account's spend; a failed
announcement-store hydration now fails open as unacknowledged — re-showing
the one-time modal beats never showing it, and in-memory state still
dismisses it for the session.

Generated-By: PostHog Code
Task-Id: 1039ea23-9930-44e2-9888-b05cf8b129ec

Author: adboio

packages/core/src/usage/usage-monitor.test.ts

@@ -1,4 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import type { AuthService } from "../auth/auth";
 import type { UsageHost } from "./identifiers";
 import { UsageMonitorEvent } from "./monitor-schemas";
 import type { UsageOutput } from "./schemas";
@@ -50,6 +51,26 @@ function makeLogger() {
 
 type GatewaySlice = Pick<UsageHost, "fetchUsage">;
 
+let emitAuthState: (currentOrgId: string | null) => void = () => {};
+
+function makeAuthService(): AuthService {
+  const listeners = new Set<
+    (state: { currentOrgId: string | null }) => void
+  >();
+  emitAuthState = (currentOrgId) => {
+    for (const listener of [...listeners]) {
+      listener({ currentOrgId });
+    }
+  };
+  return {
+    getState: () => ({ currentOrgId: "org-1" }),
+    on: (
+      _event: string,
+      listener: (state: { currentOrgId: string | null }) => void,
+    ) => listeners.add(listener),
+  } as unknown as AuthService;
+}
+
 function makeService(
   gateway: GatewaySlice,
   activity: ActivitySlice,
@@ -59,7 +80,7 @@ function makeService(
     ...activity,
     ...makeThresholdStore(),
   };
-  return new UsageMonitorService(host, makeLogger());
+  return new UsageMonitorService(host, makeLogger(), makeAuthService());
 }
 
 function makeUsage(overrides?: {
@@ -273,6 +294,65 @@ describe("UsageMonitorService", () => {
 
   // The titlebar meter keys off this bit; a subscribe flip must not wait for
   // some other field to change.
+  it("emits UsageUpdated when only the org spend numbers change", async () => {
+    const updates: UsageOutput[] = [];
+    const gateway = {
+      fetchUsage: vi
+        .fn()
+        .mockResolvedValueOnce({
+          ...makeUsage(),
+          ai_credits: { exhausted: false, used_usd: 12.4, limit_usd: 50 },
+        })
+        .mockResolvedValueOnce({
+          ...makeUsage(),
+          ai_credits: { exhausted: false, used_usd: 13.1, limit_usd: 50 },
+        }),
+    } as unknown as GatewaySlice;
+    service = makeService(gateway, makeActivityMonitor());
+    service.on(UsageMonitorEvent.UsageUpdated, (u) => updates.push(u));
+
+    await service.fetchOnce();
+    await service.fetchOnce();
+
+    expect(updates).toHaveLength(2);
+    expect(updates[1].ai_credits?.used_usd).toBe(13.1);
+  });
+
+  it("forgets the snapshot and refetches when the organization changes", async () => {
+    const gateway = mockGateway(makeUsage());
+    service = makeService(gateway, makeActivityMonitor());
+    await service.fetchOnce();
+    expect(service.getLatest()).not.toBeNull();
+
+    emitAuthState("org-2");
+
+    expect(service.getLatest()).toBeNull();
+    await vi.advanceTimersByTimeAsync(10_000);
+    expect(gateway.fetchUsage).toHaveBeenCalledTimes(2);
+    expect(service.getLatest()).not.toBeNull();
+  });
+
+  it("keeps the snapshot across same-org auth changes", async () => {
+    service = makeService(mockGateway(makeUsage()), makeActivityMonitor());
+    await service.fetchOnce();
+
+    emitAuthState("org-1");
+
+    expect(service.getLatest()).not.toBeNull();
+  });
+
+  it("clears the snapshot on sign-out without scheduling a refetch", async () => {
+    const gateway = mockGateway(makeUsage());
+    service = makeService(gateway, makeActivityMonitor());
+    await service.fetchOnce();
+
+    emitAuthState(null);
+
+    expect(service.getLatest()).toBeNull();
+    await vi.advanceTimersByTimeAsync(10_000);
+    expect(gateway.fetchUsage).toHaveBeenCalledTimes(1);
+  });
+
   it("emits UsageUpdated when only the subscription bit flips", async () => {
     const updates: UsageOutput[] = [];
     const gateway = {

packages/core/src/usage/usage-monitor.ts

@@ -1,6 +1,9 @@
 import { ROOT_LOGGER, type RootLogger } from "@posthog/di/logger";
 import { TypedEventEmitter } from "@posthog/shared";
 import { inject, injectable, postConstruct, preDestroy } from "inversify";
+import type { AuthService } from "../auth/auth";
+import { AUTH_SERVICE } from "../auth/auth.module";
+import { AuthServiceEvent } from "../auth/schemas";
 import { isCodeUsageFreeTier } from "../billing/usageDisplay";
 import { USAGE_HOST, type UsageHost, type UsageLogger } from "./identifiers";
 import {
@@ -34,10 +37,21 @@ export class UsageMonitorService extends TypedEventEmitter<UsageMonitorEvents> {
     private readonly host: UsageHost,
     @inject(ROOT_LOGGER)
     logger: RootLogger,
+    @inject(AUTH_SERVICE)
+    authService: AuthService,
   ) {
     super();
     this.log = logger.scope("usage-monitor");
     this.thresholdsSeen = { ...this.host.getThresholdsSeen() };
+    // The snapshot is identity-scoped billing data: a signed-in account must
+    // never be served the previous account's spend.
+    let orgId = authService.getState().currentOrgId;
+    authService.on(AuthServiceEvent.StateChanged, (state) => {
+      if (state.currentOrgId === orgId) return;
+      orgId = state.currentOrgId;
+      this.latestUsage = null;
+      if (state.currentOrgId !== null) this.requestRefresh();
+    });
   }
 
   private readonly log: UsageLogger;
@@ -251,6 +265,8 @@ function isSameUsage(a: UsageOutput | null, b: UsageOutput): boolean {
     a.billing_period_end === b.billing_period_end &&
     a.code_usage_subscribed === b.code_usage_subscribed &&
     a.ai_credits?.exhausted === b.ai_credits?.exhausted &&
+    a.ai_credits?.used_usd === b.ai_credits?.used_usd &&
+    a.ai_credits?.limit_usd === b.ai_credits?.limit_usd &&
     isSameBucket(a.burst, b.burst) &&
     isSameBucket(a.sustained, b.sustained)
   );

packages/ui/src/features/auth/useAuthSession.ts

@@ -1,6 +1,7 @@
 import { useHostTRPCClient } from "@posthog/host-router/react";
 import { BILLING_FLAG } from "@posthog/shared";
 import { useSeatStore } from "@posthog/ui/features/billing/seatStore";
+import { USAGE_QUERY_KEY } from "@posthog/ui/features/billing/useUsage";
 import { useFeatureFlag } from "@posthog/ui/features/feature-flags/useFeatureFlag";
 import {
   identifyUser,
@@ -113,6 +114,9 @@ function useSeatSync(
 ): void {
   const queryClient = useQueryClient();
   useEffect(() => {
+    // Usage is identity-scoped billing data — drop the cached snapshot so a
+    // new sign-in never renders the previous account's spend.
+    queryClient.removeQueries({ queryKey: USAGE_QUERY_KEY });
     if (!authIdentity || !billingEnabled) {
       useSeatStore.getState().reset();
       return;

packages/ui/src/features/billing/billingAnnouncementStore.ts

@@ -24,7 +24,14 @@ export const useBillingAnnouncementStore = create<BillingAnnouncementState>()(
       storage: electronStorage,
       partialize: (state) => ({ acknowledged: state.acknowledged }),
       onRehydrateStorage: () => (state) => {
-        state?.setHasHydrated(true);
+        if (state) {
+          state.setHasHydrated(true);
+          return;
+        }
+        // Failed storage read: fail open as unacknowledged — re-showing the
+        // one-time announcement beats never showing it, and the in-memory
+        // acknowledgment still dismisses it for the rest of the session.
+        useBillingAnnouncementStore.setState({ _hasHydrated: true });
       },
     },
   ),

packages/ui/src/features/billing/useUsage.ts

@@ -2,7 +2,7 @@ import { useHostTRPCClient } from "@posthog/host-router/react";
 import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";
 import { useCallback, useEffect } from "react";
 
-const USAGE_QUERY_KEY = ["billing", "usage", "latest"] as const;
+export const USAGE_QUERY_KEY = ["billing", "usage", "latest"] as const;
 
 export function useUsage({ enabled = true }: { enabled?: boolean } = {}) {
   const client = useHostTRPCClient();