@marimo-team/frontend 0.23.14-dev0 → 0.23.14-dev10

package/src/components/editor/chrome/panels/snippets-panel.tsx

@@ -31,6 +31,7 @@ import { cn } from "@/utils/cn";
 import { HideInKioskMode } from "../../kiosk-mode";
 import { ContributeSnippetButton } from "../components/contribute-snippet-button";
 import { usePanelOrientation, usePanelSection } from "./panel-context";
+import { getSnippetDisplay } from "./snippet-display";
 
 const extensions = [EditorView.lineWrapping];
 
@@ -187,6 +188,8 @@ const SnippetViewer: React.FC<{ snippet: Snippet; onClose: () => void }> = ({
             return null;
           }
 
+          const { language, value } = getSnippetDisplay(code);
+
           return (
             <div
               className="relative hover-actions-parent pr-2"
@@ -210,10 +213,10 @@ const SnippetViewer: React.FC<{ snippet: Snippet; onClose: () => void }> = ({
                 <LazyAnyLanguageCodeMirror
                   key={`${snippet.title}-${id}`}
                   theme={theme === "dark" ? "dark" : "light"}
-                  language="python"
+                  language={language}
                   className="cm border rounded overflow-hidden"
                   extensions={extensions}
-                  value={code}
+                  value={value}
                   readOnly={true}
                 />
               </Suspense>

package/src/components/editor/chrome/wrapper/__tests__/useOpenAiAssistant.test.ts

@@ -0,0 +1,36 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const getFeatureFlag = vi.fn<(flag: string) => boolean>();
+
+vi.mock("@/core/config/feature-flag", () => ({
+  getFeatureFlag: (flag: string) => getFeatureFlag(flag),
+}));
+
+const { resolveAiPanelTab } = await import("../useOpenAiAssistant");
+
+describe("resolveAiPanelTab", () => {
+  beforeEach(() => {
+    getFeatureFlag.mockReset();
+  });
+
+  it("honors an explicit panel regardless of the feature flag", () => {
+    getFeatureFlag.mockReturnValue(false);
+    expect(resolveAiPanelTab("agents", "chat")).toBe("agents");
+    expect(resolveAiPanelTab("chat", "agents")).toBe("chat");
+  });
+
+  it("uses the stored tab when external agents are enabled", () => {
+    getFeatureFlag.mockReturnValue(true);
+    expect(resolveAiPanelTab(undefined, "agents")).toBe("agents");
+    expect(resolveAiPanelTab(undefined, "chat")).toBe("chat");
+  });
+
+  it("falls back to chat when external agents are disabled", () => {
+    getFeatureFlag.mockReturnValue(false);
+    // A stale "agents" preference must not strand the prompt.
+    expect(resolveAiPanelTab(undefined, "agents")).toBe("chat");
+    expect(resolveAiPanelTab(undefined, "chat")).toBe("chat");
+  });
+});

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

@@ -1,47 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue } from "jotai";
-import { AlertCircleIcon } from "lucide-react";
 import type React from "react";
-import { Spinner } from "@/components/icons/spinner";
-import { Tooltip } from "@/components/ui/tooltip";
-import { wasmInitializationAtom, wasmInitStatusAtom } from "@/core/wasm/state";
+import { RuntimeStatusBadge } from "@/components/lifecycle/RuntimeStatusBadge";
 import { isWasm } from "@/core/wasm/utils";
 
 /**
- * Footer indicator that surfaces Pyodide initialization progress. Mirrors
- * the "Kernel" indicator but tracks the WASM runtime instead of the server
- * connection. Hides itself once Pyodide is ready.
+ * Footer indicator that surfaces Pyodide initialization progress. Thin
+ * wrapper around `<RuntimeStatusBadge>` so the rest of the footer chrome
+ * doesn't need to know about adapters.
  */
 export const PyodideStatus: React.FC = () => {
-  const status = useAtomValue(wasmInitStatusAtom);
-  const message = useAtomValue(wasmInitializationAtom);
-
-  if (!isWasm() || status === "ready") {
+  if (!isWasm()) {
     return null;
   }
-
-  const icon =
-    status === "error" ? (
-      <AlertCircleIcon className="w-4 h-4 text-destructive" />
-    ) : (
-      <Spinner size="small" />
-    );
-
-  const tooltip = status === "error" ? "Pyodide failed to initialize" : message;
-
-  return (
-    <Tooltip
-      content={<div className="text-sm whitespace-pre-line">{tooltip}</div>}
-      data-testid="footer-pyodide-status"
-    >
-      <div
-        className="p-1 hover:bg-accent rounded flex items-center gap-1.5 text-xs text-muted-foreground"
-        data-testid="pyodide-status"
-      >
-        {icon}
-        <span>Pyodide</span>
-      </div>
-    </Tooltip>
-  );
+  return <RuntimeStatusBadge />;
 };

package/src/components/editor/chrome/wrapper/useAiPanel.ts

@@ -6,7 +6,9 @@ import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
 const AI_PANEL_TAB_KEY = "marimo:chrome:ai-panel-tab";
 
-const aiPanelTabAtom = atomWithStorage<"chat" | "agents">(
+export type AiPanelTab = "chat" | "agents";
+
+export const aiPanelTabAtom = atomWithStorage<AiPanelTab>(
   AI_PANEL_TAB_KEY,
   "chat",
   jotaiJsonStorage,

package/src/components/editor/chrome/wrapper/useOpenAiAssistant.ts

@@ -0,0 +1,88 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { useSetAtom, useStore } from "jotai";
+import useEvent from "react-use-event-hook";
+import { agentSessionStateAtom } from "@/components/chat/acp/state";
+import { toast } from "@/components/ui/use-toast";
+import { useModelChange } from "@/core/ai/config";
+import { pendingAiPromptAtom } from "@/core/ai/state";
+import type { CopilotMode } from "@/core/ai/tools/registry";
+import { aiModelConfiguredAtom } from "@/core/config/config";
+import { getFeatureFlag } from "@/core/config/feature-flag";
+import { Logger } from "@/utils/Logger";
+import { useChromeActions } from "../state";
+import { type AiPanelTab, aiPanelTabAtom, useAiPanelTab } from "./useAiPanel";
+
+// Error fixes work best with kernel access, so we default the chat panel to
+// code mode. The agents panel manages its own mode, so this is chat-only.
+const FIX_CHAT_MODE: CopilotMode = "code_mode";
+
+export interface OpenAiAssistantOptions {
+  prompt: string;
+  submit?: boolean;
+  /** Override the user's AI sidebar tab. When omitted, the stored tab is used. */
+  panel?: AiPanelTab;
+}
+
+// Resolve which AI panel to open.
+export function resolveAiPanelTab(
+  panel: AiPanelTab | undefined,
+  storedTab: AiPanelTab,
+): AiPanelTab {
+  if (panel) {
+    return panel;
+  }
+  return getFeatureFlag("external_agents") ? storedTab : "chat";
+}
+
+/**
+ * Opens the AI assistant sidebar panel and delivers a prompt to it.
+ *
+ * If no `panel` is given, the user's configured tab is respected.
+ */
+export function useOpenAiAssistant() {
+  const { openApplication } = useChromeActions();
+  const { setAiPanelTab } = useAiPanelTab();
+  const { saveModeChange } = useModelChange();
+  const setPendingPrompt = useSetAtom(pendingAiPromptAtom);
+  const store = useStore();
+
+  return useEvent(async (opts: OpenAiAssistantOptions) => {
+    const tab = resolveAiPanelTab(opts.panel, store.get(aiPanelTabAtom));
+
+    const chatPanelReady = store.get(aiModelConfiguredAtom);
+    const agentsPanelReady =
+      getFeatureFlag("external_agents") &&
+      store.get(agentSessionStateAtom).sessions.length > 0;
+    const isReady = tab === "agents" ? agentsPanelReady : chatPanelReady;
+
+    if (!isReady) {
+      toast({
+        title: tab === "agents" ? "No agent session" : "AI not configured",
+        description:
+          tab === "agents"
+            ? "Start an agent session or switch to the chat panel."
+            : "Configure an AI provider and chat model to fix with AI.",
+        variant: "danger",
+      });
+      return;
+    }
+
+    if (opts.panel) {
+      setAiPanelTab(opts.panel);
+    }
+
+    if (tab === "chat") {
+      await saveModeChange(FIX_CHAT_MODE).catch(() =>
+        Logger.error("Failed to save mode change"),
+      );
+    }
+
+    setPendingPrompt({
+      prompt: opts.prompt,
+      submit: opts.submit ?? false,
+    });
+
+    openApplication("ai");
+  });
+}

package/src/components/editor/errors/__tests__/auto-fix.test.ts

@@ -0,0 +1,119 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import type { MarimoError } from "@/core/kernel/messages";
+
+const getDatasourceContext = vi.fn<(id: unknown) => string | null>(() => null);
+
+vi.mock("@/core/ai/context/providers/datasource", () => ({
+  getDatasourceContext: (id: unknown) => getDatasourceContext(id),
+}));
+
+const { buildFixPrompt, buildFixPromptFromText } = await import("../auto-fix");
+
+describe("buildFixPromptFromText", () => {
+  beforeEach(() => {
+    getDatasourceContext.mockReset();
+    getDatasourceContext.mockReturnValue(null);
+  });
+
+  it("includes the cell id in the header when provided", () => {
+    const prompt = buildFixPromptFromText("boom", cellId("cell-1"));
+    expect(prompt).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nboom",
+    );
+  });
+
+  it("uses a generic header when no cell id is provided", () => {
+    const prompt = buildFixPromptFromText("boom");
+    expect(prompt).toBe(
+      "My code gives the following error. Please fix it:\n\nboom",
+    );
+  });
+
+  it("appends datasource context when explicitly requested for the cell", () => {
+    getDatasourceContext.mockReturnValue("@datasource://my_db");
+    const prompt = buildFixPromptFromText("boom", cellId("cell-1"), {
+      includeDatasourceContext: true,
+    });
+    expect(prompt).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nboom\n\nDatabase schema: @datasource://my_db",
+    );
+  });
+
+  it("does not append datasource context by default", () => {
+    getDatasourceContext.mockReturnValue("@datasource://my_db");
+    const prompt = buildFixPromptFromText("boom", cellId("cell-1"));
+    expect(prompt).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nboom",
+    );
+    expect(getDatasourceContext).not.toHaveBeenCalled();
+  });
+
+  it("does not look up datasource context without a cell id", () => {
+    buildFixPromptFromText("boom", undefined, {
+      includeDatasourceContext: true,
+    });
+    expect(getDatasourceContext).not.toHaveBeenCalled();
+  });
+});
+
+describe("buildFixPrompt", () => {
+  beforeEach(() => {
+    getDatasourceContext.mockReset();
+    getDatasourceContext.mockReturnValue(null);
+  });
+
+  it("uses the error message when present", () => {
+    const errors: MarimoError[] = [
+      { type: "sql-error", msg: "syntax error", sql_statement: "SELECT" },
+    ];
+    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nsyntax error",
+    );
+  });
+
+  it("joins multiple errors with newlines", () => {
+    const errors: MarimoError[] = [
+      { type: "syntax", msg: "bad syntax" },
+      {
+        type: "exception",
+        exception_type: "ValueError",
+        msg: "bad value",
+        raising_cell: null,
+      },
+    ];
+    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nbad syntax\nbad value",
+    );
+  });
+
+  it("falls back to the error type when there is no message", () => {
+    const errors: MarimoError[] = [
+      { type: "multiple-defs", name: "foo", cells: [cellId("foo")] },
+    ];
+    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nmultiple-defs",
+    );
+  });
+
+  it("appends datasource context for SQL errors", () => {
+    getDatasourceContext.mockReturnValue("@datasource://my_db");
+    const errors: MarimoError[] = [
+      { type: "sql-error", msg: "syntax error", sql_statement: "SELECT" },
+    ];
+    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nsyntax error\n\nDatabase schema: @datasource://my_db",
+    );
+  });
+
+  it("does not append datasource context for non-SQL errors", () => {
+    getDatasourceContext.mockReturnValue("@datasource://my_db");
+    const errors: MarimoError[] = [{ type: "syntax", msg: "bad syntax" }];
+    expect(buildFixPrompt(errors, cellId("cell-1"))).toBe(
+      "My cell (id: cell-1) produced the following error. Please fix it:\n\nbad syntax",
+    );
+    expect(getDatasourceContext).not.toHaveBeenCalled();
+  });
+});

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

@@ -1,7 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { useAtomValue, useSetAtom, useStore } from "jotai";
-import { ChevronDownIcon, SparklesIcon, WrenchIcon } from "lucide-react";
+import {
+  CheckIcon,
+  ChevronDownIcon,
+  HatGlasses,
+  type LucideIcon,
+  SparklesIcon,
+  WrenchIcon,
+} from "lucide-react";
 import { Button } from "@/components/ui/button";
 import {
   DropdownMenu,
@@ -14,11 +21,46 @@ import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { aiFeaturesEnabledAtom } from "@/core/config/config";
+import { getDatasourceContext } from "@/core/ai/context/providers/datasource";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { cn } from "@/utils/cn";
+import { useOpenAiAssistant } from "../chrome/wrapper/useOpenAiAssistant";
 import { type FixMode, useFixMode } from "./fix-mode";
 
+export function buildFixPromptFromText(
+  errorText: string,
+  cellId?: CellId,
+  {
+    includeDatasourceContext = false,
+  }: { includeDatasourceContext?: boolean } = {},
+): string {
+  const header =
+    cellId != null
+      ? `My cell (id: ${cellId}) produced the following error. Please fix it:`
+      : "My code gives the following error. Please fix it:";
+  let prompt = `${header}\n\n${errorText}`;
+  if (cellId != null && includeDatasourceContext) {
+    const datasourceContext = getDatasourceContext(cellId);
+    if (datasourceContext) {
+      prompt += `\n\nDatabase schema: ${datasourceContext}`;
+    }
+  }
+  return prompt;
+}
+
+export function buildFixPrompt(errors: MarimoError[], cellId: CellId): string {
+  const errorText = errors
+    .map((error) => ("msg" in error && error.msg ? error.msg : error.type))
+    .join("\n");
+  const includeDatasourceContext = errors.some(
+    (error) => error.type === "sql-error",
+  );
+  return buildFixPromptFromText(errorText, cellId, {
+    includeDatasourceContext,
+  });
+}
+
 export const AutoFixButton = ({
   errors,
   cellId,
@@ -35,6 +77,7 @@ export const AutoFixButton = ({
     getAutoFixes(error, { aiEnabled: aiFeaturesEnabled }),
   );
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
+  const openAiAssistant = useOpenAiAssistant();
 
   if (autoFixes.length === 0) {
     return null;
@@ -67,6 +110,12 @@ export const AutoFixButton = ({
     editorView?.focus();
   };
 
+  const openAISidebar = () => {
+    openAiAssistant({
+      prompt: buildFixPrompt(errors, cellId),
+    });
+  };
+
   return (
     <div className={cn("my-2", className)}>
       {firstFix.fixType === "ai" ? (
@@ -74,6 +123,7 @@ export const AutoFixButton = ({
           tooltip={firstFix.description}
           openPrompt={() => handleFix(false)}
           applyAutofix={() => handleFix(true)}
+          openChat={openAISidebar}
         />
       ) : (
         <Tooltip content={firstFix.description} align="start">
@@ -92,23 +142,53 @@ export const AutoFixButton = ({
   );
 };
 
-const PromptIcon = SparklesIcon;
-const AutofixIcon = WrenchIcon;
+interface FixModeConfig {
+  Icon: LucideIcon;
+  title: string;
+  description: string;
+}
 
-const PromptTitle = "Suggest a prompt";
-const AutofixTitle = "Fix with AI";
+const MODE_CONFIG: Record<FixMode, FixModeConfig> = {
+  autofix: {
+    Icon: WrenchIcon,
+    title: "Inline AI Fix",
+    description: "Apply AI fixes inline in the cell",
+  },
+  chat: {
+    Icon: HatGlasses,
+    title: "Fix with AI assistant",
+    description: "Open the AI sidebar to fix",
+  },
+  prompt: {
+    Icon: SparklesIcon,
+    title: "Suggest a prompt",
+    description: "Edit the prompt before applying",
+  },
+};
+
+const FIX_MODES: FixMode[] = ["autofix", "prompt", "chat"];
 
 export const AIFixButton = ({
   tooltip,
   openPrompt,
   applyAutofix,
+  openChat,
 }: {
   tooltip: string;
   openPrompt: () => void;
   applyAutofix: () => void;
+  openChat: () => void;
 }) => {
   const { fixMode, setFixMode } = useFixMode();
 
+  let onAction = openPrompt;
+  if (fixMode === "chat") {
+    onAction = openChat;
+  } else if (fixMode === "autofix") {
+    onAction = applyAutofix;
+  }
+  const { Icon, title } = MODE_CONFIG[fixMode];
+
   return (
     <div className="flex">
       <Tooltip content={tooltip} align="start">
@@ -116,14 +196,10 @@ export const AIFixButton = ({
           size="xs"
           variant="outline"
           className="font-normal rounded-r-none border-r-0"
-          onClick={fixMode === "prompt" ? openPrompt : applyAutofix}
+          onClick={onAction}
         >
-          {fixMode === "prompt" ? (
-            <PromptIcon className="h-3 w-3 mr-2 mb-0.5" />
-          ) : (
-            <AutofixIcon className="h-3 w-3 mr-2 mb-0.5" />
-          )}
-          {fixMode === "prompt" ? PromptTitle : AutofixTitle}
+          <Icon className="h-3 w-3 mr-2 mb-0.5" />
+          {title}
         </Button>
       </Tooltip>
       <DropdownMenu>
@@ -138,40 +214,38 @@ export const AIFixButton = ({
           </Button>
         </DropdownMenuTrigger>
         <DropdownMenuContent align="end" className="w-56">
-          <DropdownMenuItem
-            className="flex items-center gap-2"
-            onClick={() => {
-              setFixMode(fixMode === "prompt" ? "autofix" : "prompt");
-            }}
-          >
-            <AiModeItem mode={fixMode === "prompt" ? "autofix" : "prompt"} />
-          </DropdownMenuItem>
+          {FIX_MODES.map((mode) => (
+            <DropdownMenuItem
+              key={mode}
+              className="flex items-center gap-2"
+              onClick={() => setFixMode(mode)}
+            >
+              <AiModeItem mode={mode} selected={mode === fixMode} />
+            </DropdownMenuItem>
+          ))}
         </DropdownMenuContent>
       </DropdownMenu>
     </div>
   );
 };
 
-const AiModeItem = ({ mode }: { mode: FixMode }) => {
-  const icon =
-    mode === "prompt" ? (
-      <PromptIcon className="h-4 w-4" />
-    ) : (
-      <AutofixIcon className="h-4 w-4" />
-    );
-  const title = mode === "prompt" ? PromptTitle : AutofixTitle;
-  const description =
-    mode === "prompt"
-      ? "Edit the prompt before applying"
-      : "Apply AI fixes automatically";
+const AiModeItem = ({
+  mode,
+  selected,
+}: {
+  mode: FixMode;
+  selected: boolean;
+}) => {
+  const { Icon, title, description } = MODE_CONFIG[mode];
 
   return (
-    <div className="flex items-center gap-2">
-      {icon}
+    <div className="flex items-center gap-2 w-full">
+      <Icon className="h-4 w-4 shrink-0" />
       <div className="flex flex-col">
         <span className="font-medium">{title}</span>
         <span className="text-xs text-muted-foreground">{description}</span>
       </div>
+      {selected && <CheckIcon className="h-4 w-4 ml-auto shrink-0" />}
     </div>
   );
 };

package/src/components/editor/errors/fix-mode.ts

@@ -4,7 +4,7 @@ import { useAtom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import { jotaiJsonStorage } from "@/utils/storage/jotai";
 
-export type FixMode = "prompt" | "autofix";
+export type FixMode = "prompt" | "autofix" | "chat";
 
 const BASE_KEY = "marimo:ai-autofix-mode";
 

package/src/components/editor/output/MarimoTracebackOutput.tsx

@@ -41,7 +41,8 @@ import {
   extractAllTracebackInfo,
   getTracebackInfo,
 } from "@/utils/traceback";
-import { AIFixButton } from "../errors/auto-fix";
+import { useOpenAiAssistant } from "../chrome/wrapper/useOpenAiAssistant";
+import { AIFixButton, buildFixPromptFromText } from "../errors/auto-fix";
 import { MangledSegments } from "../errors/mangled-local-chip";
 import { CellLinkTraceback } from "../links/cell-link";
 import type { OnRefactorWithAI } from "../Output";
@@ -71,6 +72,7 @@ export const MarimoTracebackOutput = ({
 
   const lastTracebackLine = lastLine(traceback);
   const aiFeaturesEnabled = useAtomValue(aiFeaturesEnabledAtom);
+  const openAiAssistant = useOpenAiAssistant();
 
   // Get last traceback info
   const tracebackInfo = extractAllTracebackInfo(traceback)?.at(0);
@@ -97,6 +99,12 @@ export const MarimoTracebackOutput = ({
     });
   };
 
+  const openAISidebar = () => {
+    openAiAssistant({
+      prompt: buildFixPromptFromText(lastTracebackLine, cellId),
+    });
+  };
+
   return (
     <div className="flex flex-col gap-2 min-w-full w-fit">
       <button
@@ -126,6 +134,7 @@ export const MarimoTracebackOutput = ({
             tooltip="Fix with AI"
             openPrompt={() => handleRefactorWithAI(false)}
             applyAutofix={() => handleRefactorWithAI(true)}
+            openChat={openAISidebar}
           />
         )}
         {showDebugger && (

package/src/components/editor/output/__tests__/traceback.test.tsx

@@ -1,11 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { render } from "@testing-library/react";
+import { Provider } from "jotai";
 import { beforeAll, beforeEach, describe, expect, test, vi } from "vitest";
+import { MockRequestClient } from "@/__mocks__/requests";
 import { Tracebacks } from "@/__mocks__/tracebacks";
 import { cellId } from "@/__tests__/branded";
 import { TooltipProvider } from "@/components/ui/tooltip";
 import { initialModeAtom } from "@/core/mode";
+import { requestClientAtom } from "@/core/network/requests";
 import { store } from "@/core/state/jotai";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
@@ -20,13 +23,16 @@ describe("traceback component", () => {
   beforeEach(() => {
     vi.resetAllMocks();
     store.set(initialModeAtom, "edit");
+    store.set(requestClientAtom, MockRequestClient.create());
   });
 
   test("extracts cell-link", () => {
     const traceback = (
-      <TooltipProvider>
-        <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cid} />
-      </TooltipProvider>
+      <Provider store={store}>
+        <TooltipProvider>
+          <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cid} />
+        </TooltipProvider>
+      </Provider>
     );
     const { unmount, getAllByRole } = render(traceback);
 
@@ -44,9 +50,11 @@ describe("traceback component", () => {
 
   test("renames File to Cell for relevant lines", () => {
     const traceback = (
-      <TooltipProvider>
-        <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cid} />
-      </TooltipProvider>
+      <Provider store={store}>
+        <TooltipProvider>
+          <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cid} />
+        </TooltipProvider>
+      </Provider>
     );
     const { unmount, container } = render(traceback);