fix(ui): refresh permissions after tenant switch (#12087)

ui/auth.config.test.ts

@@ -0,0 +1,197 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import { authConfig } from "./auth.config";
+import type { RolePermissionAttributes } from "./types/users";
+
+const { getUserByMeMock } = vi.hoisted(() => ({
+  getUserByMeMock: vi.fn(),
+}));
+
+vi.mock("next-auth", () => ({
+  default: vi.fn(() => ({
+    signIn: vi.fn(),
+    signOut: vi.fn(),
+    auth: vi.fn(),
+    handlers: {},
+  })),
+}));
+
+vi.mock("next-auth/providers/credentials", () => ({
+  default: vi.fn((config) => config),
+}));
+
+vi.mock("./actions/auth", () => ({
+  getToken: vi.fn(),
+  getUserByMe: getUserByMeMock,
+}));
+
+vi.mock("./lib", () => ({
+  apiBaseUrl: "https://api.example.com/api/v1",
+}));
+
+const RESTRICTED_PERMISSIONS: RolePermissionAttributes = {
+  manage_users: false,
+  manage_account: false,
+  manage_providers: false,
+  manage_scans: false,
+  manage_integrations: false,
+  manage_alerts: false,
+  unlimited_visibility: false,
+};
+
+const ELEVATED_PERMISSIONS: RolePermissionAttributes = {
+  ...RESTRICTED_PERMISSIONS,
+  manage_users: true,
+  manage_scans: true,
+};
+
+describe("authConfig JWT callback", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("should load elevated tenant permissions after switching from a restricted tenant", async () => {
+    // Given
+    const accessToken =
+      "header.eyJzdWIiOiJ1c2VyLTEiLCJ0ZW5hbnRfaWQiOiJ0ZW5hbnQtMiJ9.signature";
+    getUserByMeMock.mockResolvedValue({
+      name: "Tenant User",
+      email: "tenant@example.com",
+      company: "Tenant Company",
+      dateJoined: "2026-01-01",
+      permissions: ELEVATED_PERMISSIONS,
+    });
+    const jwtCallback = authConfig.callbacks?.jwt;
+    if (!jwtCallback) throw new Error("JWT callback is not configured");
+
+    // When
+    const result = await jwtCallback({
+      token: {
+        accessToken: "restricted-access-token",
+        refreshToken: "restricted-refresh-token",
+        tenant_id: "tenant-1",
+        user: {
+          name: "Tenant User",
+          email: "tenant@example.com",
+          dateJoined: "2026-01-01",
+          permissions: RESTRICTED_PERMISSIONS,
+        },
+      },
+      user: {} as Parameters<typeof jwtCallback>[0]["user"],
+      trigger: "update",
+      session: {
+        accessToken,
+        refreshToken: "elevated-refresh-token",
+      },
+    });
+
+    // Then
+    expect(getUserByMeMock).toHaveBeenCalledWith(accessToken);
+    expect(result.user).toEqual({
+      name: "Tenant User",
+      email: "tenant@example.com",
+      companyName: "Tenant Company",
+      dateJoined: "2026-01-01",
+      permissions: ELEVATED_PERMISSIONS,
+    });
+  });
+
+  it("should load restricted tenant permissions after switching from an elevated tenant", async () => {
+    // Given
+    const accessToken =
+      "header.eyJzdWIiOiJ1c2VyLTEiLCJ0ZW5hbnRfaWQiOiJ0ZW5hbnQtMSJ9.signature";
+    getUserByMeMock.mockResolvedValue({
+      name: "Tenant User",
+      email: "tenant@example.com",
+      company: "Tenant Company",
+      dateJoined: "2026-01-01",
+      permissions: RESTRICTED_PERMISSIONS,
+    });
+    const jwtCallback = authConfig.callbacks?.jwt;
+    if (!jwtCallback) throw new Error("JWT callback is not configured");
+
+    // When
+    const result = await jwtCallback({
+      token: {
+        accessToken: "elevated-access-token",
+        refreshToken: "elevated-refresh-token",
+        tenant_id: "tenant-2",
+        user: {
+          name: "Tenant User",
+          email: "tenant@example.com",
+          dateJoined: "2026-01-01",
+          permissions: ELEVATED_PERMISSIONS,
+        },
+      },
+      user: {} as Parameters<typeof jwtCallback>[0]["user"],
+      trigger: "update",
+      session: {
+        accessToken,
+        refreshToken: "restricted-refresh-token",
+      },
+    });
+
+    // Then
+    expect(getUserByMeMock).toHaveBeenCalledWith(accessToken);
+    expect(result.accessToken).toBe(accessToken);
+    expect(result.refreshToken).toBe("restricted-refresh-token");
+    expect(result.tenant_id).toBe("tenant-1");
+    expect(result.user).toMatchObject({
+      permissions: RESTRICTED_PERMISSIONS,
+    });
+  });
+
+  it("should report a tenant switch failure while preserving the current session", async () => {
+    // Given
+    vi.spyOn(console, "warn").mockImplementation(() => undefined);
+    getUserByMeMock.mockRejectedValue(new Error("Temporary API failure"));
+    const jwtCallback = authConfig.callbacks?.jwt;
+    if (!jwtCallback) throw new Error("JWT callback is not configured");
+    const sessionCallback = authConfig.callbacks?.session;
+    if (!sessionCallback) throw new Error("Session callback is not configured");
+    const currentToken = {
+      accessToken: "current-access-token",
+      refreshToken: "current-refresh-token",
+      tenant_id: "tenant-1",
+      user: {
+        name: "Tenant User",
+        email: "tenant@example.com",
+        dateJoined: "2026-01-01",
+        permissions: RESTRICTED_PERMISSIONS,
+      },
+    };
+
+    // When
+    const result = await jwtCallback({
+      token: currentToken,
+      user: {} as Parameters<typeof jwtCallback>[0]["user"],
+      trigger: "update",
+      session: {
+        accessToken:
+          "header.eyJzdWIiOiJ1c2VyLTEiLCJ0ZW5hbnRfaWQiOiJ0ZW5hbnQtMiJ9.signature",
+        refreshToken: "switched-refresh-token",
+      },
+    });
+    if (!result) throw new Error("JWT callback cleared the current token");
+
+    const session = await sessionCallback({
+      session: {
+        expires: "2026-12-31T23:59:59.999Z",
+        user: { name: "Tenant User" },
+      },
+      token: result,
+    } as Parameters<typeof sessionCallback>[0]);
+
+    // Then
+    expect(session).toMatchObject({
+      error: "TenantSwitchError",
+      accessToken: "current-access-token",
+      refreshToken: "current-refresh-token",
+      tenantId: "tenant-1",
+      user: {
+        permissions: RESTRICTED_PERMISSIONS,
+      },
+    });
+    expect(result.error).toBeUndefined();
+  });
+});

ui/auth.config.ts

@@ -4,7 +4,6 @@ import NextAuth, {
   type DefaultSession,
   type NextAuthConfig,
   type Session,
-  User,
 } from "next-auth";
 import type { JWT } from "next-auth/jwt";
 import Credentials from "next-auth/providers/credentials";
@@ -58,8 +57,29 @@ const DEFAULT_PERMISSIONS: RolePermissionAttributes = {
   unlimited_visibility: false,
 };
 
+const TENANT_SWITCH_ERROR = "TenantSwitchError";
+
 type TokenUserInput = Partial<TokenUser> & { company?: string };
 
+type JwtCallback = NonNullable<NonNullable<NextAuthConfig["callbacks"]>["jwt"]>;
+type JwtCallbackParams = Parameters<JwtCallback>[0];
+
+interface JwtCallbackCredentials {
+  accessToken?: string;
+  refreshToken?: string;
+}
+
+type AuthJwtUser = JwtCallbackParams["user"] &
+  TokenUserInput &
+  JwtCallbackCredentials;
+
+interface AuthJwtCallbackParams
+  extends Omit<JwtCallbackParams, "session" | "token" | "user"> {
+  session?: Partial<ExtendedSession>;
+  token: AuthToken;
+  user: AuthJwtUser;
+}
+
 const toTokenUser = (user?: TokenUserInput): TokenUser =>
   ({
     name: user?.name ?? undefined,
@@ -308,37 +328,54 @@ export const authConfig = {
       return true;
     },
 
-    jwt: async ({ token, account, user, trigger, session }) => {
-      const authToken = token as AuthToken;
-
+    jwt: async ({
+      token: authToken,
+      account,
+      user,
+      trigger,
+      session,
+    }: AuthJwtCallbackParams): Promise<AuthToken> => {
       // Handle tenant switch: update tokens from client-side useSession().update()
       if (trigger === "update" && session?.accessToken) {
-        authToken.accessToken = session.accessToken;
-        authToken.refreshToken = session.refreshToken;
-        applyDecodedClaims(authToken, authToken.accessToken, "tenant switch");
-        return authToken;
+        const newAccessToken = session.accessToken;
+
+        try {
+          const userMeResponse = await getUserByMe(newAccessToken);
+          const nextAuthToken: AuthToken = {
+            ...authToken,
+            accessToken: newAccessToken,
+            refreshToken: session.refreshToken,
+            user: tokenUserFromApi(userMeResponse),
+            error: undefined,
+          };
+
+          applyDecodedClaims(nextAuthToken, newAccessToken, "tenant switch");
+
+          return nextAuthToken;
+        } catch (error) {
+          // eslint-disable-next-line no-console
+          console.warn("Error refreshing user after tenant switch:", error);
+          return {
+            ...authToken,
+            error: TENANT_SWITCH_ERROR,
+          };
+        }
       }
 
       applyDecodedClaims(authToken, authToken.accessToken);
 
-      if (account && user) {
-        const signedInUser = user as User &
-          TokenUserInput & {
-            accessToken: string;
-            refreshToken: string;
-          };
-
+      if (account && user?.accessToken && user.refreshToken) {
         const nextAuthToken: AuthToken = {
           ...authToken,
-          accessToken: signedInUser.accessToken,
-          refreshToken: signedInUser.refreshToken,
-          user: toTokenUser(signedInUser),
+          accessToken: user.accessToken,
+          refreshToken: user.refreshToken,
+          user: toTokenUser(user),
           error: undefined,
         };
 
         applyDecodedClaims(
           nextAuthToken,
-          signedInUser.accessToken,
+          user.accessToken,
           "access token on sign-in",
         );
 
@@ -359,7 +396,7 @@ export const authConfig = {
       const authToken = token as AuthToken;
       const nextSession = { ...session } as ExtendedSession;
 
-      if (authToken?.error) {
+      if (authToken.error && authToken.error !== TENANT_SWITCH_ERROR) {
         nextSession.error = authToken.error;
         nextSession.user = undefined;
         nextSession.userId = undefined;
@@ -369,7 +406,8 @@ export const authConfig = {
         return nextSession;
       }
 
-      nextSession.error = undefined;
+      nextSession.error = authToken.error;
+      authToken.error = undefined;
       nextSession.userId = authToken.user_id ?? nextSession.userId;
       nextSession.tenantId = authToken.tenant_id ?? nextSession.tenantId;
       nextSession.accessToken =

ui/changelog.d/tenant-switch-permissions.fixed.md

@@ -0,0 +1 @@
+Tenant switches now refresh session user permissions for the selected tenant

ui/components/users/forms/switch-tenant-form.test.tsx

@@ -1,27 +1,42 @@
-import { render, screen } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 
 import { SwitchTenantForm } from "./switch-tenant-form";
 
-const mockUpdate = vi.fn();
+const { mockReloadPage, mockSwitchTenant, mockToast, mockUpdate } = vi.hoisted(
+  () => ({
+    mockReloadPage: vi.fn(),
+    mockSwitchTenant: vi.fn(),
+    mockToast: vi.fn(),
+    mockUpdate: vi.fn(),
+  }),
+);
+
 vi.mock("next-auth/react", () => ({
   useSession: () => ({ update: mockUpdate }),
 }));
 
 vi.mock("@/actions/users/tenants", () => ({
-  switchTenant: vi.fn(),
+  switchTenant: mockSwitchTenant,
 }));
 
-const mockToast = vi.fn();
 vi.mock("@/components/shadcn", async (importOriginal) => ({
   ...(await importOriginal<Record<string, unknown>>()),
   useToast: () => ({ toast: mockToast }),
 }));
 
+vi.mock("@/lib/navigation", () => ({
+  reloadPage: mockReloadPage,
+}));
+
 describe("SwitchTenantForm", () => {
   const setIsOpen = vi.fn();
 
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
   it("renders confirm and cancel buttons", () => {
     render(<SwitchTenantForm tenantId="test-uuid" setIsOpen={setIsOpen} />);
 
@@ -48,4 +63,55 @@ describe("SwitchTenantForm", () => {
     await user.click(screen.getByRole("button", { name: /cancel/i }));
     expect(setIsOpen).toHaveBeenCalledWith(false);
   });
+
+  it("shows an error when the session cannot apply the tenant switch", async () => {
+    // Given
+    const user = userEvent.setup();
+    mockSwitchTenant.mockResolvedValue({
+      success: true,
+      accessToken: "switched-access-token",
+      refreshToken: "switched-refresh-token",
+    });
+    mockUpdate.mockResolvedValue({ error: "TenantSwitchError" });
+    render(<SwitchTenantForm tenantId="test-uuid" setIsOpen={setIsOpen} />);
+
+    // When
+    await user.click(screen.getByRole("button", { name: /confirm/i }));
+
+    // Then
+    await waitFor(() =>
+      expect(mockToast).toHaveBeenCalledWith({
+        variant: "destructive",
+        title: "Oops! Something went wrong",
+        description: "Unable to switch organization. Please try again.",
+      }),
+    );
+    expect(mockReloadPage).not.toHaveBeenCalled();
+  });
+
+  it("reloads after the session applies the tenant switch", async () => {
+    // Given
+    const user = userEvent.setup();
+    mockSwitchTenant.mockResolvedValue({
+      success: true,
+      accessToken: "switched-access-token",
+      refreshToken: "switched-refresh-token",
+    });
+    mockUpdate.mockResolvedValue({
+      expires: "2026-12-31T23:59:59.999Z",
+    });
+    render(<SwitchTenantForm tenantId="test-uuid" setIsOpen={setIsOpen} />);
+
+    // When
+    await user.click(screen.getByRole("button", { name: /confirm/i }));
+
+    // Then
+    await waitFor(() =>
+      expect(mockToast).toHaveBeenCalledWith({
+        title: "Organization switched",
+        description: "The page will reload to apply the change.",
+      }),
+    );
+    expect(mockReloadPage).toHaveBeenCalledOnce();
+  });
 });

ui/components/users/forms/switch-tenant-form.tsx

@@ -23,15 +23,32 @@ export const SwitchTenantForm = ({
 
     const handleSwitch = async () => {
       if ("success" in state) {
-        await update({
-          accessToken: state.accessToken,
-          refreshToken: state.refreshToken,
-        });
-        toast({
-          title: "Organization switched",
-          description: "The page will reload to apply the change.",
-        });
-        reloadPage();
+        try {
+          const updatedSession = await update({
+            accessToken: state.accessToken,
+            refreshToken: state.refreshToken,
+          });
+
+          if (
+            !updatedSession ||
+            ("error" in updatedSession && updatedSession.error)
+          ) {
+            throw new Error("Session update failed");
+          }
+
+          toast({
+            title: "Organization switched",
+            description: "The page will reload to apply the change.",
+          });
+          reloadPage();
+        } catch {
+          toast({
+            variant: "destructive",
+            title: "Oops! Something went wrong",
+            description: "Unable to switch organization. Please try again.",
+          });
+          setIsOpen(false);
+        }
       } else {
         toast({
           variant: "destructive",