coder · aslilac · May 14, 2024 · May 13, 2024 · May 13, 2024 · May 13, 2024
diff --git a/site/src/contexts/auth/AuthProvider.tsx b/site/src/contexts/auth/AuthProvider.tsx
@@ -30,7 +30,7 @@ export type AuthContextValue = {
   isUpdatingProfile: boolean;
   user: User | undefined;
   permissions: Permissions | undefined;
-  organizationId: string | undefined;
+  organizationIds: readonly string[] | undefined;
   signInError: unknown;
   updateProfileError: unknown;
   signOut: () => void;
@@ -119,7 +119,7 @@ export const AuthProvider: FC<PropsWithChildren> = ({ children }) => {
         permissions: permissionsQuery.data as Permissions | undefined,
         signInError: loginMutation.error,
         updateProfileError: updateProfileMutation.error,
-        organizationId: userQuery.data?.organization_ids[0],
+        organizationIds: userQuery.data?.organization_ids,
       }}
     >
       {children}

diff --git a/site/src/contexts/auth/RequireAuth.test.tsx b/site/src/contexts/auth/RequireAuth.test.tsx
@@ -45,7 +45,7 @@ const createAuthWrapper = (override: Partial<AuthContextValue>) => {
     isUpdatingProfile: false,
     permissions: undefined,
     authMethods: undefined,
-    organizationId: undefined,
+    organizationIds: undefined,
     signInError: undefined,
     updateProfileError: undefined,
     signOut: jest.fn(),
@@ -95,6 +95,7 @@ describe("useAuthenticated", () => {
         wrapper: createAuthWrapper({
           user: MockUser,
           permissions: MockPermissions,
+          organizationIds: [],
         }),
       });
     }).not.toThrow();

diff --git a/site/src/contexts/auth/RequireAuth.tsx b/site/src/contexts/auth/RequireAuth.tsx
@@ -66,15 +66,18 @@ export const RequireAuth: FC = () => {
   );
 };
 
+type RequireKeys<T, R extends keyof T> = Omit<T, R> & {
+  [K in keyof Pick<T, R>]: NonNullable<T[K]>;
+};
+
 // We can do some TS magic here but I would rather to be explicit on what
 // values are not undefined when authenticated
-type NonNullableAuth = AuthContextValue & {
-  user: Exclude<AuthContextValue["user"], undefined>;
-  permissions: Exclude<AuthContextValue["permissions"], undefined>;
-  organizationId: Exclude<AuthContextValue["organizationId"], undefined>;
-};
+type AuthenticatedAuthContextValue = RequireKeys<
+  AuthContextValue,
+  "user" | "permissions" | "organizationIds"
+>;
 
-export const useAuthenticated = (): NonNullableAuth => {
+export const useAuthenticated = (): AuthenticatedAuthContextValue => {
   const auth = useAuthContext();
 
   if (!auth.user) {
@@ -85,5 +88,9 @@ export const useAuthenticated = (): NonNullableAuth => {
     throw new Error("Permissions are not available.");
   }
 
-  return auth as NonNullableAuth;
+  if (!auth.organizationIds) {
+    throw new Error("Organization ID is not available.");
+  }
+
+  return auth as AuthenticatedAuthContextValue;
 };
diff --git a/site/src/modules/dashboard/DashboardProvider.tsx b/site/src/modules/dashboard/DashboardProvider.tsx
@@ -1,4 +1,9 @@
-import { createContext, type FC, type PropsWithChildren } from "react";
+import {
+  createContext,
+  type FC,
+  type PropsWithChildren,
+  useState,
+} from "react";
 import { useQuery } from "react-query";
 import { appearance } from "api/queries/appearance";
 import { entitlements } from "api/queries/entitlements";
@@ -9,9 +14,13 @@ import type {
   Experiments,
 } from "api/typesGenerated";
 import { Loader } from "components/Loader/Loader";
+import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useEffectEvent } from "hooks/hookPolyfills";
 import { useEmbeddedMetadata } from "hooks/useEmbeddedMetadata";
 
 export interface DashboardValue {
+  organizationId: string;
+  setOrganizationId: (id: string) => void;
   entitlements: Entitlements;
   experiments: Experiments;
   appearance: AppearanceConfig;
@@ -23,20 +32,40 @@ export const DashboardContext = createContext<DashboardValue | undefined>(
 
 export const DashboardProvider: FC<PropsWithChildren> = ({ children }) => {
   const { metadata } = useEmbeddedMetadata();
+  const { user, organizationIds } = useAuthenticated();
   const entitlementsQuery = useQuery(entitlements(metadata.entitlements));
   const experimentsQuery = useQuery(experiments(metadata.experiments));
   const appearanceQuery = useQuery(appearance(metadata.appearance));
 
   const isLoading =
     !entitlementsQuery.data || !appearanceQuery.data || !experimentsQuery.data;
 
+  const lastUsedOrganizationId = localStorage.getItem(
+    `user:${user.id}.lastUsedOrganizationId`,
+  );
+  const [activeOrganizationId, setActiveOrganizationId] = useState(() =>
+    lastUsedOrganizationId && organizationIds.includes(lastUsedOrganizationId)
+      ? lastUsedOrganizationId
+      : organizationIds[0],
+  );
+
+  const setOrganizationId = useEffectEvent((id: string) => {
+    if (!organizationIds.includes(id)) {
+      throw new ReferenceError("Invalid organization ID");
+    }
+    localStorage.setItem(`user:${user.id}.lastUsedOrganizationId`, id);
+    setActiveOrganizationId(id);
+  });
+
   if (isLoading) {
     return <Loader fullscreen />;
   }
 
   return (
     <DashboardContext.Provider
       value={{
+        organizationId: activeOrganizationId,
+        setOrganizationId: setOrganizationId,
         entitlements: entitlementsQuery.data,
         experiments: experimentsQuery.data,
         appearance: appearanceQuery.data,

diff --git a/site/src/modules/workspaces/WorkspaceStatusBadge/WorkspaceStatusBadge.stories.tsx b/site/src/modules/workspaces/WorkspaceStatusBadge/WorkspaceStatusBadge.stories.tsx
@@ -1,21 +1,18 @@
 import type { Meta, StoryObj } from "@storybook/react";
-import { DashboardContext } from "modules/dashboard/DashboardProvider";
 import {
-  MockAppearanceConfig,
   MockBuildInfo,
   MockCanceledWorkspace,
   MockCancelingWorkspace,
   MockDeletedWorkspace,
   MockDeletingWorkspace,
-  MockEntitlementsWithScheduling,
-  MockExperiments,
   MockFailedWorkspace,
   MockPendingWorkspace,
   MockStartingWorkspace,
   MockStoppedWorkspace,
   MockStoppingWorkspace,
   MockWorkspace,
 } from "testHelpers/entities";
+import { withDashboardProvider } from "testHelpers/storybook";
 import { WorkspaceStatusBadge } from "./WorkspaceStatusBadge";
 
 const meta: Meta<typeof WorkspaceStatusBadge> = {
@@ -29,19 +26,7 @@ const meta: Meta<typeof WorkspaceStatusBadge> = {
       },
     ],
   },
-  decorators: [
-    (Story) => (
-      <DashboardContext.Provider
-        value={{
-          entitlements: MockEntitlementsWithScheduling,
-          experiments: MockExperiments,
-          appearance: MockAppearanceConfig,
-        }}
-      >
-        <Story />
-      </DashboardContext.Provider>
-    ),
-  ],
+  decorators: [withDashboardProvider],
 };
 
 export default meta;

diff --git a/site/src/pages/CreateTemplatePage/DuplicateTemplateView.tsx b/site/src/pages/CreateTemplatePage/DuplicateTemplateView.tsx
@@ -10,7 +10,6 @@ import {
 } from "api/queries/templates";
 import { ErrorAlert } from "components/Alert/ErrorAlert";
 import { Loader } from "components/Loader/Loader";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
 import { useDashboard } from "modules/dashboard/useDashboard";
 import { CreateTemplateForm } from "./CreateTemplateForm";
 import type { CreateTemplatePageViewProps } from "./types";
@@ -24,7 +23,7 @@ export const DuplicateTemplateView: FC<CreateTemplatePageViewProps> = ({
   isCreating,
 }) => {
   const navigate = useNavigate();
-  const { organizationId } = useAuthenticated();
+  const { entitlements, organizationId } = useDashboard();
   const [searchParams] = useSearchParams();
   const templateByNameQuery = useQuery(
     templateByName(organizationId, searchParams.get("fromTemplate")!),
@@ -47,8 +46,7 @@ export const DuplicateTemplateView: FC<CreateTemplatePageViewProps> = ({
     templateVersionQuery.error ||
     templateVersionVariablesQuery.error;
 
-  const dashboard = useDashboard();
-  const formPermissions = getFormPermissions(dashboard.entitlements);
+  const formPermissions = getFormPermissions(entitlements);
 
   const isJobError = error instanceof JobError;
   const templateVersionLogsQuery = useQuery({

diff --git a/site/src/pages/CreateTemplatePage/ImportStarterTemplateView.tsx b/site/src/pages/CreateTemplatePage/ImportStarterTemplateView.tsx
@@ -9,7 +9,6 @@ import {
 } from "api/queries/templates";
 import { ErrorAlert } from "components/Alert/ErrorAlert";
 import { Loader } from "components/Loader/Loader";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
 import { useDashboard } from "modules/dashboard/useDashboard";
 import { CreateTemplateForm } from "./CreateTemplateForm";
 import type { CreateTemplatePageViewProps } from "./types";
@@ -27,7 +26,7 @@ export const ImportStarterTemplateView: FC<CreateTemplatePageViewProps> = ({
   isCreating,
 }) => {
   const navigate = useNavigate();
-  const { organizationId } = useAuthenticated();
+  const { entitlements, organizationId } = useDashboard();
   const [searchParams] = useSearchParams();
   const templateExamplesQuery = useQuery(templateExamples(organizationId));
   const templateExample = templateExamplesQuery.data?.find(
@@ -37,8 +36,7 @@ export const ImportStarterTemplateView: FC<CreateTemplatePageViewProps> = ({
   const isLoading = templateExamplesQuery.isLoading;
   const loadingError = templateExamplesQuery.error;
 
-  const dashboard = useDashboard();
-  const formPermissions = getFormPermissions(dashboard.entitlements);
+  const formPermissions = getFormPermissions(entitlements);
 
   const isJobError = error instanceof JobError;
   const templateVersionLogsQuery = useQuery({

diff --git a/site/src/pages/CreateTemplatePage/UploadTemplateView.tsx b/site/src/pages/CreateTemplatePage/UploadTemplateView.tsx
@@ -7,7 +7,6 @@ import {
   JobError,
   templateVersionVariables,
 } from "api/queries/templates";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
 import { useDashboard } from "modules/dashboard/useDashboard";
 import { CreateTemplateForm } from "./CreateTemplateForm";
 import type { CreateTemplatePageViewProps } from "./types";
@@ -21,10 +20,9 @@ export const UploadTemplateView: FC<CreateTemplatePageViewProps> = ({
   error,
 }) => {
   const navigate = useNavigate();
-  const { organizationId } = useAuthenticated();
 
-  const dashboard = useDashboard();
-  const formPermissions = getFormPermissions(dashboard.entitlements);
+  const { entitlements, organizationId } = useDashboard();
+  const formPermissions = getFormPermissions(entitlements);
 
   const uploadFileMutation = useMutation(uploadFile());
   const uploadedFile = uploadFileMutation.data;

diff --git a/site/src/pages/CreateUserPage/CreateUserPage.tsx b/site/src/pages/CreateUserPage/CreateUserPage.tsx
@@ -5,7 +5,7 @@ import { useNavigate } from "react-router-dom";
 import { authMethods, createUser } from "api/queries/users";
 import { displaySuccess } from "components/GlobalSnackbar/utils";
 import { Margins } from "components/Margins/Margins";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { pageTitle } from "utils/page";
 import { CreateUserForm } from "./CreateUserForm";
 
@@ -14,7 +14,7 @@ export const Language = {
 };
 
 export const CreateUserPage: FC = () => {
-  const { organizationId } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const navigate = useNavigate();
   const queryClient = useQueryClient();
   const createUserMutation = useMutation(createUser(queryClient));

diff --git a/site/src/pages/CreateWorkspacePage/CreateWorkspacePage.tsx b/site/src/pages/CreateWorkspacePage/CreateWorkspacePage.tsx
@@ -35,10 +35,10 @@ export type ExternalAuthPollingState = "idle" | "polling" | "abandoned";
 
 const CreateWorkspacePage: FC = () => {
   const { template: templateName } = useParams() as { template: string };
-  const { user: me, organizationId } = useAuthenticated();
+  const { user: me } = useAuthenticated();
   const navigate = useNavigate();
   const [searchParams, setSearchParams] = useSearchParams();
-  const { experiments } = useDashboard();
+  const { experiments, organizationId } = useDashboard();
 
   const customVersionId = searchParams.get("version") ?? undefined;
   const defaultName = searchParams.get("name");

diff --git a/site/src/pages/GroupsPage/CreateGroupPage.tsx b/site/src/pages/GroupsPage/CreateGroupPage.tsx
@@ -3,14 +3,14 @@ import { Helmet } from "react-helmet-async";
 import { useMutation, useQueryClient } from "react-query";
 import { useNavigate } from "react-router-dom";
 import { createGroup } from "api/queries/groups";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { pageTitle } from "utils/page";
 import CreateGroupPageView from "./CreateGroupPageView";
 
 export const CreateGroupPage: FC = () => {
   const queryClient = useQueryClient();
   const navigate = useNavigate();
-  const { organizationId } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const createGroupMutation = useMutation(createGroup(queryClient));
 
   return (

diff --git a/site/src/pages/GroupsPage/GroupsPage.tsx b/site/src/pages/GroupsPage/GroupsPage.tsx
@@ -5,12 +5,14 @@ import { getErrorMessage } from "api/errors";
 import { groups } from "api/queries/groups";
 import { displayError } from "components/GlobalSnackbar/utils";
 import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { useFeatureVisibility } from "modules/dashboard/useFeatureVisibility";
 import { pageTitle } from "utils/page";
 import GroupsPageView from "./GroupsPageView";
 
 export const GroupsPage: FC = () => {
-  const { organizationId, permissions } = useAuthenticated();
+  const { permissions } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const { createGroup: canCreateGroup } = permissions;
   const { template_rbac: isTemplateRBACEnabled } = useFeatureVisibility();
   const groupsQuery = useQuery(groups(organizationId));

diff --git a/site/src/pages/StarterTemplatePage/StarterTemplatePage.tsx b/site/src/pages/StarterTemplatePage/StarterTemplatePage.tsx
@@ -3,13 +3,13 @@ import { Helmet } from "react-helmet-async";
 import { useQuery } from "react-query";
 import { useParams } from "react-router-dom";
 import { templateExamples } from "api/queries/templates";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { pageTitle } from "utils/page";
 import { StarterTemplatePageView } from "./StarterTemplatePageView";
 
 const StarterTemplatePage: FC = () => {
   const { exampleId } = useParams() as { exampleId: string };
-  const { organizationId } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const templateExamplesQuery = useQuery(templateExamples(organizationId));
   const starterTemplate = templateExamplesQuery.data?.find(
     (example) => example.id === exampleId,

diff --git a/site/src/pages/StarterTemplatesPage/StarterTemplatesPage.tsx b/site/src/pages/StarterTemplatesPage/StarterTemplatesPage.tsx
@@ -3,13 +3,13 @@ import { Helmet } from "react-helmet-async";
 import { useQuery } from "react-query";
 import { templateExamples } from "api/queries/templates";
 import type { TemplateExample } from "api/typesGenerated";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { pageTitle } from "utils/page";
 import { getTemplatesByTag } from "utils/starterTemplates";
 import { StarterTemplatesPageView } from "./StarterTemplatesPageView";
 
 const StarterTemplatesPage: FC = () => {
-  const { organizationId } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const templateExamplesQuery = useQuery(templateExamples(organizationId));
   const starterTemplatesByTag = templateExamplesQuery.data
     ? // Currently, the scratch template should not be displayed on the starter templates page.

diff --git a/site/src/pages/TemplatePage/TemplateFilesPage/TemplateFilesPage.tsx b/site/src/pages/TemplatePage/TemplateFilesPage/TemplateFilesPage.tsx
@@ -3,13 +3,13 @@ import { Helmet } from "react-helmet-async";
 import { useQuery } from "react-query";
 import { previousTemplateVersion, templateFiles } from "api/queries/templates";
 import { Loader } from "components/Loader/Loader";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { TemplateFiles } from "modules/templates/TemplateFiles/TemplateFiles";
 import { useTemplateLayoutContext } from "pages/TemplatePage/TemplateLayout";
 import { getTemplatePageTitle } from "../utils";
 
 const TemplateFilesPage: FC = () => {
-  const { organizationId } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const { template, activeVersion } = useTemplateLayoutContext();
   const { data: currentFiles } = useQuery(
     templateFiles(activeVersion.job.file_id),

diff --git a/site/src/pages/TemplatePage/TemplateLayout.tsx b/site/src/pages/TemplatePage/TemplateLayout.tsx
@@ -13,7 +13,7 @@ import { ErrorAlert } from "components/Alert/ErrorAlert";
 import { Loader } from "components/Loader/Loader";
 import { Margins } from "components/Margins/Margins";
 import { TAB_PADDING_Y, TabLink, Tabs, TabsList } from "components/Tabs/Tabs";
-import { useAuthenticated } from "contexts/auth/RequireAuth";
+import { useDashboard } from "modules/dashboard/useDashboard";
 import { TemplatePageHeader } from "./TemplatePageHeader";
 
 const templatePermissions = (
@@ -71,7 +71,7 @@ export const TemplateLayout: FC<PropsWithChildren> = ({
   children = <Outlet />,
 }) => {
   const navigate = useNavigate();
-  const { organizationId } = useAuthenticated();
+  const { organizationId } = useDashboard();
   const { template: templateName } = useParams() as { template: string };
   const { data, error, isLoading } = useQuery({
     queryKey: ["template", templateName],