@marimo-team/islands 0.23.9-dev34 → 0.23.9-dev36

package/src/components/data-table/__tests__/header-items.test.tsx

@@ -1,6 +1,11 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { Column, SortingState } from "@tanstack/react-table";
+import type {
+  Column,
+  SortDirection,
+  SortingState,
+  Table,
+} from "@tanstack/react-table";
 import { fireEvent, render, screen } from "@testing-library/react";
 import { describe, expect, it, vi } from "vitest";
 import {
@@ -8,7 +13,23 @@ import {
   DropdownMenuContent,
   DropdownMenuTrigger,
 } from "@/components/ui/dropdown-menu";
-import { HideColumn } from "../header-items";
+import {
+  ColumnPinning,
+  ColumnWrapping,
+  CopyColumn,
+  DataType,
+  FormatOptions,
+  HideColumn,
+  Sorts,
+} from "../header-items";
+
+const renderInMenu = (node: React.ReactNode) =>
+  render(
+    <DropdownMenu open={true}>
+      <DropdownMenuTrigger />
+      <DropdownMenuContent>{node}</DropdownMenuContent>
+    </DropdownMenu>,
+  );
 
 describe("multi-column sorting logic", () => {
   // Extract the core sorting logic to test in isolation
@@ -167,14 +188,6 @@ describe("HideColumn", () => {
       toggleVisibility,
     }) as unknown as Column<unknown, unknown>;
 
-  const renderInMenu = (node: React.ReactNode) =>
-    render(
-      <DropdownMenu open={true}>
-        <DropdownMenuTrigger />
-        <DropdownMenuContent>{node}</DropdownMenuContent>
-      </DropdownMenu>,
-    );
-
   it("renders 'Hide column' when canHide is true", () => {
     renderInMenu(<HideColumn column={makeColumn()} />);
     expect(screen.getByText("Hide column")).toBeInTheDocument();
@@ -192,3 +205,200 @@ describe("HideColumn", () => {
     expect(toggleVisibility).toHaveBeenCalledWith(false);
   });
 });
+
+describe("DataType", () => {
+  const makeColumn = (dtype?: string) =>
+    ({
+      columnDef: { meta: dtype === undefined ? {} : { dtype } },
+    }) as unknown as Column<unknown, unknown>;
+
+  it("renders the dtype label when present", () => {
+    renderInMenu(<DataType column={makeColumn("int64")} />);
+    expect(screen.getByText("int64")).toBeInTheDocument();
+  });
+
+  it("returns null when dtype is absent", () => {
+    renderInMenu(<DataType column={makeColumn()} />);
+    expect(screen.queryByText("int64")).toBeNull();
+  });
+});
+
+describe("Sorts", () => {
+  const makeColumn = ({
+    canSort = true,
+    sorted = false,
+    sortIndex = 0,
+  }: {
+    canSort?: boolean;
+    sorted?: false | SortDirection;
+    sortIndex?: number;
+  } = {}) =>
+    ({
+      getCanSort: () => canSort,
+      getIsSorted: () => sorted,
+      getSortIndex: () => sortIndex,
+      clearSorting: vi.fn(),
+      toggleSorting: vi.fn(),
+    }) as unknown as Column<unknown, unknown>;
+
+  const makeTable = (sorting: SortingState) =>
+    ({
+      getState: () => ({ sorting }),
+      resetSorting: vi.fn(),
+    }) as unknown as Table<unknown>;
+
+  it("returns null when the column cannot sort", () => {
+    renderInMenu(<Sorts column={makeColumn({ canSort: false })} />);
+    expect(screen.queryByText("Asc")).toBeNull();
+  });
+
+  it("renders Asc and Desc items", () => {
+    renderInMenu(<Sorts column={makeColumn()} />);
+    expect(screen.getByText("Asc")).toBeInTheDocument();
+    expect(screen.getByText("Desc")).toBeInTheDocument();
+  });
+
+  it("offers single-column 'Clear sort' when sorted without multi-sort", () => {
+    renderInMenu(<Sorts column={makeColumn({ sorted: "asc" })} />);
+    expect(screen.getByText("Clear sort")).toBeInTheDocument();
+  });
+
+  it("offers 'Clear all sorts' when the table has multiple sorts", () => {
+    renderInMenu(
+      <Sorts
+        column={makeColumn({ sorted: "asc" })}
+        table={makeTable([
+          { id: "a", desc: false },
+          { id: "b", desc: true },
+        ])}
+      />,
+    );
+    expect(screen.getByText("Clear all sorts")).toBeInTheDocument();
+  });
+});
+
+describe("CopyColumn", () => {
+  const makeColumn = ({
+    canCopy = true,
+    id = "name",
+  }: {
+    canCopy?: boolean;
+    id?: string;
+  } = {}) =>
+    ({
+      id,
+      getCanCopy: () => canCopy,
+    }) as unknown as Column<unknown, unknown>;
+
+  it("renders 'Copy column name' when copyable", () => {
+    renderInMenu(<CopyColumn column={makeColumn()} />);
+    expect(screen.getByText("Copy column name")).toBeInTheDocument();
+  });
+
+  it("returns null when the column cannot be copied", () => {
+    renderInMenu(<CopyColumn column={makeColumn({ canCopy: false })} />);
+    expect(screen.queryByText("Copy column name")).toBeNull();
+  });
+});
+
+describe("ColumnPinning", () => {
+  const makeColumn = ({
+    canPin = true,
+    pinned = false,
+  }: {
+    canPin?: boolean;
+    pinned?: false | "left" | "right";
+  } = {}) =>
+    ({
+      getCanPin: () => canPin,
+      getIsPinned: () => pinned,
+      pin: vi.fn(),
+    }) as unknown as Column<unknown, unknown>;
+
+  it("returns null when the column cannot be pinned", () => {
+    renderInMenu(<ColumnPinning column={makeColumn({ canPin: false })} />);
+    expect(screen.queryByText("Freeze left")).toBeNull();
+  });
+
+  it("offers freeze options when unpinned", () => {
+    renderInMenu(<ColumnPinning column={makeColumn()} />);
+    expect(screen.getByText("Freeze left")).toBeInTheDocument();
+    expect(screen.getByText("Freeze right")).toBeInTheDocument();
+  });
+
+  it("offers 'Unfreeze' when pinned", () => {
+    renderInMenu(<ColumnPinning column={makeColumn({ pinned: "left" })} />);
+    expect(screen.getByText("Unfreeze")).toBeInTheDocument();
+  });
+});
+
+describe("ColumnWrapping", () => {
+  const makeColumn = ({
+    canWrap = true,
+    wrapping = "nowrap",
+  }: {
+    canWrap?: boolean;
+    wrapping?: "wrap" | "nowrap";
+  } = {}) =>
+    ({
+      getCanWrap: () => canWrap,
+      getColumnWrapping: () => wrapping,
+      toggleColumnWrapping: vi.fn(),
+    }) as unknown as Column<unknown, unknown>;
+
+  it("returns null when the column cannot wrap", () => {
+    renderInMenu(<ColumnWrapping column={makeColumn({ canWrap: false })} />);
+    expect(screen.queryByText("Wrap text")).toBeNull();
+  });
+
+  it("offers 'Wrap text' when not wrapping", () => {
+    renderInMenu(<ColumnWrapping column={makeColumn()} />);
+    expect(screen.getByText("Wrap text")).toBeInTheDocument();
+  });
+
+  it("offers 'No wrap text' when wrapping", () => {
+    renderInMenu(<ColumnWrapping column={makeColumn({ wrapping: "wrap" })} />);
+    expect(screen.getByText("No wrap text")).toBeInTheDocument();
+  });
+});
+
+describe("FormatOptions", () => {
+  const makeColumn = ({
+    dataType = "number",
+    canFormat = true,
+  }: {
+    dataType?: string;
+    canFormat?: boolean;
+  } = {}) =>
+    ({
+      columnDef: { meta: { dataType } },
+      getCanFormat: () => canFormat,
+      getColumnFormatting: () => undefined,
+      setColumnFormatting: vi.fn(),
+    }) as unknown as Column<unknown, unknown>;
+
+  it("renders the 'Format' submenu trigger for formattable columns", () => {
+    renderInMenu(<FormatOptions column={makeColumn()} locale="en-US" />);
+    expect(screen.getByText("Format")).toBeInTheDocument();
+  });
+
+  it("returns null when the column cannot be formatted", () => {
+    renderInMenu(
+      <FormatOptions
+        column={makeColumn({ canFormat: false })}
+        locale="en-US"
+      />,
+    );
+    expect(screen.queryByText("Format")).toBeNull();
+  });
+
+  it("returns null when the data type has no format options", () => {
+    renderInMenu(
+      <FormatOptions
+        column={makeColumn({ dataType: "unknown" })}
+        locale="en-US"
+      />,
+    );
+    expect(screen.queryByText("Format")).toBeNull();
+  });
+});

package/src/components/data-table/column-header.tsx

@@ -17,14 +17,14 @@ import { useFilterEditor } from "./filter-editor-context";
 import { EDITABLE_FILTER_TYPES, isMembershipFilterType } from "./filters";
 import {
   ClearFilterMenuItem,
+  ColumnPinning,
+  ColumnWrapping,
+  CopyColumn,
+  DataType,
+  FormatOptions,
   HideColumn,
-  renderColumnPinning,
-  renderColumnWrapping,
-  renderCopyColumn,
-  renderDataType,
-  renderFormatOptions,
   renderSortIcon,
-  renderSorts,
+  Sorts,
 } from "./header-items";
 
 interface DataTableColumnHeaderProps<
@@ -36,6 +36,11 @@ interface DataTableColumnHeaderProps<
   subheader?: React.ReactNode;
   justify?: "left" | "center" | "right";
   calculateTopKRows?: CalculateTopKRows;
+  /**
+   * Optional: only used to surface multi-column sort actions ("Clear all
+   * sorts"). Omitted by call sites that define their header inside column
+   * definitions, where the table instance isn't yet available.
+   */
   table?: Table<TData>;
 }
 
@@ -119,12 +124,12 @@ export const DataTableColumnHeader = <TData, TValue>({
             </button>
           </DropdownMenuTrigger>
           <DropdownMenuContent align="start">
-            {renderDataType(column)}
-            {renderSorts(column, table)}
-            {renderCopyColumn(column)}
-            {renderColumnPinning(column)}
-            {renderColumnWrapping(column)}
-            {renderFormatOptions(column, locale)}
+            <DataType column={column} />
+            <Sorts column={column} table={table} />
+            <CopyColumn column={column} />
+            <ColumnPinning column={column} />
+            <ColumnWrapping column={column} />
+            <FormatOptions column={column} locale={locale} />
             <HideColumn column={column} />
             {canEditFilter && <DropdownMenuSeparator />}
             {canEditFilter && (

package/src/components/data-table/header-items.tsx

@@ -29,10 +29,13 @@ import { formattingExample } from "./column-formatting/feature";
 import { formatOptions } from "./column-formatting/types";
 import { NAMELESS_COLUMN_PREFIX } from "./columns";
 
-export function renderFormatOptions<TData, TValue>(
-  column: Column<TData, TValue>,
-  locale: string,
-) {
+export function FormatOptions<TData, TValue>({
+  column,
+  locale,
+}: {
+  column: Column<TData, TValue>;
+  locale: string;
+}) {
   const dataType: DataType | undefined = column.columnDef.meta?.dataType;
   const columnFormatOptions = dataType ? formatOptions[dataType] : [];
 
@@ -83,9 +86,11 @@ export function renderFormatOptions<TData, TValue>(
   );
 }
 
-export function renderColumnWrapping<TData, TValue>(
-  column: Column<TData, TValue>,
-) {
+export function ColumnWrapping<TData, TValue>({
+  column,
+}: {
+  column: Column<TData, TValue>;
+}) {
   if (!column.getCanWrap?.() || !column.getColumnWrapping) {
     return null;
   }
@@ -108,9 +113,11 @@ export function renderColumnWrapping<TData, TValue>(
   );
 }
 
-export function renderColumnPinning<TData, TValue>(
-  column: Column<TData, TValue>,
-) {
+export function ColumnPinning<TData, TValue>({
+  column,
+}: {
+  column: Column<TData, TValue>;
+}) {
   if (!column.getCanPin?.() || !column.getIsPinned) {
     return null;
   }
@@ -157,7 +164,11 @@ export function HideColumn<TData, TValue>({
   );
 }
 
-export function renderCopyColumn<TData, TValue>(column: Column<TData, TValue>) {
+export function CopyColumn<TData, TValue>({
+  column,
+}: {
+  column: Column<TData, TValue>;
+}) {
   if (!column.getCanCopy?.()) {
     return null;
   }
@@ -177,10 +188,19 @@ export function renderCopyColumn<TData, TValue>(column: Column<TData, TValue>) {
 const AscIcon = ArrowUpNarrowWideIcon;
 const DescIcon = ArrowDownWideNarrowIcon;
 
-export function renderSorts<TData, TValue>(
-  column: Column<TData, TValue>,
-  table?: Table<TData>,
-) {
+/**
+ * `table` is optional: it is only needed to detect multi-column sorting and
+ * offer "Clear all sorts". Call sites that build their header inside column
+ * definitions (where the table instance isn't yet in scope) omit it and fall
+ * back to single-column "Clear sort".
+ */
+export function Sorts<TData, TValue>({
+  column,
+  table,
+}: {
+  column: Column<TData, TValue>;
+  table?: Table<TData>;
+}) {
   if (!column.getCanSort()) {
     return null;
   }
@@ -271,7 +291,11 @@ export function renderSortIcon<TData, TValue>(column: Column<TData, TValue>) {
   return <Icon className="h-3 w-3" />;
 }
 
-export function renderDataType<TData, TValue>(column: Column<TData, TValue>) {
+export function DataType<TData, TValue>({
+  column,
+}: {
+  column: Column<TData, TValue>;
+}) {
   const dtype: string | undefined = column.columnDef.meta?.dtype;
   if (!dtype) {
     return null;

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

@@ -48,7 +48,7 @@ import {
 import { switchLanguage } from "@/core/codemirror/language/extension";
 import { MARKDOWN_INITIAL_HIDE_CODE } from "@/core/codemirror/language/languages/markdown";
 import {
-  aiEnabledAtom,
+  aiFeaturesEnabledAtom,
   appWidthAtom,
   autoInstantiateAtom,
 } from "@/core/config/config";
@@ -100,7 +100,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
   const deleteCell = useDeleteCellCallback();
   const { openModal } = useImperativeModal();
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
-  const aiEnabled = useAtomValue(aiEnabledAtom);
+  const aiFeaturesEnabled = useAtomValue(aiFeaturesEnabledAtom);
   const autoInstantiate = useAtomValue(autoInstantiateAtom);
   const kioskMode = useAtomValue(kioskModeAtom);
   const appWidth = useAtomValue(appWidthAtom);
@@ -162,7 +162,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
       {
         icon: <SparklesIcon size={13} strokeWidth={1.5} />,
         label: "Refactor with AI",
-        hidden: !aiEnabled,
+        hidden: !aiFeaturesEnabled,
         handle: () => {
           setAiCompletionCell((current) =>
             current?.cellId === cellId ? null : { cellId },

package/src/components/editor/cell/code/cell-editor.tsx

@@ -26,7 +26,7 @@ import {
   connectedDocAtom,
   realTimeCollaboration,
 } from "@/core/codemirror/rtc/extension";
-import { autoInstantiateAtom, isAiEnabled } from "@/core/config/config";
+import { autoInstantiateAtom, isAiFeatureEnabled } from "@/core/config/config";
 import type { UserConfig } from "@/core/config/config-schema";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
 import { connectionAtom } from "@/core/network/connection";
@@ -173,13 +173,13 @@ const CellEditorInternal = ({
     });
   });
 
-  const aiEnabled = isAiEnabled(userConfig);
+  const aiFeaturesEnabled = isAiFeatureEnabled(userConfig);
 
   const extensions = useMemo(() => {
     const extensions = setupCodeMirror({
       cellId,
       showPlaceholder,
-      enableAI: aiEnabled,
+      enableAI: aiFeaturesEnabled,
       cellActions: {
         ...cellActions,
         afterToggleMarkdown,
@@ -201,6 +201,9 @@ const CellEditorInternal = ({
         splitCell,
         toggleHideCode,
         aiCellCompletion: () => {
+          if (!aiFeaturesEnabled) {
+            return false;
+          }
           let closed = false;
           setAiCompletionCell((v) => {
             // Toggle close
@@ -271,7 +274,7 @@ const CellEditorInternal = ({
     userConfig.display,
     userConfig.diagnostics,
     userConfig.ai?.inline_tooltip,
-    aiEnabled,
+    aiFeaturesEnabled,
     theme,
     showPlaceholder,
     cellActions,

package/src/components/editor/chrome/types.ts

@@ -201,16 +201,23 @@ export const PANEL_MAP = new Map<PanelType, PanelDescriptor>(
 );
 
 /**
- * Check if a panel should be hidden based on its `hidden` property
- * and `requiredCapability`.
+ * Check if a panel should be hidden based on its descriptor and runtime state.
  */
-export function isPanelHidden(
-  panel: PanelDescriptor,
-  capabilities: Capabilities,
-): boolean {
+export function isPanelHidden({
+  panel,
+  capabilities,
+  aiEnabled,
+}: {
+  panel: PanelDescriptor;
+  capabilities: Capabilities;
+  aiEnabled: boolean;
+}): boolean {
   if (panel.hidden) {
     return true;
   }
+  if (panel.type === "ai" && !aiEnabled) {
+    return true;
+  }
   if (panel.requiredCapability && !capabilities[panel.requiredCapability]) {
     return true;
   }

package/src/components/editor/chrome/wrapper/app-chrome.tsx

@@ -27,6 +27,7 @@ import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { LazyActivity } from "@/components/utils/lazy-mount";
 import { cellErrorCount } from "@/core/cells/cells";
 import { capabilitiesAtom } from "@/core/config/capabilities";
+import { aiEnabledAtom } from "@/core/config/config";
 import { getFeatureFlag } from "@/core/config/feature-flag";
 import { cn } from "@/utils/cn";
 import { ErrorBoundary } from "../../boundary/ErrorBoundary";
@@ -93,18 +94,19 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
   const [panelLayout, setPanelLayout] = useAtom(panelLayoutAtom);
   // Subscribe to capabilities to re-render when they change (e.g., terminal capability)
   const capabilities = useAtomValue(capabilitiesAtom);
+  const aiEnabled = useAtomValue(aiEnabledAtom);
 
   // Convert current developer panel items to PanelDescriptors
   // Filter out hidden panels (e.g., terminal when capability is not available)
   const devPanelItems = useMemo(() => {
     return panelLayout.developerPanel.flatMap((id) => {
       const panel = PANEL_MAP.get(id);
-      if (!panel || isPanelHidden(panel, capabilities)) {
+      if (!panel || isPanelHidden({ panel, capabilities, aiEnabled })) {
         return [];
       }
       return [panel];
     });
-  }, [panelLayout.developerPanel, capabilities]);
+  }, [panelLayout.developerPanel, capabilities, aiEnabled]);
 
   const handleSetDevPanelItems = (items: PanelDescriptor[]) => {
     setPanelLayout((prev) => ({
@@ -141,7 +143,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
   const availableDevPanels = useMemo(() => {
     const sidebarIds = new Set(panelLayout.sidebar);
     return PANELS.filter((p) => {
-      if (isPanelHidden(p, capabilities)) {
+      if (isPanelHidden({ panel: p, capabilities, aiEnabled })) {
         return false;
       }
       // Exclude panels that are in the sidebar
@@ -150,7 +152,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
       }
       return true;
     });
-  }, [panelLayout.sidebar, capabilities]);
+  }, [panelLayout.sidebar, capabilities, aiEnabled]);
 
   const emitResizeEvent = useEvent(() => {
     // HACK: Unfortunately, we have to do this twice to make sure the

package/src/components/editor/chrome/wrapper/footer-items/ai-status.tsx

@@ -4,18 +4,27 @@ import { useAtomValue } from "jotai";
 import { SparklesIcon } from "lucide-react";
 import React from "react";
 import { useOpenSettingsToTab } from "@/components/app-config/state";
-import { aiAtom, aiEnabledAtom } from "@/core/config/config";
+import {
+  aiAtom,
+  aiEnabledAtom,
+  aiModelConfiguredAtom,
+} from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FooterItem } from "../footer-item";
 
 export const AIStatusIcon: React.FC = () => {
   const ai = useAtomValue(aiAtom);
   const aiEnabled = useAtomValue(aiEnabledAtom);
+  const aiModelConfigured = useAtomValue(aiModelConfiguredAtom);
   const chatModel = ai?.models?.chat_model || DEFAULT_AI_MODEL;
   const editModel = ai?.models?.edit_model || chatModel;
   const { handleClick } = useOpenSettingsToTab();
 
   if (!aiEnabled) {
+    return null;
+  }
+
+  if (!aiModelConfigured) {
     return (
       <FooterItem
         tooltip="Assist is disabled"

package/src/components/editor/chrome/wrapper/sidebar.tsx

@@ -12,6 +12,7 @@ import {
   notebookQueuedOrRunningCountAtom,
 } from "@/core/cells/cells";
 import { capabilitiesAtom } from "@/core/config/capabilities";
+import { aiEnabledAtom } from "@/core/config/config";
 import { cn } from "@/utils/cn";
 import { FeedbackButton } from "../components/feedback-button";
 import { panelLayoutAtom, useChromeActions, useChromeState } from "../state";
@@ -30,6 +31,7 @@ export const Sidebar: React.FC = () => {
   const [panelLayout, setPanelLayout] = useAtom(panelLayoutAtom);
   // Subscribe to capabilities to re-render when they change
   const capabilities = useAtomValue(capabilitiesAtom);
+  const aiEnabled = useAtomValue(aiEnabledAtom);
 
   const renderIcon = ({ Icon }: PanelDescriptor, className?: string) => {
     return <Icon className={cn("h-5 w-5", className)} />;
@@ -40,7 +42,7 @@ export const Sidebar: React.FC = () => {
   const availableSidebarPanels = useMemo(() => {
     const devPanelIds = new Set(panelLayout.developerPanel);
     return PANELS.filter((p) => {
-      if (isPanelHidden(p, capabilities)) {
+      if (isPanelHidden({ panel: p, capabilities, aiEnabled })) {
         return false;
       }
       // Exclude panels that are in the developer panel
@@ -49,19 +51,19 @@ export const Sidebar: React.FC = () => {
       }
       return true;
     });
-  }, [panelLayout.developerPanel, capabilities]);
+  }, [panelLayout.developerPanel, capabilities, aiEnabled]);
 
   // Convert current sidebar items to PanelDescriptors
   // Filter out hidden panels (e.g., when capability is not available)
   const sidebarItems = useMemo(() => {
     return panelLayout.sidebar.flatMap((id) => {
       const panel = PANEL_MAP.get(id);
-      if (!panel || isPanelHidden(panel, capabilities)) {
+      if (!panel || isPanelHidden({ panel, capabilities, aiEnabled })) {
         return [];
       }
       return [panel];
     });
-  }, [panelLayout.sidebar, capabilities]);
+  }, [panelLayout.sidebar, capabilities, aiEnabled]);
 
   const handleSetSidebarItems = (items: PanelDescriptor[]) => {
     setPanelLayout((prev) => ({
@@ -218,7 +220,7 @@ const SidebarItem: React.FC<
   // Render as div when not clickable (e.g., inside ReorderableList)
   // This avoids nested interactive elements which break react-aria's drag behavior
   const content = onClick ? (
-    <button className={itemClassName} onClick={onClick}>
+    <button type="button" className={itemClassName} onClick={onClick}>
       {children}
     </button>
   ) : (

package/src/components/editor/errors/auto-fix.tsx

@@ -13,7 +13,7 @@ import { Tooltip } from "@/components/ui/tooltip";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { aiEnabledAtom } from "@/core/config/config";
+import { aiFeaturesEnabledAtom } from "@/core/config/config";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { cn } from "@/utils/cn";
@@ -30,9 +30,9 @@ export const AutoFixButton = ({
 }) => {
   const store = useStore();
   const { createNewCell } = useCellActions();
-  const aiEnabled = useAtomValue(aiEnabledAtom);
+  const aiFeaturesEnabled = useAtomValue(aiFeaturesEnabledAtom);
   const autoFixes = errors.flatMap((error) =>
-    getAutoFixes(error, { aiEnabled }),
+    getAutoFixes(error, { aiEnabled: aiFeaturesEnabled }),
   );
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
 

package/src/components/editor/navigation/__tests__/navigation.test.ts

@@ -1207,6 +1207,21 @@ describe("useCellNavigationProps", () => {
   });
 
   describe("AI completion functionality", () => {
+    beforeEach(() => {
+      const config = defaultUserConfig();
+      store.set(userConfigAtom, {
+        ...config,
+        ai: {
+          ...config.ai,
+          models: {
+            displayed_models: [],
+            custom_models: [],
+            edit_model: "openai/gpt-4o",
+          },
+        },
+      });
+    });
+
     it("should toggle AI completion when shortcut is pressed", () => {
       const { result } = renderWithProvider(() =>
         useCellNavigationProps(cellId1, optionsWithMockEditor),