fix: Misc billing, error handling and UX improvements

apps/code/src/main/services/auth/service.ts

@@ -113,6 +113,16 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
     return this.getState();
   }
   async getValidAccessToken(): Promise<ValidAccessTokenOutput> {
+    const override = process.env.VITE_POSTHOG_ACCESS_TOKEN_OVERRIDE;
+    if (override) {
+      await this.initialize();
+      const region = this.session?.cloudRegion ?? "us";
+      return {
+        accessToken: override,
+        apiHost: getCloudUrlFromRegion(region),
+      };
+    }
+
     await this.initialize();
 
     const session = await this.ensureValidSession();
@@ -122,6 +132,16 @@ export class AuthService extends TypedEventEmitter<AuthServiceEvents> {
     };
   }
   async refreshAccessToken(): Promise<ValidAccessTokenOutput> {
+    const override = process.env.VITE_POSTHOG_ACCESS_TOKEN_OVERRIDE;
+    if (override) {
+      await this.initialize();
+      const region = this.session?.cloudRegion ?? "us";
+      return {
+        accessToken: override,
+        apiHost: getCloudUrlFromRegion(region),
+      };
+    }
+
     await this.initialize();
 
     const session = await this.ensureValidSession(true);

apps/code/src/renderer/api/fetcher.test.ts

@@ -13,11 +13,19 @@ describe("buildApiFetcher", () => {
     status: 200,
     json: () => Promise.resolve(data),
   });
-  const err = (status: number) => ({
-    ok: false,
-    status,
-    json: () => Promise.resolve({ error: status }),
-  });
+  const err = (status: number) => {
+    const response = {
+      ok: false,
+      status,
+      statusText: `Error ${status}`,
+      json: () => Promise.resolve({ error: status }),
+      clone: () => ({
+        ...response,
+        text: () => Promise.resolve(`Error ${status}`),
+      }),
+    };
+    return response;
+  };
 
   beforeEach(() => {
     vi.resetAllMocks();

apps/code/src/renderer/api/fetcher.ts

@@ -64,15 +64,25 @@ export const buildApiFetcher: (config: {
             await config.refreshAccessToken(),
           );
         } catch {
-          const errorResponse = await response.json();
+          const cloned = response.clone();
+          const errorResponse = await response
+            .json()
+            .catch(() =>
+              cloned.text().then((t) => ({ error: t || `${response.status}` })),
+            );
           throw new Error(
             `Failed request: [${response.status}] ${JSON.stringify(errorResponse)}`,
           );
         }
       }
 
       if (!response.ok) {
-        const errorResponse = await response.json();
+        const cloned = response.clone();
+        const errorResponse = await response
+          .json()
+          .catch(() =>
+            cloned.text().then((t) => ({ error: t || `${response.status}` })),
+          );
         throw new Error(
           `Failed request: [${response.status}] ${JSON.stringify(errorResponse)}`,
         );

apps/code/src/renderer/api/posthogClient.ts

@@ -2460,11 +2460,15 @@ export class PostHogAPIClient {
     return data.results ?? data ?? [];
   }
 
-  async getMySeat(): Promise<SeatData | null> {
+  async getMySeat(
+    options: { best?: boolean } = { best: true },
+  ): Promise<SeatData | null> {
     try {
       const url = new URL(`${this.api.baseUrl}/api/seats/me/`);
       url.searchParams.set("product_key", SEAT_PRODUCT_KEY);
-      url.searchParams.set("best", "true");
+      if (options.best) {
+        url.searchParams.set("best", "true");
+      }
       const response = await this.api.fetcher.fetch({
         method: "get",
         url,

apps/code/src/renderer/components/FullScreenLayout.tsx

@@ -1,3 +1,4 @@
+import { UpdateBanner } from "@features/sidebar/components/UpdateBanner";
 import { Lifebuoy } from "@phosphor-icons/react";
 import { Button, Flex, Theme } from "@radix-ui/themes";
 import phWordmark from "@renderer/assets/images/wordmark.svg";
@@ -63,20 +64,23 @@ export function FullScreenLayout({
             className="absolute right-[32px] bottom-[20px] left-[32px] z-[2]"
           >
             {footerLeft ?? (
-              <Button
-                size="1"
-                variant="ghost"
-                color="gray"
-                onClick={() =>
-                  trpcClient.os.openExternal.mutate({
-                    url: EXTERNAL_LINKS.discord,
-                  })
-                }
-                className="opacity-50"
-              >
-                <Lifebuoy size={14} />
-                Get support
-              </Button>
+              <Flex align="center" gap="3">
+                <Button
+                  size="1"
+                  variant="ghost"
+                  color="gray"
+                  onClick={() =>
+                    trpcClient.os.openExternal.mutate({
+                      url: EXTERNAL_LINKS.discord,
+                    })
+                  }
+                  className="opacity-50"
+                >
+                  <Lifebuoy size={14} />
+                  Get support
+                </Button>
+                <UpdateBanner variant="compact" />
+              </Flex>
             )}
             {footerRight ?? <div />}
           </Flex>

apps/code/src/renderer/features/auth/hooks/useAuthSession.ts

@@ -14,6 +14,7 @@ import { trpcClient } from "@renderer/trpc/client";
 import { BILLING_FLAG } from "@shared/constants";
 import { identifyUser, resetUser } from "@utils/analytics";
 import { logger } from "@utils/logger";
+import { queryClient } from "@utils/queryClient";
 import { useEffect } from "react";
 
 const log = logger.scope("auth-session");
@@ -93,9 +94,8 @@ function useSeatSync(
       return;
     }
 
-    void useSeatStore.getState().fetchSeat({
-      autoProvision: true,
-    });
+    void useSeatStore.getState().fetchSeat({ autoProvision: true });
+    void queryClient.invalidateQueries({ queryKey: [["llmGateway"]] });
   }, [authIdentity, billingEnabled]);
 }
 

apps/code/src/renderer/features/billing/stores/seatStore.test.ts

@@ -85,6 +85,7 @@ describe("seatStore", () => {
     vi.clearAllMocks();
     useSeatStore.setState({
       seat: null,
+      orgSeat: null,
       isLoading: false,
       error: null,
       redirectUrl: null,

apps/code/src/renderer/features/billing/stores/seatStore.ts

@@ -14,6 +14,7 @@ const log = logger.scope("seat-store");
 
 interface SeatStoreState {
   seat: SeatData | null;
+  orgSeat: SeatData | null;
   isLoading: boolean;
   error: string | null;
   redirectUrl: string | null;
@@ -40,6 +41,25 @@ async function getClient() {
   return client;
 }
 
+async function fetchAndProvision(
+  client: Awaited<ReturnType<typeof getClient>>,
+  options: { best: boolean; autoProvision: boolean },
+): Promise<SeatData | null> {
+  let seat = await client.getMySeat({ best: options.best });
+  if (!seat && options.autoProvision) {
+    log.info("No seat found, auto-provisioning free plan", {
+      best: options.best,
+    });
+    try {
+      seat = await client.createSeat(PLAN_FREE);
+    } catch {
+      log.info("Auto-provision failed, re-fetching seat");
+      seat = await client.getMySeat({ best: options.best });
+    }
+  }
+  return seat;
+}
+
 function handleSeatError(
   error: unknown,
   set: (state: Partial<SeatStoreState>) => void,
@@ -77,6 +97,7 @@ function invalidatePlanCache(): void {
 
 const initialState: SeatStoreState = {
   seat: null,
+  orgSeat: null,
   isLoading: false,
   error: null,
   redirectUrl: null,
@@ -90,18 +111,14 @@ export const useSeatStore = create<SeatStore>()((set, get) => ({
     set({ isLoading: true, error: null, redirectUrl: null });
     try {
       const client = await getClient();
-      let seat = await client.getMySeat();
-      if (!seat && options?.autoProvision) {
-        log.info("No seat found, auto-provisioning free plan");
-        try {
-          seat = await client.createSeat(PLAN_FREE);
-        } catch {
-          log.info("Auto-provision failed, re-fetching seat");
-          seat = await client.getMySeat();
-        }
-      }
+      const autoProvision = options?.autoProvision ?? false;
+      const [seat, orgSeat] = await Promise.all([
+        fetchAndProvision(client, { best: true, autoProvision }),
+        fetchAndProvision(client, { best: false, autoProvision }),
+      ]);
       set({
         seat,
+        orgSeat,
         isLoading: false,
         billingOrgId: seat?.organization_id ?? null,
       });
@@ -165,6 +182,7 @@ export const useSeatStore = create<SeatStore>()((set, get) => ({
         const seat = await client.upgradeSeat(PLAN_PRO);
         set({
           seat,
+          orgSeat: seat,
           isLoading: false,
           billingOrgId: seat.organization_id ?? null,
         });
@@ -174,6 +192,7 @@ export const useSeatStore = create<SeatStore>()((set, get) => ({
       const seat = await client.createSeat(PLAN_PRO);
       set({
         seat,
+        orgSeat: seat,
         isLoading: false,
         billingOrgId: seat.organization_id ?? null,
       });
@@ -191,6 +210,7 @@ export const useSeatStore = create<SeatStore>()((set, get) => ({
       const seat = await client.getMySeat();
       set({
         seat,
+        orgSeat: seat,
         isLoading: false,
         billingOrgId: seat?.organization_id ?? null,
       });
@@ -207,6 +227,7 @@ export const useSeatStore = create<SeatStore>()((set, get) => ({
       const seat = await client.reactivateSeat();
       set({
         seat,
+        orgSeat: seat,
         isLoading: false,
         billingOrgId: seat.organization_id ?? null,
       });

apps/code/src/renderer/features/onboarding/components/ProjectSelectStep.tsx

@@ -7,6 +7,7 @@ import {
   useAuthStateValue,
   useCurrentUser,
 } from "@features/auth/hooks/authQueries";
+import { useSeatStore } from "@features/billing/stores/seatStore";
 import {
   type ProjectInfo,
   useProjects,
@@ -107,6 +108,7 @@ export function ProjectSelectStep({ onNext, onBack }: ProjectSelectStepProps) {
       await queryClient.invalidateQueries({
         queryKey: authKeys.currentUsers(),
       });
+      void useSeatStore.getState().fetchSeat({ autoProvision: true });
     },
     onMutate: () => {
       setIsSwitchingOrg(true);