chore(site): refactor popover to make it easier to extend (#13611)

site/src/components/HelpTooltip/HelpTooltip.tsx

@@ -160,7 +160,7 @@ export const HelpTooltipAction: FC<HelpTooltipActionProps> = ({
       onClick={(event) => {
         event.stopPropagation();
         onClick();
-        popover.setIsOpen(false);
+        popover.setOpen(false);
       }}
     >
       <Icon css={styles.actionIcon} />

site/src/components/IconField/IconField.tsx

@@ -3,7 +3,7 @@ import Button from "@mui/material/Button";
 import InputAdornment from "@mui/material/InputAdornment";
 import TextField, { type TextFieldProps } from "@mui/material/TextField";
 import { visuallyHidden } from "@mui/utils";
-import { type FC, lazy, Suspense } from "react";
+import { type FC, lazy, Suspense, useState } from "react";
 import { DropdownArrow } from "components/DropdownArrow/DropdownArrow";
 import { ExternalImage } from "components/ExternalImage/ExternalImage";
 import { Loader } from "components/Loader/Loader";
@@ -37,6 +37,7 @@ export const IconField: FC<IconFieldProps> = ({
 
   const theme = useTheme();
   const hasIcon = textFieldProps.value && textFieldProps.value !== "";
+  const [open, setOpen] = useState(false);
 
   return (
     <Stack spacing={1}>
@@ -86,9 +87,7 @@ export const IconField: FC<IconFieldProps> = ({
           }
         `}
       />
-      <Popover>
-        {(popover) => (
-          <>
+      <Popover open={open} onOpenChange={setOpen}>
         <PopoverTrigger>
           <Button fullWidth endIcon={<DropdownArrow />}>
             Select emoji
@@ -101,16 +100,13 @@ export const IconField: FC<IconFieldProps> = ({
           <Suspense fallback={<Loader />}>
             <EmojiPicker
               onEmojiSelect={(emoji) => {
-                    const value =
-                      emoji.src ?? urlFromUnifiedCode(emoji.unified);
+                const value = emoji.src ?? urlFromUnifiedCode(emoji.unified);
                 onPickEmoji(value);
-                    popover.setIsOpen(false);
+                setOpen(false);
               }}
             />
           </Suspense>
         </PopoverContent>
-          </>
-        )}
       </Popover>
 
       {/*

site/src/components/Popover/Popover.tsx

@@ -12,7 +12,6 @@ import {
   type ReactNode,
   type RefObject,
   useContext,
-  useEffect,
   useId,
   useRef,
   useState,
@@ -25,14 +24,12 @@ type TriggerRef = RefObject<HTMLElement>;
 type TriggerElement = ReactElement<{
   ref: TriggerRef;
   onClick?: () => void;
-  "aria-haspopup"?: boolean;
-  "aria-owns"?: string | undefined;
 }>;
 
 type PopoverContextValue = {
   id: string;
-  isOpen: boolean;
-  setIsOpen: React.Dispatch<React.SetStateAction<boolean>>;
+  open: boolean;
+  setOpen: (open: boolean) => void;
   triggerRef: TriggerRef;
   mode: TriggerMode;
 };
@@ -41,32 +38,41 @@ const PopoverContext = createContext<PopoverContextValue | undefined>(
   undefined,
 );
 
-export interface PopoverProps {
-  children: ReactNode | ((popover: PopoverContextValue) => ReactNode); // Allows inline usage
+type BasePopoverProps = {
+  children: ReactNode;
   mode?: TriggerMode;
-  isDefaultOpen?: boolean;
-}
+};
 
-export const Popover: FC<PopoverProps> = ({
-  children,
-  mode,
-  isDefaultOpen,
-}) => {
+// By separating controlled and uncontrolled props, we achieve more accurate
+// type inference.
+type UncontrolledPopoverProps = BasePopoverProps & {
+  open?: undefined;
+  onOpenChange?: undefined;
+};
+
+type ControlledPopoverProps = BasePopoverProps & {
+  open: boolean;
+  onOpenChange: (open: boolean) => void;
+};
+
+export type PopoverProps = UncontrolledPopoverProps | ControlledPopoverProps;
+
+export const Popover: FC<PopoverProps> = (props) => {
   const hookId = useId();
-  const [isOpen, setIsOpen] = useState(isDefaultOpen ?? false);
-  const triggerRef = useRef<HTMLElement>(null);
+  const [uncontrolledOpen, setUncontrolledOpen] = useState(false);
+  const triggerRef: TriggerRef = useRef(null);
 
   const value: PopoverContextValue = {
-    isOpen,
-    setIsOpen,
     triggerRef,
     id: `${hookId}-popover`,
-    mode: mode ?? "click",
+    mode: props.mode ?? "click",
+    open: props.open ?? uncontrolledOpen,
+    setOpen: props.onOpenChange ?? setUncontrolledOpen,
   };
 
   return (
     <PopoverContext.Provider value={value}>
-      {typeof children === "function" ? children(value) : children}
+      {props.children}
     </PopoverContext.Provider>
   );
 };
@@ -82,23 +88,25 @@ export const usePopover = () => {
 };
 
 export const PopoverTrigger = (
-  props: HTMLAttributes<HTMLElement> & { children: TriggerElement },
+  props: HTMLAttributes<HTMLElement> & {
+    children: TriggerElement;
+  },
 ) => {
   const popover = usePopover();
   const { children, ...elementProps } = props;
 
   const clickProps = {
     onClick: () => {
-      popover.setIsOpen((isOpen) => !isOpen);
+      popover.setOpen(true);
     },
   };
 
   const hoverProps = {
     onPointerEnter: () => {
-      popover.setIsOpen(true);
+      popover.setOpen(true);
     },
     onPointerLeave: () => {
-      popover.setIsOpen(false);
+      popover.setOpen(false);
     },
   };
 
@@ -106,7 +114,8 @@ export const PopoverTrigger = (
     ...elementProps,
     ...(popover.mode === "click" ? clickProps : hoverProps),
     "aria-haspopup": true,
-    "aria-owns": popover.isOpen ? popover.id : undefined,
+    "aria-owns": popover.id,
+    "aria-expanded": popover.open,
     ref: popover.triggerRef,
   });
 };
@@ -125,22 +134,8 @@ export const PopoverContent: FC<PopoverContentProps> = ({
   ...popoverProps
 }) => {
   const popover = usePopover();
-  const [isReady, setIsReady] = useState(false);
   const hoverMode = popover.mode === "hover";
 
-  // This is a hack to make sure the popover is not rendered until the trigger
-  // is ready. This is a limitation on MUI that does not support defaultIsOpen
-  // on Popover but we need it to storybook the component.
-  useEffect(() => {
-    if (!isReady && popover.triggerRef.current !== null) {
-      setIsReady(true);
-    }
-  }, [isReady, popover.triggerRef]);
-
-  if (!popover.triggerRef.current) {
-    return null;
-  }
-
   return (
     <MuiPopover
       disablePortal
@@ -161,8 +156,8 @@ export const PopoverContent: FC<PopoverContentProps> = ({
       {...modeProps(popover)}
       {...popoverProps}
       id={popover.id}
-      open={popover.isOpen}
-      onClose={() => popover.setIsOpen(false)}
+      open={popover.open}
+      onClose={() => popover.setOpen(false)}
       anchorEl={popover.triggerRef.current}
     />
   );
@@ -172,10 +167,10 @@ const modeProps = (popover: PopoverContextValue) => {
   if (popover.mode === "hover") {
     return {
       onPointerEnter: () => {
-        popover.setIsOpen(true);
+        popover.setOpen(true);
       },
       onPointerLeave: () => {
-        popover.setIsOpen(false);
+        popover.setOpen(false);
       },
     };
   }

site/src/modules/dashboard/Navbar/DeploymentDropdown.tsx

@@ -87,7 +87,7 @@ const DeploymentDropdownContent: FC<DeploymentDropdownProps> = ({
 }) => {
   const popover = usePopover();
 
-  const onPopoverClose = () => popover.setIsOpen(false);
+  const onPopoverClose = () => popover.setOpen(false);
 
   return (
     <nav>

site/src/modules/dashboard/Navbar/ProxyMenu.stories.tsx

@@ -11,6 +11,7 @@ import {
   MockUser,
   MockWorkspaceProxies,
 } from "testHelpers/entities";
+import { withDesktopViewport } from "testHelpers/storybook";
 import { ProxyMenu } from "./ProxyMenu";
 
 const defaultProxyContextValue = {
@@ -36,11 +37,7 @@ const meta: Meta<typeof ProxyMenu> = {
         <Story />
       </AuthProvider>
     ),
-    (Story) => (
-      <div css={{ width: 1200, height: 800 }}>
-        <Story />
-      </div>
-    ),
+    withDesktopViewport,
   ],
   parameters: {
     queries: [

site/src/modules/dashboard/Navbar/UserDropdown/UserDropdown.tsx

@@ -1,6 +1,6 @@
 import { css, type Interpolation, type Theme, useTheme } from "@emotion/react";
 import Badge from "@mui/material/Badge";
-import type { FC } from "react";
+import { useState, type FC } from "react";
 import type * as TypesGen from "api/typesGenerated";
 import { DropdownArrow } from "components/DropdownArrow/DropdownArrow";
 import {
@@ -26,11 +26,10 @@ export const UserDropdown: FC<UserDropdownProps> = ({
   onSignOut,
 }) => {
   const theme = useTheme();
+  const [open, setOpen] = useState(false);
 
   return (
-    <Popover>
-      {(popover) => (
-        <>
+    <Popover open={open} onOpenChange={setOpen}>
       <PopoverTrigger>
         <button css={styles.button} data-testid="user-dropdown-trigger">
           <div css={styles.badgeContainer}>
@@ -43,7 +42,7 @@ export const UserDropdown: FC<UserDropdownProps> = ({
             </Badge>
             <DropdownArrow
               color={theme.experimental.l2.fill.solid}
-                  close={popover.isOpen}
+              close={open}
             />
           </div>
         </button>
@@ -66,8 +65,6 @@ export const UserDropdown: FC<UserDropdownProps> = ({
           onSignOut={onSignOut}
         />
       </PopoverContent>
-        </>
-      )}
     </Popover>
   );
 };

site/src/modules/dashboard/Navbar/UserDropdown/UserDropdownContent.tsx

@@ -43,7 +43,7 @@ export const UserDropdownContent: FC<UserDropdownContentProps> = ({
   const popover = usePopover();
 
   const onPopoverClose = () => {
-    popover.setIsOpen(false);
+    popover.setOpen(false);
   };
 
   const renderMenuIcon = (icon: string): JSX.Element => {

site/src/modules/resources/SSHButton/SSHButton.stories.tsx

@@ -1,5 +1,7 @@
 import type { Meta, StoryObj } from "@storybook/react";
+import { userEvent, within } from "@storybook/test";
 import { MockWorkspace, MockWorkspaceAgent } from "testHelpers/entities";
+import { withDesktopViewport } from "testHelpers/storybook";
 import { SSHButton } from "./SSHButton";
 
 const meta: Meta<typeof SSHButton> = {
@@ -22,7 +24,12 @@ export const Opened: Story = {
   args: {
     workspaceName: MockWorkspace.name,
     agentName: MockWorkspaceAgent.name,
-    isDefaultOpen: true,
     sshPrefix: "coder.",
   },
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    const button = canvas.getByRole("button");
+    await userEvent.click(button);
+  },
+  decorators: [withDesktopViewport],
 };

site/src/modules/resources/SSHButton/SSHButton.tsx

@@ -20,20 +20,18 @@ import { docs } from "utils/docs";
 export interface SSHButtonProps {
   workspaceName: string;
   agentName: string;
-  isDefaultOpen?: boolean;
   sshPrefix?: string;
 }
 
 export const SSHButton: FC<SSHButtonProps> = ({
   workspaceName,
   agentName,
-  isDefaultOpen = false,
   sshPrefix,
 }) => {
   const paper = useClassName(classNames.paper, []);
 
   return (
-    <Popover isDefaultOpen={isDefaultOpen}>
+    <Popover>
       <PopoverTrigger>
         <Button
           size="small"

site/src/modules/workspaces/WorkspaceOutdatedTooltip/WorkspaceOutdatedTooltip.tsx

@@ -56,7 +56,7 @@ export const WorkspaceOutdatedTooltipContent: FC<TooltipProps> = ({
   const popover = usePopover();
   const { data: activeVersion } = useQuery({
     ...templateVersion(latestVersionId),
-    enabled: popover.isOpen,
+    enabled: popover.open,
   });
   const theme = useTheme();
 

site/src/pages/TemplatePage/TemplateInsightsPage/DateRange.tsx

@@ -49,11 +49,10 @@ export const DateRange: FC<DateRangeProps> = ({ value, onChange }) => {
       key: "selection",
     },
   ]);
+  const [open, setOpen] = useState(false);
 
   return (
-    <Popover>
-      {(popover) => (
-        <>
+    <Popover open={open} onOpenChange={setOpen}>
       <PopoverTrigger>
         <Button>
           <span>{format(value.startDate, "MMM d, Y")}</span>
@@ -87,7 +86,7 @@ export const DateRange: FC<DateRangeProps> = ({ value, onChange }) => {
                 ? startOfHour(addHours(now, 1))
                 : startOfDay(addDays(endDate, 1)),
             });
-                popover.setIsOpen(false);
+            setOpen(false);
           }}
           moveRangeOnFirstSelection={false}
           months={2}
@@ -133,8 +132,6 @@ export const DateRange: FC<DateRangeProps> = ({ value, onChange }) => {
           ])}
         />
       </PopoverContent>
-        </>
-      )}
     </Popover>
   );
 };

site/src/pages/UsersPage/UsersTable/EditRolesButton.stories.tsx

@@ -1,43 +1,43 @@
 import type { Meta, StoryObj } from "@storybook/react";
+import { userEvent, within } from "@storybook/test";
 import {
   MockOwnerRole,
   MockSiteRoles,
   MockUserAdminRole,
 } from "testHelpers/entities";
+import { withDesktopViewport } from "testHelpers/storybook";
 import { EditRolesButton } from "./EditRolesButton";
 
 const meta: Meta<typeof EditRolesButton> = {
   title: "pages/UsersPage/EditRolesButton",
   component: EditRolesButton,
   args: {
-    isDefaultOpen: true,
+    selectedRoleNames: new Set([MockUserAdminRole.name, MockOwnerRole.name]),
+    roles: MockSiteRoles,
   },
+  decorators: [withDesktopViewport],
 };
 
 export default meta;
 type Story = StoryObj<typeof EditRolesButton>;
 
-const selectedRoleNames = new Set([MockUserAdminRole.name, MockOwnerRole.name]);
+export const Closed: Story = {};
 
 export const Open: Story = {
-  args: {
-    selectedRoleNames,
-    roles: MockSiteRoles,
-  },
-  parameters: {
-    chromatic: { delay: 300 },
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    await userEvent.click(canvas.getByRole("button"));
   },
 };
 
 export const Loading: Story = {
   args: {
     isLoading: true,
-    selectedRoleNames,
-    roles: MockSiteRoles,
     userLoginType: "password",
     oidcRoleSync: false,
   },
-  parameters: {
-    chromatic: { delay: 300 },
+  play: async ({ canvasElement }) => {
+    const canvas = within(canvasElement);
+    await userEvent.click(canvas.getByRole("button"));
   },
 };

site/src/pages/UsersPage/UsersTable/EditRolesButton.tsx

@@ -72,7 +72,6 @@ export interface EditRolesButtonProps {
   roles: readonly SlimRole[];
   selectedRoleNames: Set<string>;
   onChange: (roles: SlimRole["name"][]) => void;
-  isDefaultOpen?: boolean;
   oidcRoleSync: boolean;
   userLoginType: string;
 }
@@ -82,7 +81,6 @@ export const EditRolesButton: FC<EditRolesButtonProps> = ({
   selectedRoleNames,
   onChange,
   isLoading,
-  isDefaultOpen = false,
   userLoginType,
   oidcRoleSync,
 }) => {
@@ -116,7 +114,7 @@ export const EditRolesButton: FC<EditRolesButtonProps> = ({
   }
 
   return (
-    <Popover isDefaultOpen={isDefaultOpen}>
+    <Popover>
       <PopoverTrigger>
         <IconButton
           size="small"

site/src/pages/WorkspacePage/WorkspaceActions/BuildParametersPopover.tsx

@@ -117,7 +117,7 @@ const BuildParametersPopoverContent: FC<BuildParametersPopoverContentProps> = ({
               <Form
                 onSubmit={(buildParameters) => {
                   onSubmit(buildParameters);
-                  popover.setIsOpen(false);
+                  popover.setOpen(false);
                 }}
                 ephemeralParameters={ephemeralParameters}
                 buildParameters={buildParameters.map(

site/src/pages/WorkspacePage/WorkspaceActions/DownloadLogsDialog.stories.tsx

@@ -2,6 +2,7 @@ import type { Meta, StoryObj } from "@storybook/react";
 import { expect, userEvent, fn, waitFor, within } from "@storybook/test";
 import { agentLogsKey, buildLogsKey } from "api/queries/workspaces";
 import { MockWorkspace, MockWorkspaceAgent } from "testHelpers/entities";
+import { withDesktopViewport } from "testHelpers/storybook";
 import { DownloadLogsDialog } from "./DownloadLogsDialog";
 
 const meta: Meta<typeof DownloadLogsDialog> = {
@@ -24,13 +25,7 @@ const meta: Meta<typeof DownloadLogsDialog> = {
       },
     ],
   },
-  decorators: [
-    (Story) => (
-      <div css={{ width: 1200, height: 800 }}>
-        <Story />
-      </div>
-    ),
-  ],
+  decorators: [withDesktopViewport],
 };
 
 export default meta;

site/src/pages/WorkspacePage/WorkspaceActions/WorkspaceActions.stories.tsx

@@ -2,6 +2,7 @@ import type { Meta, StoryObj } from "@storybook/react";
 import { userEvent, within, expect } from "@storybook/test";
 import { buildLogsKey, agentLogsKey } from "api/queries/workspaces";
 import * as Mocks from "testHelpers/entities";
+import { withDesktopViewport } from "testHelpers/storybook";
 import { WorkspaceActions } from "./WorkspaceActions";
 
 const meta: Meta<typeof WorkspaceActions> = {
@@ -10,13 +11,7 @@ const meta: Meta<typeof WorkspaceActions> = {
   args: {
     isUpdating: false,
   },
-  decorators: [
-    (Story) => (
-      <div css={{ width: 1200, height: 800 }}>
-        <Story />
-      </div>
-    ),
-  ],
+  decorators: [withDesktopViewport],
 };
 
 export default meta;

site/src/pages/WorkspacePage/WorkspaceNotifications/Notifications.tsx

@@ -70,7 +70,7 @@ const NotificationPill: FC<NotificationsProps> = ({
       icon={icon}
       css={(theme) => ({
         "& svg": { color: theme.roles[severity].outline },
-        borderColor: popover.isOpen ? theme.roles[severity].outline : undefined,
+        borderColor: popover.open ? theme.roles[severity].outline : undefined,
       })}
     >
       {items.length}

site/src/testHelpers/storybook.tsx

@@ -78,3 +78,9 @@ export const withWebSocket = (Story: FC, { parameters }: StoryContext) => {
 
   return <Story />;
 };
+
+export const withDesktopViewport = (Story: FC) => (
+  <div style={{ width: 1200, height: 800 }}>
+    <Story />
+  </div>
+);