fix(web): open billing for the workspace that needs payment

apps/web/src/components/auth/WorkspaceRouteGuard.tsx

@@ -21,8 +21,11 @@ export function WorkspaceRouteGuard() {
   );
 
   React.useEffect(() => {
+    // While a switch is in flight the route still names the previous
+    // workspace. Following it here would undo the switch midway.
     if (
       isLoading ||
+      isSwitching ||
       !workspaceId ||
       !hasMembership ||
       activeWorkspaceId === workspaceId
@@ -35,6 +38,7 @@ export function WorkspaceRouteGuard() {
     activeWorkspaceId,
     hasMembership,
     isLoading,
+    isSwitching,
     switchWorkspace,
     workspaceId,
   ]);

apps/web/src/lib/api/client.test.ts

@@ -147,3 +147,48 @@ describe("fetchWithAuth refresh handling", () => {
     expect(getAccessToken()).toBe("expired");
   });
 });
+
+describe("payment-required redirect", () => {
+  test("bills the workspace the request was sent for", async () => {
+    // Switching workspaces can change the active workspace while a request is
+    // in flight. Its 402 belongs to the workspace in its X-Company-ID header.
+    const replaced: string[] = [];
+    const originalWindow = (globalThis as { window?: unknown }).window;
+    const originalBillingUrl = process.env.VITE_BILLING_URL;
+    (globalThis as { window?: unknown }).window = {
+      location: {
+        origin: "https://app.example.com",
+        replace: (url: string) => replaced.push(url),
+      },
+    };
+    process.env.VITE_BILLING_URL = "/billing";
+
+    let respond: (response: Response) => void = () => undefined;
+    let sentCompanyId = null as string | null;
+    globalThis.fetch = ((_input: RequestInfo | URL, init?: RequestInit) => {
+      sentCompanyId = (init?.headers as Record<string, string>)["X-Company-ID"];
+      return new Promise<Response>((resolve) => {
+        respond = resolve;
+      });
+    }) as unknown as typeof fetch;
+
+    try {
+      setCompanyId("new-workspace");
+      const request = fetchWithAuth("/companies");
+      setCompanyId("previous-workspace");
+      respond(
+        Response.json({ error: "SUBSCRIPTION_RESTRICTED" }, { status: 402 }),
+      );
+
+      await expect(request).rejects.toThrow();
+      expect(sentCompanyId).toBe("new-workspace");
+      expect(replaced).toEqual([
+        "/billing?companyId=new-workspace&mode=onboarding",
+      ]);
+    } finally {
+      (globalThis as { window?: unknown }).window = originalWindow;
+      if (originalBillingUrl === undefined) delete process.env.VITE_BILLING_URL;
+      else process.env.VITE_BILLING_URL = originalBillingUrl;
+    }
+  });
+});

apps/web/src/lib/api/client.ts

@@ -18,17 +18,20 @@ let paymentRedirectStarted = false;
 
 const COMPANY_ID_STORAGE_KEY = "company_id";
 
-function redirectToBillingOnPaymentRequired(status: number): void {
+function redirectToBillingOnPaymentRequired(
+  status: number,
+  workspaceId: string | null = companyId,
+): void {
   if (
     status !== 402 ||
     paymentRedirectStarted ||
-    !companyId ||
+    !workspaceId ||
     typeof window === "undefined"
   ) {
     return;
   }
 
-  const billingUrl = getWorkspaceBillingUrl(companyId, { onboarding: true });
+  const billingUrl = getWorkspaceBillingUrl(workspaceId, { onboarding: true });
   if (!billingUrl) return;
 
   paymentRedirectStarted = true;
@@ -141,12 +144,17 @@ function sessionRefreshUnavailableError(): ApiRequestError {
 }
 
 // Response handler
-export async function handleResponse<T>(response: Response): Promise<T> {
+export async function handleResponse<T>(
+  response: Response,
+  requestCompanyId: string | null = companyId,
+): Promise<T> {
   if (!response.ok) {
     // A private deployment may require billing before workspace APIs become
     // available. Keep the OSS client commercial-logic-free: HTTP 402 plus the
     // generic configured billing URL is the complete redirect contract.
-    redirectToBillingOnPaymentRequired(response.status);
+    // Bill the workspace the request was sent for: the active workspace can
+    // change while it is in flight, e.g. when switching to a new workspace.
+    redirectToBillingOnPaymentRequired(response.status, requestCompanyId);
     let errorData: {
       code: string;
       message: string;
@@ -277,8 +285,9 @@ export async function fetchWithAuth<T>(
   }
 
   // Add company ID header for multi-tenant support
-  if (companyId) {
-    (headers as Record<string, string>)["X-Company-ID"] = companyId;
+  const requestCompanyId = companyId;
+  if (requestCompanyId) {
+    (headers as Record<string, string>)["X-Company-ID"] = requestCompanyId;
   }
 
   const response = await fetch(url, {
@@ -299,14 +308,14 @@ export async function fetchWithAuth<T>(
         headers,
         credentials: "include",
       });
-      return handleResponse<T>(retryResponse);
+      return handleResponse<T>(retryResponse, requestCompanyId);
     }
     // Report the outage rather than the 401 that triggered it, so a caller
     // does not read a deployment blip as a failed sign-in.
     if (outcome === "unavailable") throw sessionRefreshUnavailableError();
   }
 
-  return handleResponse<T>(response);
+  return handleResponse<T>(response, requestCompanyId);
 }
 
 /** Fetch an authenticated binary response, including token-refresh retry. */
@@ -319,7 +328,8 @@ export async function fetchBlobWithAuth(
     ...(options.headers as Record<string, string> | undefined),
   };
   if (accessToken) headers.Authorization = `Bearer ${accessToken}`;
-  if (companyId) headers["X-Company-ID"] = companyId;
+  const requestCompanyId = companyId;
+  if (requestCompanyId) headers["X-Company-ID"] = requestCompanyId;
 
   let response = await fetch(url, {
     ...options,
@@ -339,7 +349,7 @@ export async function fetchBlobWithAuth(
       throw sessionRefreshUnavailableError();
     }
   }
-  if (!response.ok) await handleResponse<never>(response);
+  if (!response.ok) await handleResponse<never>(response, requestCompanyId);
   return response.blob();
 }
 
@@ -358,8 +368,9 @@ export async function fetchFormDataWithAuth<T>(
     headers.Authorization = `Bearer ${accessToken}`;
   }
 
-  if (companyId) {
-    headers["X-Company-ID"] = companyId;
+  const requestCompanyId = companyId;
+  if (requestCompanyId) {
+    headers["X-Company-ID"] = requestCompanyId;
   }
 
   const response = await fetch(url, {
@@ -380,12 +391,12 @@ export async function fetchFormDataWithAuth<T>(
         body: formData,
         credentials: "include",
       });
-      return handleResponse<T>(retryResponse);
+      return handleResponse<T>(retryResponse, requestCompanyId);
     }
     if (outcome === "unavailable") throw sessionRefreshUnavailableError();
   }
 
-  return handleResponse<T>(response);
+  return handleResponse<T>(response, requestCompanyId);
 }
 
 // Build query string from params