From 4639dc3255b0ccb11c17cfd54d1ec98ffcd2a3ad Mon Sep 17 00:00:00 2001
From: Dustin Healy <54083382+dustinhealy@users.noreply.github.com>
Date: Mon, 28 Jul 2025 09:26:11 -0700
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9C=20fix:=20Forward=20Ref=20to=20`MCP?=
 =?UTF-8?q?SubMenu`=20and=20`ArtifactsSubMenu`=20(#8696)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

ToolsDropdown uses a menu library that passes refs to submenu items. Function components can't receive refs by default though, so we get  "Function components cannot be given refs" warnings in the console. React.forwardRef() allows them to properly handle ref forwarding by wrapping the component and attaching the ref to the outer div element.
---
 .../Chat/Input/ArtifactsSubMenu.tsx           | 247 +++++++++---------
 .../src/components/Chat/Input/MCPSubMenu.tsx  | 208 +++++++--------
 2 files changed, 232 insertions(+), 223 deletions(-)

diff --git a/client/src/components/Chat/Input/ArtifactsSubMenu.tsx b/client/src/components/Chat/Input/ArtifactsSubMenu.tsx
index 654c129fba..e27fa43c0e 100644
--- a/client/src/components/Chat/Input/ArtifactsSubMenu.tsx
+++ b/client/src/components/Chat/Input/ArtifactsSubMenu.tsx
@@ -15,133 +15,142 @@ interface ArtifactsSubMenuProps {
   handleCustomToggle: () => void;
 }
 
-const ArtifactsSubMenu = ({
-  isArtifactsPinned,
-  setIsArtifactsPinned,
-  artifactsMode,
-  handleArtifactsToggle,
-  handleShadcnToggle,
-  handleCustomToggle,
-  ...props
-}: ArtifactsSubMenuProps) => {
-  const localize = useLocalize();
+const ArtifactsSubMenu = React.forwardRef<HTMLDivElement, ArtifactsSubMenuProps>(
+  (
+    {
+      isArtifactsPinned,
+      setIsArtifactsPinned,
+      artifactsMode,
+      handleArtifactsToggle,
+      handleShadcnToggle,
+      handleCustomToggle,
+      ...props
+    },
+    ref,
+  ) => {
+    const localize = useLocalize();
 
-  const menuStore = Ariakit.useMenuStore({
-    focusLoop: true,
-    showTimeout: 100,
-    placement: 'right',
-  });
+    const menuStore = Ariakit.useMenuStore({
+      focusLoop: true,
+      showTimeout: 100,
+      placement: 'right',
+    });
 
-  const isEnabled = artifactsMode !== '' && artifactsMode !== undefined;
-  const isShadcnEnabled = artifactsMode === ArtifactModes.SHADCNUI;
-  const isCustomEnabled = artifactsMode === ArtifactModes.CUSTOM;
+    const isEnabled = artifactsMode !== '' && artifactsMode !== undefined;
+    const isShadcnEnabled = artifactsMode === ArtifactModes.SHADCNUI;
+    const isCustomEnabled = artifactsMode === ArtifactModes.CUSTOM;
 
-  return (
-    <Ariakit.MenuProvider store={menuStore}>
-      <Ariakit.MenuItem
-        {...props}
-        hideOnClick={false}
-        render={
-          <Ariakit.MenuButton
-            onClick={(e: React.MouseEvent<HTMLButtonElement>) => {
-              e.stopPropagation();
-              handleArtifactsToggle();
-            }}
-            onMouseEnter={() => {
-              if (isEnabled) {
-                menuStore.show();
-              }
-            }}
-            className="flex w-full cursor-pointer items-center justify-between rounded-lg p-2 hover:bg-surface-hover"
-          />
-        }
-      >
-        <div className="flex items-center gap-2">
-          <WandSparkles className="icon-md" />
-          <span>{localize('com_ui_artifacts')}</span>
-          {isEnabled && <ChevronRight className="ml-auto h-3 w-3" />}
-        </div>
-        <button
-          type="button"
-          onClick={(e) => {
-            e.stopPropagation();
-            setIsArtifactsPinned(!isArtifactsPinned);
-          }}
-          className={cn(
-            'rounded p-1 transition-all duration-200',
-            'hover:bg-surface-tertiary hover:shadow-sm',
-            !isArtifactsPinned && 'text-text-secondary hover:text-text-primary',
-          )}
-          aria-label={isArtifactsPinned ? 'Unpin' : 'Pin'}
-        >
-          <div className="h-4 w-4">
-            <PinIcon unpin={isArtifactsPinned} />
-          </div>
-        </button>
-      </Ariakit.MenuItem>
-
-      {isEnabled && (
-        <Ariakit.Menu
-          portal={true}
-          unmountOnHide={true}
-          className={cn(
-            'animate-popover-left z-50 ml-3 flex min-w-[250px] flex-col rounded-xl',
-            'border border-border-light bg-surface-secondary px-1.5 py-1 shadow-lg',
-          )}
-        >
-          <div className="px-2 py-1.5">
-            <div className="mb-2 text-xs font-medium text-text-secondary">
-              {localize('com_ui_artifacts_options')}
+    return (
+      <div ref={ref}>
+        <Ariakit.MenuProvider store={menuStore}>
+          <Ariakit.MenuItem
+            {...props}
+            hideOnClick={false}
+            render={
+              <Ariakit.MenuButton
+                onClick={(e: React.MouseEvent<HTMLButtonElement>) => {
+                  e.stopPropagation();
+                  handleArtifactsToggle();
+                }}
+                onMouseEnter={() => {
+                  if (isEnabled) {
+                    menuStore.show();
+                  }
+                }}
+                className="flex w-full cursor-pointer items-center justify-between rounded-lg p-2 hover:bg-surface-hover"
+              />
+            }
+          >
+            <div className="flex items-center gap-2">
+              <WandSparkles className="icon-md" />
+              <span>{localize('com_ui_artifacts')}</span>
+              {isEnabled && <ChevronRight className="ml-auto h-3 w-3" />}
             </div>
-
-            {/* Include shadcn/ui Option */}
-            <Ariakit.MenuItem
-              hideOnClick={false}
-              onClick={(event) => {
-                event.preventDefault();
-                event.stopPropagation();
-                handleShadcnToggle();
-              }}
-              disabled={isCustomEnabled}
-              className={cn(
-                'mb-1 flex items-center justify-between rounded-lg px-2 py-2',
-                'cursor-pointer text-text-primary outline-none transition-colors',
-                'hover:bg-black/[0.075] dark:hover:bg-white/10',
-                'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
-                isCustomEnabled && 'cursor-not-allowed opacity-50',
-              )}
-            >
-              <div className="flex items-center gap-2">
-                <Ariakit.MenuItemCheck checked={isShadcnEnabled} />
-                <span className="text-sm">{localize('com_ui_include_shadcnui' as any)}</span>
-              </div>
-            </Ariakit.MenuItem>
-
-            {/* Custom Prompt Mode Option */}
-            <Ariakit.MenuItem
-              hideOnClick={false}
-              onClick={(event) => {
-                event.preventDefault();
-                event.stopPropagation();
-                handleCustomToggle();
+            <button
+              type="button"
+              onClick={(e) => {
+                e.stopPropagation();
+                setIsArtifactsPinned(!isArtifactsPinned);
               }}
               className={cn(
-                'flex items-center justify-between rounded-lg px-2 py-2',
-                'cursor-pointer text-text-primary outline-none transition-colors',
-                'hover:bg-black/[0.075] dark:hover:bg-white/10',
-                'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
+                'rounded p-1 transition-all duration-200',
+                'hover:bg-surface-tertiary hover:shadow-sm',
+                !isArtifactsPinned && 'text-text-secondary hover:text-text-primary',
+              )}
+              aria-label={isArtifactsPinned ? 'Unpin' : 'Pin'}
+            >
+              <div className="h-4 w-4">
+                <PinIcon unpin={isArtifactsPinned} />
+              </div>
+            </button>
+          </Ariakit.MenuItem>
+
+          {isEnabled && (
+            <Ariakit.Menu
+              portal={true}
+              unmountOnHide={true}
+              className={cn(
+                'animate-popover-left z-50 ml-3 flex min-w-[250px] flex-col rounded-xl',
+                'border border-border-light bg-surface-secondary px-1.5 py-1 shadow-lg',
               )}
             >
-              <div className="flex items-center gap-2">
-                <Ariakit.MenuItemCheck checked={isCustomEnabled} />
-                <span className="text-sm">{localize('com_ui_custom_prompt_mode' as any)}</span>
+              <div className="px-2 py-1.5">
+                <div className="mb-2 text-xs font-medium text-text-secondary">
+                  {localize('com_ui_artifacts_options')}
+                </div>
+
+                {/* Include shadcn/ui Option */}
+                <Ariakit.MenuItem
+                  hideOnClick={false}
+                  onClick={(event) => {
+                    event.preventDefault();
+                    event.stopPropagation();
+                    handleShadcnToggle();
+                  }}
+                  disabled={isCustomEnabled}
+                  className={cn(
+                    'mb-1 flex items-center justify-between rounded-lg px-2 py-2',
+                    'cursor-pointer text-text-primary outline-none transition-colors',
+                    'hover:bg-black/[0.075] dark:hover:bg-white/10',
+                    'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
+                    isCustomEnabled && 'cursor-not-allowed opacity-50',
+                  )}
+                >
+                  <div className="flex items-center gap-2">
+                    <Ariakit.MenuItemCheck checked={isShadcnEnabled} />
+                    <span className="text-sm">{localize('com_ui_include_shadcnui' as any)}</span>
+                  </div>
+                </Ariakit.MenuItem>
+
+                {/* Custom Prompt Mode Option */}
+                <Ariakit.MenuItem
+                  hideOnClick={false}
+                  onClick={(event) => {
+                    event.preventDefault();
+                    event.stopPropagation();
+                    handleCustomToggle();
+                  }}
+                  className={cn(
+                    'flex items-center justify-between rounded-lg px-2 py-2',
+                    'cursor-pointer text-text-primary outline-none transition-colors',
+                    'hover:bg-black/[0.075] dark:hover:bg-white/10',
+                    'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
+                  )}
+                >
+                  <div className="flex items-center gap-2">
+                    <Ariakit.MenuItemCheck checked={isCustomEnabled} />
+                    <span className="text-sm">{localize('com_ui_custom_prompt_mode' as any)}</span>
+                  </div>
+                </Ariakit.MenuItem>
               </div>
-            </Ariakit.MenuItem>
-          </div>
-        </Ariakit.Menu>
-      )}
-    </Ariakit.MenuProvider>
-  );
-};
+            </Ariakit.Menu>
+          )}
+        </Ariakit.MenuProvider>
+      </div>
+    );
+  },
+);
+
+ArtifactsSubMenu.displayName = 'ArtifactsSubMenu';
 
 export default React.memo(ArtifactsSubMenu);
diff --git a/client/src/components/Chat/Input/MCPSubMenu.tsx b/client/src/components/Chat/Input/MCPSubMenu.tsx
index 1d704f2a6d..3628b732a1 100644
--- a/client/src/components/Chat/Input/MCPSubMenu.tsx
+++ b/client/src/components/Chat/Input/MCPSubMenu.tsx
@@ -11,115 +11,115 @@ interface MCPSubMenuProps {
   placeholder?: string;
 }
 
-const MCPSubMenu = ({ placeholder, ...props }: MCPSubMenuProps) => {
-  const {
-    configuredServers,
-    mcpValues,
-    isPinned,
-    setIsPinned,
-    placeholderText,
-    toggleServerSelection,
-    getServerStatusIconProps,
-    getConfigDialogProps,
-  } = useMCPServerManager();
+const MCPSubMenu = React.forwardRef<HTMLDivElement, MCPSubMenuProps>(
+  ({ placeholder, ...props }, ref) => {
+    const {
+      configuredServers,
+      mcpValues,
+      isPinned,
+      setIsPinned,
+      placeholderText,
+      toggleServerSelection,
+      getServerStatusIconProps,
+      getConfigDialogProps,
+    } = useMCPServerManager();
 
-  const menuStore = Ariakit.useMenuStore({
-    focusLoop: true,
-    showTimeout: 100,
-    placement: 'right',
-  });
+    const menuStore = Ariakit.useMenuStore({
+      focusLoop: true,
+      showTimeout: 100,
+      placement: 'right',
+    });
 
-  // Don't render if no MCP servers are configured
-  if (!configuredServers || configuredServers.length === 0) {
-    return null;
-  }
+    // Don't render if no MCP servers are configured
+    if (!configuredServers || configuredServers.length === 0) {
+      return null;
+    }
 
-  const configDialogProps = getConfigDialogProps();
+    const configDialogProps = getConfigDialogProps();
 
-  return (
-    <>
-      <Ariakit.MenuProvider store={menuStore}>
-        <Ariakit.MenuItem
-          {...props}
-          render={
-            <Ariakit.MenuButton
-              onClick={(e: React.MouseEvent<HTMLButtonElement>) => {
-                e.stopPropagation();
-                menuStore.toggle();
-              }}
-              className="flex w-full cursor-pointer items-center justify-between rounded-lg p-2 hover:bg-surface-hover"
-            />
-          }
-        >
-          <div className="flex items-center gap-2">
-            <MCPIcon className="icon-md" />
-            <span>{placeholder || placeholderText}</span>
-            <ChevronRight className="ml-auto h-3 w-3" />
-          </div>
-          <button
-            type="button"
-            onClick={(e) => {
-              e.stopPropagation();
-              setIsPinned(!isPinned);
-            }}
-            className={cn(
-              'rounded p-1 transition-all duration-200',
-              'hover:bg-surface-tertiary hover:shadow-sm',
-              !isPinned && 'text-text-secondary hover:text-text-primary',
-            )}
-            aria-label={isPinned ? 'Unpin' : 'Pin'}
-          >
-            <div className="h-4 w-4">
-              <PinIcon unpin={isPinned} />
-            </div>
-          </button>
-        </Ariakit.MenuItem>
-        <Ariakit.Menu
-          portal={true}
-          unmountOnHide={true}
-          className={cn(
-            'animate-popover-left z-50 ml-3 flex min-w-[200px] flex-col rounded-xl',
-            'border border-border-light bg-surface-secondary p-1 shadow-lg',
-          )}
-        >
-          {configuredServers.map((serverName) => {
-            const statusIconProps = getServerStatusIconProps(serverName);
-            const isSelected = mcpValues?.includes(serverName) ?? false;
-
-            const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />;
-
-            return (
-              <Ariakit.MenuItem
-                key={serverName}
-                onClick={(event) => {
-                  event.preventDefault();
-                  toggleServerSelection(serverName);
+    return (
+      <div ref={ref}>
+        <Ariakit.MenuProvider store={menuStore}>
+          <Ariakit.MenuItem
+            {...props}
+            render={
+              <Ariakit.MenuButton
+                onClick={(e: React.MouseEvent<HTMLButtonElement>) => {
+                  e.stopPropagation();
+                  menuStore.toggle();
                 }}
-                className={cn(
-                  'flex items-center gap-2 rounded-lg px-2 py-1.5 text-text-primary hover:cursor-pointer',
-                  'scroll-m-1 outline-none transition-colors',
-                  'hover:bg-black/[0.075] dark:hover:bg-white/10',
-                  'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
-                  'w-full min-w-0 justify-between text-sm',
-                )}
-              >
-                <button
-                  type="button"
-                  className="flex flex-grow items-center gap-2 rounded bg-transparent p-0 text-left transition-colors focus:outline-none"
-                  tabIndex={0}
+                className="flex w-full cursor-pointer items-center justify-between rounded-lg p-2 hover:bg-surface-hover"
+              />
+            }
+          >
+            <div className="flex items-center gap-2">
+              <MCPIcon className="icon-md" />
+              <span>{placeholder || placeholderText}</span>
+              <ChevronRight className="ml-auto h-3 w-3" />
+            </div>
+            <button
+              type="button"
+              onClick={(e) => {
+                e.stopPropagation();
+                setIsPinned(!isPinned);
+              }}
+              className={cn(
+                'rounded p-1 transition-all duration-200',
+                'hover:bg-surface-tertiary hover:shadow-sm',
+                !isPinned && 'text-text-secondary hover:text-text-primary',
+              )}
+              aria-label={isPinned ? 'Unpin' : 'Pin'}
+            >
+              <div className="h-4 w-4">
+                <PinIcon unpin={isPinned} />
+              </div>
+            </button>
+          </Ariakit.MenuItem>
+          <Ariakit.Menu
+            portal={true}
+            unmountOnHide={true}
+            className={cn(
+              'animate-popover-left z-50 ml-3 flex min-w-[200px] flex-col rounded-xl',
+              'border border-border-light bg-surface-secondary p-1 shadow-lg',
+            )}
+          >
+            {configuredServers.map((serverName) => {
+              const statusIconProps = getServerStatusIconProps(serverName);
+              const isSelected = mcpValues?.includes(serverName) ?? false;
+
+              const statusIcon = statusIconProps && <MCPServerStatusIcon {...statusIconProps} />;
+
+              return (
+                <Ariakit.MenuItem
+                  key={serverName}
+                  onClick={(event) => {
+                    event.preventDefault();
+                    toggleServerSelection(serverName);
+                  }}
+                  className={cn(
+                    'flex items-center gap-2 rounded-lg px-2 py-1.5 text-text-primary hover:cursor-pointer',
+                    'scroll-m-1 outline-none transition-colors',
+                    'hover:bg-black/[0.075] dark:hover:bg-white/10',
+                    'data-[active-item]:bg-black/[0.075] dark:data-[active-item]:bg-white/10',
+                    'w-full min-w-0 justify-between text-sm',
+                  )}
                 >
-                  <Ariakit.MenuItemCheck checked={isSelected} />
-                  <span>{serverName}</span>
-                </button>
-                {statusIcon && <div className="ml-2 flex items-center">{statusIcon}</div>}
-              </Ariakit.MenuItem>
-            );
-          })}
-        </Ariakit.Menu>
-      </Ariakit.MenuProvider>
-      {configDialogProps && <MCPConfigDialog {...configDialogProps} />}
-    </>
-  );
-};
+                  <div className="flex flex-grow items-center gap-2">
+                    <Ariakit.MenuItemCheck checked={isSelected} />
+                    <span>{serverName}</span>
+                  </div>
+                  {statusIcon && <div className="ml-2 flex items-center">{statusIcon}</div>}
+                </Ariakit.MenuItem>
+              );
+            })}
+          </Ariakit.Menu>
+        </Ariakit.MenuProvider>
+        {configDialogProps && <MCPConfigDialog {...configDialogProps} />}
+      </div>
+    );
+  },
+);
+
+MCPSubMenu.displayName = 'MCPSubMenu';
 
 export default React.memo(MCPSubMenu);