fix(billing): re-clear usage cache on org switch

The UI usage-cache clear keyed on authIdentity (region + project), but
usage is org-scoped — the host UsageMonitorService keys its snapshot on
currentOrgId. Switching between two orgs that both have no selected
project leaves authIdentity at "region:none", so the effect never
re-ran and a failed post-switch refresh could leave the previous org's
spend cached. Depend on currentOrgId too, matching the host.

Addresses the Greptile P1 review comment on #3507.

Generated-By: PostHog Code
Task-Id: 90f5b50d-504e-433d-82d9-2b9ead0bb1e5

Author: adboio

packages/ui/src/features/auth/useAuthSession.test.tsx

@@ -0,0 +1,77 @@
+import { USAGE_QUERY_KEY } from "@posthog/ui/features/billing/useUsage";
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { renderHook } from "@testing-library/react";
+import type { ReactNode } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { useUsageIdentitySync } from "./useAuthSession";
+
+let queryClient: QueryClient;
+function wrapper({ children }: { children: ReactNode }) {
+  return (
+    <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+  );
+}
+
+describe("useUsageIdentitySync", () => {
+  beforeEach(() => {
+    queryClient = new QueryClient({
+      defaultOptions: { queries: { retry: false } },
+    });
+  });
+
+  it("drops the cached usage snapshot on mount", () => {
+    const removeSpy = vi.spyOn(queryClient, "removeQueries");
+
+    renderHook(() => useUsageIdentitySync("us:none", "org-a"), { wrapper });
+
+    expect(removeSpy).toHaveBeenCalledWith({ queryKey: USAGE_QUERY_KEY });
+    expect(removeSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-clears the snapshot when the org changes even if authIdentity is unchanged", () => {
+    // Regression: usage is org-scoped, but authIdentity keys on region + project.
+    // Switching between two orgs that both have no selected project leaves
+    // authIdentity at "region:none", so keying only on it would leave the
+    // previous org's spend cached after a failed refresh.
+    const removeSpy = vi.spyOn(queryClient, "removeQueries");
+
+    const { rerender } = renderHook(
+      ({ authIdentity, orgId }) => useUsageIdentitySync(authIdentity, orgId),
+      { wrapper, initialProps: { authIdentity: "us:none", orgId: "org-a" } },
+    );
+    removeSpy.mockClear();
+
+    rerender({ authIdentity: "us:none", orgId: "org-b" });
+
+    expect(removeSpy).toHaveBeenCalledWith({ queryKey: USAGE_QUERY_KEY });
+    expect(removeSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it("re-clears the snapshot when authIdentity changes", () => {
+    const removeSpy = vi.spyOn(queryClient, "removeQueries");
+
+    const { rerender } = renderHook(
+      ({ authIdentity, orgId }) => useUsageIdentitySync(authIdentity, orgId),
+      { wrapper, initialProps: { authIdentity: "us:1", orgId: "org-a" } },
+    );
+    removeSpy.mockClear();
+
+    rerender({ authIdentity: "us:2", orgId: "org-a" });
+
+    expect(removeSpy).toHaveBeenCalledTimes(1);
+  });
+
+  it("does not re-clear when neither authIdentity nor org changes", () => {
+    const removeSpy = vi.spyOn(queryClient, "removeQueries");
+
+    const { rerender } = renderHook(
+      ({ authIdentity, orgId }) => useUsageIdentitySync(authIdentity, orgId),
+      { wrapper, initialProps: { authIdentity: "us:1", orgId: "org-a" } },
+    );
+    removeSpy.mockClear();
+
+    rerender({ authIdentity: "us:1", orgId: "org-a" });
+
+    expect(removeSpy).not.toHaveBeenCalled();
+  });
+});

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

@@ -105,15 +105,22 @@ function useAuthAnalyticsIdentity(
   ]);
 }
 
-function useUsageIdentitySync(authIdentity: string | null): void {
+export function useUsageIdentitySync(
+  authIdentity: string | null,
+  orgId: string | null,
+): void {
   const queryClient = useQueryClient();
-  // biome-ignore lint/correctness/useExhaustiveDependencies: re-run on identity change
+  // Usage is org-scoped billing data — drop the cached snapshot on any sign-in,
+  // sign-out, region, or org switch so a new org never renders the previous
+  // org's spend. authIdentity keys on region + project, which can stay constant
+  // across an org switch (e.g. between two orgs with no selected project), so
+  // depend on currentOrgId too — matching the host UsageMonitorService, which
+  // keys its snapshot on currentOrgId. Without it, a post-switch refresh that
+  // fails would leave the stale value cached.
+  // biome-ignore lint/correctness/useExhaustiveDependencies: re-run on identity/org change
   useEffect(() => {
-    // Usage is identity-scoped billing data — drop the cached snapshot on any
-    // sign-in, sign-out, or org/project switch so a new identity never renders
-    // the previous account's spend.
     queryClient.removeQueries({ queryKey: USAGE_QUERY_KEY });
-  }, [authIdentity, queryClient]);
+  }, [authIdentity, orgId, queryClient]);
 }
 
 export function useAuthSession() {
@@ -125,7 +132,7 @@ export function useAuthSession() {
   useAuthSubscriptionSync();
   useAuthIdentitySync(authState);
   useAuthAnalyticsIdentity(authIdentity, authState, currentUser);
-  useUsageIdentitySync(authIdentity);
+  useUsageIdentitySync(authIdentity, authState.currentOrgId);
 
   return {
     authState,