@marimo-team/frontend 0.24.1-dev41 → 0.24.1-dev43

package/src/components/editor/actions/useNotebookActions.tsx

@@ -73,6 +73,7 @@ import { useRequestClient } from "@/core/network/requests";
 import { useFilename } from "@/core/saving/filename";
 import { createShareableLink } from "@/core/wasm/share";
 import { isWasm } from "@/core/wasm/utils";
+import { useDetectedDataSources } from "@/hooks/useDataSourceDiscovery";
 import { copyToClipboard } from "@/utils/copy";
 import { Objects } from "@/utils/objects";
 import { Strings } from "@/utils/strings";
@@ -81,6 +82,7 @@ import { useRunAllCells } from "../cell/useRunCells";
 import { useChromeActions, useChromeState } from "../chrome/state";
 import { isPanelHidden, PANELS } from "../chrome/types";
 import { AddConnectionDialogContent } from "../connections/add-connection-dialog";
+import { useAddDetectedDataSource } from "../connections/components";
 import { keyboardShortcutsAtom } from "../controls/keyboard-shortcuts";
 import { commandPaletteAtom } from "../controls/state";
 import { displayLayoutName, getLayoutIcon } from "../renderers/layout-select";
@@ -135,6 +137,9 @@ export function useNotebookActions({
   const setKeyboardShortcutsOpen = useSetAtom(keyboardShortcutsAtom);
   const setExportOptions = useSetAtom(exportOptionsAtom);
   const { readCode, saveCellConfig } = useRequestClient();
+  const detectedDatabaseSources = useDetectedDataSources("database");
+  const detectedStorageSources = useDetectedDataSources("storage");
+  const addDetectedDataSource = useAddDetectedDataSource();
 
   const hasDisabledCells = useAtomValue(hasDisabledCellsAtom);
   const canUndoDeletes = useAtomValue(canUndoDeletesAtom);
@@ -455,6 +460,27 @@ export function useNotebookActions({
       handle: () => {
         openModal(<AddConnectionDialogContent onClose={closeModal} />);
       },
+      dropdown:
+        detectedDatabaseSources.length === 0
+          ? undefined
+          : [
+              ...detectedDatabaseSources.map((source) => ({
+                icon: <SparklesIcon size={14} strokeWidth={1.5} />,
+                label: `Add ${source.displayName}`,
+                description: "Detected in your environment",
+                handle: () => addDetectedDataSource(source),
+              })),
+              {
+                divider: true,
+                icon: <DatabaseIcon size={14} strokeWidth={1.5} />,
+                label: "Browse all connections",
+                handle: () => {
+                  openModal(
+                    <AddConnectionDialogContent onClose={closeModal} />,
+                  );
+                },
+              },
+            ],
     },
     {
       icon: <HardDrive size={14} strokeWidth={1.5} />,
@@ -467,6 +493,30 @@ export function useNotebookActions({
           />,
         );
       },
+      dropdown:
+        detectedStorageSources.length === 0
+          ? undefined
+          : [
+              ...detectedStorageSources.map((source) => ({
+                icon: <SparklesIcon size={14} strokeWidth={1.5} />,
+                label: `Add ${source.displayName}`,
+                description: "Detected in your environment",
+                handle: () => addDetectedDataSource(source),
+              })),
+              {
+                divider: true,
+                icon: <HardDrive size={14} strokeWidth={1.5} />,
+                label: "Browse all connections",
+                handle: () => {
+                  openModal(
+                    <AddConnectionDialogContent
+                      defaultTab="storage"
+                      onClose={closeModal}
+                    />,
+                  );
+                },
+              },
+            ],
     },
     {
       icon: <Undo2Icon size={14} strokeWidth={1.5} />,

package/src/components/editor/chrome/panels/components.tsx

@@ -105,7 +105,7 @@ const DiscoveredSourcesBadge = ({
   return (
     <Tooltip content={content}>
       <span>
-        <PanelBadge className="gap-0.5 border-(--blue-6) bg-(--blue-3) text-(--blue-11)">
+        <PanelBadge className="gap-0.5 border-(--blue-6) bg-(--blue-3) text-(--blue-11) hover:bg-(--blue-4) dark:border-(--blue-8) dark:bg-(--blue-3) dark:hover:bg-(--blue-4)">
           <SparklesIcon className="h-2.5 w-2.5" />
           {count}
         </PanelBadge>

package/src/components/editor/connections/__tests__/add-connection-button.test.tsx

@@ -0,0 +1,129 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { fireEvent, render, screen } from "@testing-library/react";
+import { Provider } from "jotai";
+import type React from "react";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { ModalProvider } from "@/components/modal/ImperativeModal";
+import { TooltipProvider } from "@/components/ui/tooltip";
+import type { DetectedDataSource } from "@/core/datasets/data-source-discovery";
+import { store } from "@/core/state/jotai";
+import { AddConnectionButton } from "../add-connection-button";
+
+const mocks = vi.hoisted(() => ({
+  addDetectedDataSource: vi.fn(),
+  sources: [] as DetectedDataSource[],
+}));
+
+const MYSQL_SOURCE: DetectedDataSource = {
+  id: "mysql-environment",
+  integration: "mysql",
+  category: "database",
+  displayName: "MySQL",
+  confidence: "high",
+  origins: [{ type: "environment", label: "Kernel environment" }],
+  configuration: [],
+  code: "engine = create_engine()",
+  hidesWhen: { kind: "dialect", substrings: ["mysql"] },
+};
+
+vi.mock("@/hooks/useDataSourceDiscovery", () => ({
+  useDetectedDataSources: () => mocks.sources,
+}));
+
+vi.mock("../components", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("../components")>()),
+  useAddDetectedDataSource: () => mocks.addDetectedDataSource,
+}));
+
+function wrapper({ children }: { children: React.ReactNode }) {
+  return (
+    <Provider store={store}>
+      <TooltipProvider>
+        <ModalProvider>{children}</ModalProvider>
+      </TooltipProvider>
+    </Provider>
+  );
+}
+
+describe("AddConnectionButton", () => {
+  beforeEach(() => {
+    mocks.sources = [];
+    mocks.addDetectedDataSource.mockClear();
+    vi.stubGlobal("PointerEvent", MouseEvent);
+  });
+
+  afterEach(() => vi.unstubAllGlobals());
+
+  it("opens the connection dialog directly without suggestions", () => {
+    render(
+      <AddConnectionButton
+        group="database"
+        label="Add database or catalog"
+        variant="outline"
+      />,
+      { wrapper },
+    );
+
+    fireEvent.click(
+      screen.getByRole("button", { name: "Add database or catalog" }),
+    );
+
+    expect(
+      screen.getByRole("dialog", { name: "Add Connection" }),
+    ).toBeVisible();
+  });
+
+  it("opens an ordered suggestion menu when sources are detected", async () => {
+    mocks.sources = [MYSQL_SOURCE];
+    render(
+      <AddConnectionButton
+        group="database"
+        label="Add database or catalog"
+        variant="outline"
+      />,
+      { wrapper },
+    );
+
+    fireEvent.pointerDown(
+      screen.getByRole("button", {
+        name: "View detected database connections",
+      }),
+      { button: 0, ctrlKey: false },
+    );
+
+    const suggestion = await screen.findByRole("menuitem", {
+      name: "Add MySQL",
+    });
+    const browseAll = screen.getByRole("menuitem", {
+      name: "Browse all connections",
+    });
+    expect(
+      suggestion.compareDocumentPosition(browseAll) &
+        Node.DOCUMENT_POSITION_FOLLOWING,
+    ).toBeTruthy();
+
+    fireEvent.click(suggestion);
+    expect(mocks.addDetectedDataSource).toHaveBeenCalledWith(MYSQL_SOURCE);
+  });
+
+  it("keeps the primary action direct when sources are detected", () => {
+    mocks.sources = [MYSQL_SOURCE];
+    render(
+      <AddConnectionButton
+        group="database"
+        label="Add database or catalog"
+        variant="outline"
+      />,
+      { wrapper },
+    );
+
+    fireEvent.click(
+      screen.getByRole("button", { name: "Add database or catalog" }),
+    );
+
+    expect(
+      screen.getByRole("dialog", { name: "Add Connection" }),
+    ).toBeVisible();
+  });
+});

package/src/components/editor/connections/add-connection-button.tsx

@@ -0,0 +1,132 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import {
+  DatabaseIcon,
+  HardDriveIcon,
+  PlusIcon,
+  SparklesIcon,
+} from "lucide-react";
+import type React from "react";
+import { useImperativeModal } from "@/components/modal/ImperativeModal";
+import { Button } from "@/components/ui/button";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuSeparator,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
+import { Tooltip } from "@/components/ui/tooltip";
+import type { DataSourceDiscoveryGroup } from "@/core/datasets/data-source-discovery";
+import { useDetectedDataSources } from "@/hooks/useDataSourceDiscovery";
+import { cn } from "@/utils/cn";
+import {
+  AddConnectionDialog,
+  AddConnectionDialogContent,
+} from "./add-connection-dialog";
+import { useAddDetectedDataSource } from "./components";
+
+type AddConnectionButtonProps = Omit<
+  React.ComponentProps<typeof Button>,
+  "children"
+> & {
+  group: DataSourceDiscoveryGroup;
+  label: string;
+  compact?: boolean;
+};
+
+/**
+ * Opens the connection dialog directly and, when useful suggestions exist,
+ * adds a secondary menu for detected sources.
+ */
+export const AddConnectionButton: React.FC<AddConnectionButtonProps> = ({
+  group,
+  label,
+  compact = false,
+  className,
+  ...buttonProps
+}) => {
+  const sources = useDetectedDataSources(group);
+  const addDetectedDataSource = useAddDetectedDataSource();
+  const { openModal, closeModal } = useImperativeModal();
+  const hasSuggestions = sources.length > 0;
+  const defaultTab = group === "storage" ? "storage" : "databases";
+
+  const primaryButton = (
+    <Button
+      {...buttonProps}
+      className={cn(className, hasSuggestions && !compact && "rounded-r-none")}
+      aria-label={compact ? label : buttonProps["aria-label"]}
+    >
+      {!compact && label}
+      <PlusIcon className={cn("h-4 w-4", !compact && "ml-2")} />
+    </Button>
+  );
+
+  const primaryAction = (
+    <AddConnectionDialog defaultTab={defaultTab}>
+      {primaryButton}
+    </AddConnectionDialog>
+  );
+
+  if (!hasSuggestions) {
+    return primaryAction;
+  }
+
+  const BrowseIcon = group === "storage" ? HardDriveIcon : DatabaseIcon;
+  const suggestionsLabel =
+    group === "storage"
+      ? "View detected remote storage connections"
+      : "View detected database connections";
+
+  return (
+    <div className="inline-flex">
+      {primaryAction}
+      <DropdownMenu>
+        <Tooltip content={suggestionsLabel}>
+          <span className="inline-flex">
+            <DropdownMenuTrigger asChild={true}>
+              <Button
+                variant={buttonProps.variant}
+                size={buttonProps.size}
+                disabled={buttonProps.disabled}
+                className={cn(
+                  "rounded-l-none border-l-0 px-2 text-(--blue-10) hover:bg-(--blue-3)",
+                  compact && "rounded-none",
+                )}
+                aria-label={suggestionsLabel}
+              >
+                <SparklesIcon className="h-4 w-4" />
+              </Button>
+            </DropdownMenuTrigger>
+          </span>
+        </Tooltip>
+        <DropdownMenuContent align="end">
+          {sources.map((source) => (
+            <DropdownMenuItem
+              key={source.id}
+              onSelect={() => addDetectedDataSource(source)}
+            >
+              <SparklesIcon className="mr-2 h-4 w-4 text-(--blue-9)" />
+              Add {source.displayName}
+            </DropdownMenuItem>
+          ))}
+          <DropdownMenuSeparator />
+          <DropdownMenuItem
+            onSelect={() =>
+              openModal(
+                <AddConnectionDialogContent
+                  defaultTab={defaultTab}
+                  onClose={closeModal}
+                />,
+              )
+            }
+          >
+            <BrowseIcon className="mr-2 h-4 w-4" />
+            Browse all connections
+          </DropdownMenuItem>
+        </DropdownMenuContent>
+      </DropdownMenu>
+    </div>
+  );
+};

package/src/components/editor/connections/components.tsx

@@ -21,6 +21,7 @@ import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
 import { useCellActions } from "@/core/cells/cells";
 import { useLastFocusedCellId } from "@/core/cells/focus";
 import { autoInstantiateAtom } from "@/core/config/config";
+import type { DetectedDataSource } from "@/core/datasets/data-source-discovery";
 import {
   ENV_RENDERER,
   SECRET_TEXTAREA_RENDERER,
@@ -134,6 +135,16 @@ export function useInsertCode() {
   };
 }
 
+/** Insert the configured cell supplied by a detected data source. */
+export function useAddDetectedDataSource(onAdd?: () => void) {
+  const insertCode = useInsertCode();
+
+  return (source: DetectedDataSource) => {
+    insertCode(source.code);
+    onAdd?.();
+  };
+}
+
 /**
  * Generic connection form: Zod-driven form with secrets support, a library
  * picker, and Back/Add buttons. Used by both database and storage forms.

package/src/components/editor/connections/quick-add-data-sources.tsx

@@ -1,13 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { PlusIcon, SparklesIcon } from "lucide-react";
+import { useId } from "react";
 import { Tooltip, TooltipProvider } from "@/components/ui/tooltip";
 import type {
   DataSourceDiscoveryGroup,
   DetectedDataSource,
 } from "@/core/datasets/data-source-discovery";
 import { useDetectedDataSources } from "@/hooks/useDataSourceDiscovery";
-import { useInsertCode } from "./components";
+import { useAddDetectedDataSource } from "./components";
 import { cn } from "@/utils/cn";
 
 export const QuickAddDataSources: React.FC<{
@@ -15,22 +16,24 @@ export const QuickAddDataSources: React.FC<{
   sources: DetectedDataSource[];
   onAdd: (source: DetectedDataSource) => void;
 }> = ({ className, sources, onAdd }) => {
+  const titleId = useId();
+
   if (sources.length === 0) {
     return null;
   }
 
   return (
     <section
-      aria-labelledby="quick-add-data-sources-title"
+      aria-labelledby={titleId}
       className={cn(
-        "rounded-full bg-[linear-gradient(135deg,var(--blue-2),var(--purple-3))] px-3 py-2",
+        "rounded-full bg-[linear-gradient(135deg,var(--blue-2),var(--purple-3))] px-3 py-2 dark:bg-[linear-gradient(135deg,var(--blue-2),var(--purple-2))]",
         className,
       )}
     >
       <div className="flex flex-wrap items-center gap-2">
         <div className="mr-1 flex items-center gap-1.5">
-          <SparklesIcon className="h-3.5 w-3.5 text-(--blue-9)" />
-          <h3 id="quick-add-data-sources-title" className="text-sm">
+          <SparklesIcon className="h-3.5 w-3.5 text-(--blue-9) dark:text-(--blue-10)" />
+          <h3 id={titleId} className="text-sm">
             Quick add
           </h3>
         </div>
@@ -44,7 +47,7 @@ export const QuickAddDataSources: React.FC<{
               <button
                 type="button"
                 aria-label={`Add ${source.displayName} connection`}
-                className="inline-flex items-center gap-1 rounded-full border border-border/60 bg-background px-2.5 py-1 text-xs font-medium text-foreground/90 transition-colors hover:border-(--blue-7) focus:outline-hidden focus:ring-2 focus:ring-ring focus:ring-offset-2"
+                className="inline-flex items-center gap-1 rounded-full border border-border/60 bg-background px-2.5 py-1 text-xs font-medium text-foreground/90 transition-colors hover:border-(--blue-7) focus:outline-hidden focus:ring-2 focus:ring-ring focus:ring-offset-2 dark:border-border dark:bg-background/60 dark:hover:bg-accent/80"
                 onClick={() => onAdd(source)}
               >
                 <PlusIcon className="h-3 w-3 text-muted-foreground" />
@@ -101,17 +104,14 @@ export const AutoDiscoveredDataSources: React.FC<{
   className?: string;
   group?: DataSourceDiscoveryGroup;
 }> = ({ onSubmit, className, group }) => {
-  const insertCode = useInsertCode();
+  const addDetectedDataSource = useAddDetectedDataSource(onSubmit);
   const sources = useDetectedDataSources(group);
 
   return (
     <QuickAddDataSources
       className={className}
       sources={sources}
-      onAdd={(source) => {
-        insertCode(source.code);
-        onSubmit?.();
-      }}
+      onAdd={addDetectedDataSource}
     />
   );
 };

package/src/components/editor/controls/__tests__/notebook-menu-dropdown.test.tsx

@@ -1,18 +1,31 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import {
+  act,
+  fireEvent,
+  render,
+  screen,
+  waitFor,
+} from "@testing-library/react";
 import { Provider } from "jotai";
 import type React from "react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { MockRequestClient } from "@/__mocks__/requests";
 import { ModalProvider } from "@/components/modal/ImperativeModal";
 import { TooltipProvider } from "@/components/ui/tooltip";
+import {
+  type DetectedDataSource,
+  invalidateDataSourceDiscovery,
+} from "@/core/datasets/data-source-discovery";
+import { DiscoverDataSources } from "@/core/datasets/request-registry";
 import { layoutStateAtom } from "@/core/layout/layout";
 import { kioskModeAtom, viewStateAtom } from "@/core/mode";
+import { connectionAtom } from "@/core/network/connection";
 import { requestClientAtom } from "@/core/network/requests";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
 import { isWasm } from "@/core/wasm/utils";
+import { WebSocketState } from "@/core/websocket/types";
 import {
   DEFAULT_EXPORT_OPTIONS,
   exportOptionsAtom,
@@ -20,6 +33,34 @@ import {
 } from "../../actions/export-dialog/state";
 import { NotebookMenuDropdown } from "../notebook-menu-dropdown";
 
+const POSTGRES_SOURCE: DetectedDataSource = {
+  id: "postgres-libpq-environment",
+  integration: "postgres",
+  category: "database",
+  displayName: "PostgreSQL",
+  confidence: "high",
+  origins: [{ type: "environment", label: "Kernel environment" }],
+  configuration: [],
+  code: "engine = create_engine()",
+  hidesWhen: { kind: "dialect", substrings: ["postgres"] },
+};
+
+const S3_SOURCE: DetectedDataSource = {
+  id: "aws-s3-environment",
+  integration: "aws",
+  category: "object-storage",
+  displayName: "AWS S3",
+  confidence: "high",
+  origins: [{ type: "environment", label: "Kernel environment" }],
+  configuration: [],
+  code: "storage = create_s3_storage()",
+  hidesWhen: {
+    kind: "storage",
+    protocols: ["s3"],
+    backendTypes: ["s3"],
+  },
+};
+
 vi.mock("@/core/wasm/utils", async (importOriginal) => {
   const actual = await importOriginal<typeof import("@/core/wasm/utils")>();
   return { ...actual, isWasm: vi.fn(() => false) };
@@ -59,6 +100,7 @@ describe("NotebookMenuDropdown", () => {
     store.set(requestClientAtom, MockRequestClient.create());
     store.set(filenameAtom, "/project/notebook.py");
     store.set(viewStateAtom, { mode: "edit", cellAnchor: null });
+    store.set(connectionAtom, { state: WebSocketState.OPEN });
     store.set(kioskModeAtom, false);
     store.set(layoutStateAtom, {
       selectedLayout: "vertical",
@@ -73,11 +115,77 @@ describe("NotebookMenuDropdown", () => {
       addEventListener: vi.fn(),
       removeEventListener: vi.fn(),
     }));
+    invalidateDataSourceDiscovery();
+    vi.spyOn(DiscoverDataSources, "request").mockResolvedValue({
+      request_id: "request-id",
+      sources: [],
+    });
   });
 
   afterEach(() => {
     document.title = "";
     vi.unstubAllGlobals();
+    vi.restoreAllMocks();
+  });
+
+  it("shows detected databases in the add connection submenu", async () => {
+    vi.mocked(DiscoverDataSources.request).mockResolvedValue({
+      request_id: "request-id",
+      sources: [POSTGRES_SOURCE],
+    });
+
+    render(<NotebookMenuDropdown />, { wrapper });
+    await waitFor(() =>
+      expect(DiscoverDataSources.request).toHaveBeenCalledOnce(),
+    );
+    await act(async () => Promise.resolve());
+    openNotebookMenu();
+    fireEvent.click(
+      await screen.findByRole("menuitem", { name: "Add database connection" }),
+    );
+
+    const suggestion = await screen.findByRole("menuitem", {
+      name: "Add PostgreSQL",
+    });
+    const browseAll = screen.getByRole("menuitem", {
+      name: "Browse all connections",
+    });
+    expect(suggestion).toBeVisible();
+    expect(browseAll).toBeVisible();
+    expect(
+      suggestion.compareDocumentPosition(browseAll) &
+        Node.DOCUMENT_POSITION_FOLLOWING,
+    ).toBeTruthy();
+  });
+
+  it("shows detected storage in the add remote storage submenu", async () => {
+    vi.mocked(DiscoverDataSources.request).mockResolvedValue({
+      request_id: "request-id",
+      sources: [S3_SOURCE],
+    });
+
+    render(<NotebookMenuDropdown />, { wrapper });
+    await waitFor(() =>
+      expect(DiscoverDataSources.request).toHaveBeenCalledOnce(),
+    );
+    await act(async () => Promise.resolve());
+    openNotebookMenu();
+    fireEvent.click(
+      await screen.findByRole("menuitem", { name: "Add remote storage" }),
+    );
+
+    const suggestion = await screen.findByRole("menuitem", {
+      name: "Add AWS S3",
+    });
+    const browseAll = screen.getByRole("menuitem", {
+      name: "Browse all connections",
+    });
+    expect(suggestion).toBeVisible();
+    expect(browseAll).toBeVisible();
+    expect(
+      suggestion.compareDocumentPosition(browseAll) &
+        Node.DOCUMENT_POSITION_FOLLOWING,
+    ).toBeTruthy();
   });
 
   it("returns focus to the notebook menu when the dialog closes", async () => {

package/src/components/storage/storage-inspector.tsx

@@ -8,7 +8,6 @@ import {
   HardDriveIcon,
   HelpCircleIcon,
   LoaderCircle,
-  PlusIcon,
   ViewIcon,
   XIcon,
 } from "lucide-react";
@@ -17,7 +16,7 @@ import { useLocale } from "react-aria";
 import { EngineVariable } from "@/components/databases/engine-variable";
 import { useAddCodeToNewCell } from "@/components/editor/cell/useAddCell";
 import { PanelEmptyState } from "@/components/editor/chrome/panels/empty-state";
-import { AddConnectionDialog } from "@/components/editor/connections/add-connection-dialog";
+import { AddConnectionButton } from "@/components/editor/connections/add-connection-button";
 import {
   FILE_ICON_COLOR,
   renderFileIcon,
@@ -1011,12 +1010,12 @@ export const StorageInspector: React.FC = () => {
           </span>
         }
         action={
-          <AddConnectionDialog defaultTab="storage">
-            <Button variant="outline" size="sm">
-              Add remote storage
-              <PlusIcon className="h-4 w-4 ml-2" />
-            </Button>
-          </AddConnectionDialog>
+          <AddConnectionButton
+            group="storage"
+            label="Add remote storage"
+            variant="outline"
+            size="sm"
+          />
         }
         icon={<HardDriveIcon className="h-8 w-8" />}
       />
@@ -1076,15 +1075,14 @@ export const StorageInspector: React.FC = () => {
           >
             <HelpCircleIcon className="h-3.5 w-3.5 shrink-0 cursor-help text-muted-foreground hover:text-foreground mr-2" />
           </Tooltip>
-          <AddConnectionDialog defaultTab="storage">
-            <Button
-              variant="ghost"
-              size="sm"
-              className="px-2 border-0 border-l border-muted-background rounded-none focus-visible:ring-0 focus-visible:ring-offset-0"
-            >
-              <PlusIcon className="h-4 w-4" />
-            </Button>
-          </AddConnectionDialog>
+          <AddConnectionButton
+            group="storage"
+            label="Add remote storage"
+            compact={true}
+            variant="ghost"
+            size="sm"
+            className="px-2 border-0 border-l border-muted-background rounded-none focus-visible:ring-0 focus-visible:ring-offset-0"
+          />
         </div>
         <CommandList className="flex flex-col">
           {namespaces.map((ns) => {