@marimo-team/islands 0.17.0 → 0.17.2

package/src/components/chat/chat-panel.tsx

@@ -5,7 +5,7 @@ import { useChat } from "@ai-sdk/react";
 import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import { DefaultChatTransport, type ToolUIPart } from "ai";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
+import { useAtom, useAtomValue, useSetAtom, useStore } from "jotai";
 import {
   AtSignIcon,
   BotMessageSquareIcon,
@@ -30,13 +30,19 @@ import {
 import { replaceMessagesInChat } from "@/core/ai/chat-utils";
 import { useModelChange } from "@/core/ai/config";
 import { AiModelId, type ProviderId } from "@/core/ai/ids/ids";
+import { useStagedAICellsActions } from "@/core/ai/staged-cells";
 import {
   activeChatAtom,
   type Chat,
   type ChatId,
   chatStateAtom,
 } from "@/core/ai/state";
-import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
+import type { ToolNotebookContext } from "@/core/ai/tools/base";
+import {
+  type CopilotMode,
+  FRONTEND_TOOL_REGISTRY,
+} from "@/core/ai/tools/registry";
+import { useCellActions } from "@/core/cells/cells";
 import { aiAtom, aiEnabledAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FeatureFlagged } from "@/core/config/feature-flag";
@@ -45,7 +51,6 @@ import { useRuntimeManager } from "@/core/runtime/config";
 import { ErrorBanner } from "@/plugins/impl/common/error-banner";
 import { cn } from "@/utils/cn";
 import { Logger } from "@/utils/Logger";
-
 import { AIModelDropdown } from "../ai/ai-model-dropdown";
 import { useOpenSettingsToTab } from "../app-config/state";
 import { PromptInput } from "../editor/ai/add-cell-with-ai";
@@ -301,7 +306,11 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
 
     const { saveModeChange } = useModelChange();
 
-    const modeOptions = [
+    const modeOptions: {
+      value: CopilotMode;
+      label: string;
+      subtitle: string;
+    }[] = [
       {
         value: "ask",
         label: "Ask",
@@ -313,6 +322,11 @@ const ChatInputFooter: React.FC<ChatInputFooterProps> = memo(
         label: "Manual",
         subtitle: "Pure chat, no tool usage",
       },
+      {
+        value: "agent",
+        label: "Agent (beta)",
+        subtitle: "Use AI with access to read and write tools",
+      },
     ];
 
     const isAttachmentSupported =
@@ -511,9 +525,20 @@ const ChatPanelBody = () => {
   const fileInputRef = useRef<HTMLInputElement>(null);
   const messagesEndRef = useRef<HTMLDivElement>(null);
   const runtimeManager = useRuntimeManager();
-  const { invokeAiTool } = useRequestClient();
+  const { invokeAiTool, sendRun } = useRequestClient();
 
   const activeChatId = activeChat?.id;
+  const store = useStore();
+
+  const { addStagedCell } = useStagedAICellsActions();
+  const { createNewCell, prepareForRun } = useCellActions();
+  const toolContext: ToolNotebookContext = {
+    addStagedCell,
+    createNewCell,
+    prepareForRun,
+    sendRun,
+    store,
+  };
 
   const {
     messages,
@@ -555,6 +580,13 @@ const ChatPanelBody = () => {
       });
     },
     onToolCall: async ({ toolCall }) => {
+      // Dynamic tool calls will throw an error for toolName
+      // https://ai-sdk.dev/docs/ai-sdk-ui/chatbot-tool-usage#client-side-page
+      if (toolCall.dynamic) {
+        Logger.debug("Skipping dynamic tool call", toolCall);
+        return;
+      }
+
       await handleToolCall({
         invokeAiTool,
         addToolResult,
@@ -563,6 +595,7 @@ const ChatPanelBody = () => {
           toolCallId: toolCall.toolCallId,
           input: toolCall.input as Record<string, never>,
         },
+        toolContext,
       });
     },
     onError: (error) => {

package/src/components/chat/chat-utils.ts

@@ -2,6 +2,7 @@
 
 import type { components } from "@marimo-team/marimo-api";
 import type { FileUIPart, ToolUIPart, UIMessage } from "ai";
+import type { ToolNotebookContext } from "@/core/ai/tools/base";
 import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
 import type {
   InvokeAiToolRequest,
@@ -106,8 +107,9 @@ interface AddToolResult {
 
 export async function handleToolCall({
   invokeAiTool,
-  addToolResult,
+  addToolResult, // Important that we don't await addToolResult to prevent potential deadlocks
   toolCall,
+  toolContext,
 }: {
   invokeAiTool: (request: InvokeAiToolRequest) => Promise<InvokeAiToolResponse>;
   addToolResult: (result: AddToolResult) => Promise<void>;
@@ -116,6 +118,7 @@ export async function handleToolCall({
     toolCallId: string;
     input: Record<string, never>;
   };
+  toolContext: ToolNotebookContext;
 }) {
   try {
     if (FRONTEND_TOOL_REGISTRY.has(toolCall.toolName)) {
@@ -123,6 +126,7 @@ export async function handleToolCall({
       const response = await FRONTEND_TOOL_REGISTRY.invoke(
         toolCall.toolName,
         toolCall.input,
+        toolContext,
       );
       addToolResult({
         tool: toolCall.toolName,

package/src/components/chat/markdown-renderer.css

@@ -33,6 +33,7 @@
 
   code {
     font-size: inherit !important;
+    white-space: pre-wrap;
   }
 
   li.task-list-item {

package/src/components/chat/tool-call-accordion.tsx

@@ -119,6 +119,8 @@ const ToolArgsRenderer: React.FC<{ input: unknown }> = ({ input }) => {
     return null;
   }
 
+  const isEmptyInput = isEmpty(input);
+
   const isObject =
     typeof input === "object" &&
     Object.keys(input as Record<string, unknown>).length > 0;
@@ -129,7 +131,11 @@ const ToolArgsRenderer: React.FC<{ input: unknown }> = ({ input }) => {
         Tool Request
       </h3>
       <pre className="bg-[var(--slate-2)] p-2 text-muted-foreground border border-[var(--slate-4)] rounded text-xs overflow-auto scrollbar-thin max-h-64">
-        {isObject ? JSON.stringify(input, null, 2) : String(input)}
+        {isEmptyInput
+          ? "{}"
+          : isObject
+            ? JSON.stringify(input, null, 2)
+            : String(input)}
       </pre>
     </div>
   );

package/src/components/data-table/charts/__tests__/storage.test.ts

@@ -2,8 +2,18 @@
 
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import { asMock, SetupMocks } from "@/__mocks__/common";
+
+// Mock the storage module to use our mock storage
+vi.mock("@/utils/storage/storage", () => ({
+  availableStorage: {
+    getItem: vi.fn(),
+    setItem: vi.fn(),
+    removeItem: vi.fn(),
+  },
+}));
+
 import type { CellId } from "@/core/cells/ids";
+import { availableStorage } from "@/utils/storage/storage";
 import { ChartSchema } from "../schemas";
 import type { TabName } from "../storage";
 import { KEY, tabsStorageAtom } from "../storage";
@@ -12,7 +22,6 @@ import { ChartType } from "../types";
 describe("Chart Transforms Storage", () => {
   beforeEach(() => {
     vi.clearAllMocks();
-    SetupMocks.localStorage();
     // Reset the store before each test
     const store = getDefaultStore();
     store.set(tabsStorageAtom, new Map());
@@ -87,14 +96,11 @@ describe("Chart Transforms Storage", () => {
 
   describe("LocalStorage integration", () => {
     it("should call localStorage.setItem when setting data", () => {
-      // Reset the mock to ensure clean state
-      asMock(window.localStorage.setItem).mockReset();
-
       const store = getDefaultStore();
       const newMap = new Map();
       store.set(tabsStorageAtom, newMap);
 
-      expect(window.localStorage.setItem).toHaveBeenCalledWith(
+      expect(availableStorage.setItem).toHaveBeenCalledWith(
         KEY,
         JSON.stringify([]),
       );

package/src/components/data-table/charts/charts.tsx

@@ -48,6 +48,7 @@ const CHART_HEIGHT = 290;
 
 export interface TablePanelProps {
   cellId: CellId | null;
+  data: unknown[];
   dataTable: JSX.Element;
   displayHeader: boolean;
   getDataUrl?: GetDataUrl;
@@ -56,6 +57,7 @@ export interface TablePanelProps {
 
 export const TablePanel: React.FC<TablePanelProps> = ({
   cellId,
+  data,
   dataTable,
   getDataUrl,
   fieldTypes,
@@ -211,6 +213,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
         return (
           <TabsContent key={idx} value={tab.tabName} className="h-[400px] mt-1">
             <ChartPanel
+              tableData={data}
               chartConfig={tab.config}
               chartType={tab.chartType}
               saveChart={saveChart}
@@ -228,6 +231,7 @@ export const TablePanel: React.FC<TablePanelProps> = ({
 const CHART_PLACEHOLDER_CODE = "X and Y columns are not set";
 
 export const ChartPanel: React.FC<{
+  tableData: unknown[];
   chartConfig: ChartSchemaType | null;
   chartType: ChartType;
   saveChart: (formValues: ChartSchemaType) => void;
@@ -235,6 +239,7 @@ export const ChartPanel: React.FC<{
   getDataUrl?: GetDataUrl;
   fieldTypes?: FieldTypesWithExternalType | null;
 }> = ({
+  tableData,
   chartConfig,
   chartType,
   saveChart,
@@ -255,7 +260,7 @@ export const ChartPanel: React.FC<{
   const { ref: chartContainerRef } = useResizeObserver();
 
   const { data, isPending, error } = useAsyncData(async () => {
-    if (!getDataUrl) {
+    if (!getDataUrl || tableData.length === 0) {
       return [];
     }
 
@@ -276,7 +281,8 @@ export const ChartPanel: React.FC<{
       },
     );
     return chartData;
-  }, []);
+    // Re-run when the data table changes
+  }, [tableData]);
 
   const formValues = form.watch();
 

package/src/components/data-table/charts/storage.ts

@@ -5,7 +5,7 @@ import { capitalize } from "lodash-es";
 import { z } from "zod";
 import type { CellId } from "@/core/cells/ids";
 import { Logger } from "@/utils/Logger";
-import { NotebookScopedLocalStorage } from "@/utils/localStorage";
+import { NotebookScopedLocalStorage } from "@/utils/storage/typed";
 import type { TypedString } from "@/utils/typed";
 import { ChartSchema, type ChartSchemaType } from "./schemas";
 import type { ChartType } from "./types";

package/src/components/data-table/column-wrapping/feature.ts

@@ -13,6 +13,8 @@ import type {
   ColumnWrappingTableState,
 } from "./types";
 
+export const COLUMN_WRAPPING_STYLES = "whitespace-pre-wrap min-w-[200px]";
+
 export const ColumnWrappingFeature: TableFeature = {
   getInitialState: (state): ColumnWrappingTableState => {
     return {

package/src/components/data-table/columns.tsx

@@ -21,6 +21,7 @@ import { Tooltip } from "../ui/tooltip";
 import { DataTableColumnHeader } from "./column-header";
 import type { ColumnChartSpecModel } from "./column-summary/chart-spec-model";
 import { TableColumnSummary } from "./column-summary/column-summary";
+import { COLUMN_WRAPPING_STYLES } from "./column-wrapping/feature";
 import { DatePopover } from "./date-popover";
 import type { FilterType } from "./filters";
 import { getMimeValues, MimeCell } from "./mime-cell";
@@ -315,6 +316,7 @@ const PopoutColumn = ({
   rawStringValue,
   contentClassName,
   buttonText,
+  wrapped,
   children,
 }: {
   cellStyles?: string;
@@ -322,6 +324,7 @@ const PopoutColumn = ({
   rawStringValue: string;
   contentClassName?: string;
   buttonText?: string;
+  wrapped?: boolean;
   children: React.ReactNode;
 }) => {
   return (
@@ -336,7 +339,10 @@ const PopoutColumn = ({
           }}
         >
           <span
-            className="cursor-pointer hover:text-link"
+            className={cn(
+              "cursor-pointer hover:text-link",
+              wrapped && COLUMN_WRAPPING_STYLES,
+            )}
             title={rawStringValue}
           >
             {rawStringValue}
@@ -408,7 +414,7 @@ function getCellStyleClass(
     "truncate",
     justify === "center" && "text-center",
     justify === "right" && "text-right",
-    wrapped && "whitespace-pre-wrap min-w-[200px] break-words",
+    wrapped && `${COLUMN_WRAPPING_STYLES} break-words`,
   );
 }
 
@@ -517,6 +523,7 @@ export function renderCellValue<TData, TValue>({
         rawStringValue={stringValue}
         contentClassName="max-h-64 overflow-auto whitespace-pre-wrap break-words text-sm"
         buttonText="X"
+        wrapped={column.getColumnWrapping?.() === "wrap"}
       >
         <UrlDetector parts={parts} />
       </PopoutColumn>
@@ -558,6 +565,7 @@ export function renderCellValue<TData, TValue>({
         cellStyles={cellStyles}
         selectCell={selectCell}
         rawStringValue={rawStringValue}
+        wrapped={column.getColumnWrapping?.() === "wrap"}
       >
         <JsonOutput data={value} format="tree" className="max-h-64" />
       </PopoutColumn>

package/src/components/data-table/renderers.tsx

@@ -20,6 +20,7 @@ import {
   TableRow,
 } from "@/components/ui/table";
 import { cn } from "@/utils/cn";
+import { COLUMN_WRAPPING_STYLES } from "./column-wrapping/feature";
 import { CellRangeSelectionIndicator } from "./range-focus/cell-selection-indicator";
 import { useCellRangeSelection } from "./range-focus/use-cell-range-selection";
 import { useScrollIntoViewOnFocus } from "./range-focus/use-scroll-into-view";
@@ -130,8 +131,8 @@ export const DataTableBody = <TData,>({
           className={cn(
             "whitespace-pre truncate max-w-[300px] outline-hidden",
             cell.column.getColumnWrapping &&
-              cell.column.getColumnWrapping() === "wrap" &&
-              "whitespace-pre-wrap min-w-[200px]",
+              cell.column.getColumnWrapping?.() === "wrap" &&
+              COLUMN_WRAPPING_STYLES,
             "px-1.5 py-[0.18rem]",
             className,
           )}

package/src/components/editor/Output.tsx

@@ -26,8 +26,8 @@ import {
   ExpandIcon,
 } from "lucide-react";
 import { useExpandedOutput } from "@/core/cells/outputs";
+import { useIframeCapabilities } from "@/hooks/useIframeCapabilities";
 import { renderHTML } from "@/plugins/core/RenderHTML";
-import { LazyAnyLanguageCodeMirror } from "@/plugins/impl/code/LazyAnyLanguageCodeMirror";
 import { Banner } from "@/plugins/impl/common/error-banner";
 import type { TopLevelFacetedUnitSpec } from "@/plugins/impl/data-explorer/queries/types";
 import { useTheme } from "@/theme/useTheme";
@@ -93,7 +93,14 @@ export const OutputRenderer: React.FC<{
         typeof data === "string",
         `Expected string data for mime=${mimetype}. Got ${typeof data}`,
       );
-      return <HtmlOutput className={channel} html={data} />;
+      // We don't sanitize HTML in text/html to allow for iframes or rich javascript content.
+      return (
+        <HtmlOutput
+          className={channel}
+          html={data}
+          alwaysSanitizeHtml={false}
+        />
+      );
 
     case "text/plain":
       invariant(
@@ -124,7 +131,7 @@ export const OutputRenderer: React.FC<{
         typeof data === "string",
         `Expected string data for mime=${mimetype}. Got ${typeof data}`,
       );
-      return renderHTML({ html: data });
+      return renderHTML({ html: data, alwaysSanitizeHtml: true });
 
     case "video/mp4":
     case "video/mpeg":
@@ -164,13 +171,7 @@ export const OutputRenderer: React.FC<{
         `Expected string data for mime=${mimetype}. Got ${typeof data}`,
       );
       return (
-        <LazyAnyLanguageCodeMirror
-          theme={theme === "dark" ? "dark" : "light"}
-          value={data}
-          readOnly={true}
-          editable={false}
-          language="markdown"
-        />
+        <HtmlOutput className={channel} html={data} alwaysSanitizeHtml={true} />
       );
     case "application/vnd.vegalite.v5+json":
     case "application/vnd.vega.v5+json":
@@ -378,6 +379,7 @@ const ExpandableOutput = React.memo(
     const containerRef = useRef<HTMLDivElement>(null);
     const [isExpanded, setIsExpanded] = useExpandedOutput(cellId);
     const [isOverflowing, setIsOverflowing] = useState(false);
+    const { hasFullscreen } = useIframeCapabilities();
 
     // Create resize observer to detect overflow
     useEffect(() => {
@@ -403,20 +405,22 @@ const ExpandableOutput = React.memo(
         <div>
           <div className="relative print:hidden">
             <div className="absolute -right-9 top-1 z-1 flex flex-col gap-1">
-              <Tooltip content="Fullscreen" side="left">
-                <Button
-                  data-testid="fullscreen-output-button"
-                  className="hover-action hover:bg-muted p-1 hover:border-border border border-transparent"
-                  onClick={async () => {
-                    await containerRef.current?.requestFullscreen();
-                  }}
-                  onMouseDown={Events.preventFocus}
-                  size="xs"
-                  variant="text"
-                >
-                  <ExpandIcon className="size-4" strokeWidth={1.25} />
-                </Button>
-              </Tooltip>
+              {hasFullscreen && (
+                <Tooltip content="Fullscreen" side="left">
+                  <Button
+                    data-testid="fullscreen-output-button"
+                    className="hover-action hover:bg-muted p-1 hover:border-border border border-transparent"
+                    onClick={async () => {
+                      await containerRef.current?.requestFullscreen();
+                    }}
+                    onMouseDown={Events.preventFocus}
+                    size="xs"
+                    variant="text"
+                  >
+                    <ExpandIcon className="size-4" strokeWidth={1.25} />
+                  </Button>
+                </Tooltip>
+              )}
               {(isOverflowing || isExpanded) && !forceExpand && (
                 <Button
                   data-testid="expand-output-button"

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

@@ -39,7 +39,11 @@ import { Switch } from "@/components/ui/switch";
 import { toast } from "@/components/ui/use-toast";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { maybeAddMarimoImport } from "@/core/cells/add-missing-import";
-import { hasOnlyOneCellAtom, useCellActions } from "@/core/cells/cells";
+import {
+  hasOnlyOneCellAtom,
+  SETUP_CELL_ID,
+  useCellActions,
+} from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import type { CellData } from "@/core/cells/types";
 import { formatEditorViews } from "@/core/codemirror/format";
@@ -136,6 +140,8 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
     }
   };
 
+  const isSetupCell = cellId === SETUP_CELL_ID;
+
   // Actions
   const actions: ActionButton[][] = [
     [
@@ -181,6 +187,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
             onEnterKey={() => closePopover?.()}
           />
         ),
+        hidden: isSetupCell,
       },
       {
         icon: <PlayIcon size={13} strokeWidth={1.5} />,
@@ -254,6 +261,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
           />
         ),
         handle: toggleDisabled,
+        hidden: isSetupCell,
       },
       {
         icon: <XCircleIcon size={13} strokeWidth={1.5} />,
@@ -282,6 +290,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
             keepCodeAsIs: false,
           });
         },
+        hidden: isSetupCell,
       },
       {
         icon: <DatabaseIcon size={13} strokeWidth={1.5} />,
@@ -297,6 +306,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
             keepCodeAsIs: false,
           });
         },
+        hidden: isSetupCell,
       },
       {
         icon: <PythonIcon />,
@@ -309,6 +319,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
           maybeAddMarimoImport({ autoInstantiate, createNewCell: createCell });
           toggleToLanguage(editorView, "python", { force: true });
         },
+        hidden: isSetupCell,
       },
     ],
 
@@ -324,6 +335,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         label: "Create cell above",
         hotkey: "cell.createAbove",
         handle: () => createCell({ cellId, before: true }),
+        hidden: isSetupCell,
       },
       {
         icon: (
@@ -341,26 +353,28 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         label: "Move cell up",
         hotkey: "cell.moveUp",
         handle: () => moveCell({ cellId, before: true }),
+        hidden: isSetupCell,
       },
       {
         icon: <ChevronDownIcon size={13} strokeWidth={1.5} />,
         label: "Move cell down",
         hotkey: "cell.moveDown",
         handle: () => moveCell({ cellId, before: false }),
+        hidden: isSetupCell,
       },
       {
         icon: <ChevronLeftIcon size={13} strokeWidth={1.5} />,
         label: "Move cell left",
         hotkey: "cell.moveLeft",
         handle: () => moveCell({ cellId, direction: "left" }),
-        hidden: appWidth !== "columns",
+        hidden: appWidth !== "columns" || isSetupCell,
       },
       {
         icon: <ChevronRightIcon size={13} strokeWidth={1.5} />,
         label: "Move cell right",
         hotkey: "cell.moveRight",
         handle: () => moveCell({ cellId, direction: "right" }),
-        hidden: appWidth !== "columns",
+        hidden: appWidth !== "columns" || isSetupCell,
       },
       {
         icon: <ChevronsUpIcon size={13} strokeWidth={1.5} />,
@@ -369,6 +383,7 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         // When using the cell menu, likely the user doesn't want to scroll
         // and instead just wants to get the cell out of the way
         handle: () => sendToTop({ cellId, scroll: false }),
+        hidden: isSetupCell,
       },
       {
         icon: <ChevronsDownIcon size={13} strokeWidth={1.5} />,
@@ -377,13 +392,14 @@ export function useCellActionButtons({ cell, closePopover }: Props) {
         // When using the cell menu, likely the user doesn't want to scroll
         // and instead just wants to get the cell out of the way
         handle: () => sendToBottom({ cellId, scroll: false }),
+        hidden: isSetupCell,
       },
       {
         icon: <Columns2Icon size={13} strokeWidth={1.5} />,
         label: "Break into new column",
         hotkey: "cell.addColumnBreakpoint",
-        hidden: appWidth !== "columns",
         handle: () => addColumnBreakpoint({ cellId }),
+        hidden: appWidth !== "columns" || isSetupCell,
       },
     ],
 

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -44,6 +44,8 @@ import {
 } from "@/components/ui/dropdown-menu";
 import { toast } from "@/components/ui/use-toast";
 import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
+import type { ToolNotebookContext } from "@/core/ai/tools/base";
+import { useCellActions } from "@/core/cells/cells";
 import { resourceExtension } from "@/core/codemirror/ai/resources";
 import { useRequestClient } from "@/core/network/requests";
 import type { AiCompletionRequest } from "@/core/network/types";
@@ -51,7 +53,8 @@ import { useRuntimeManager } from "@/core/runtime/config";
 import { useTheme } from "@/theme/useTheme";
 import { cn } from "@/utils/cn";
 import { prettyError } from "@/utils/errors";
-import { ZodLocalStorage } from "@/utils/localStorage";
+import { jotaiJsonStorage } from "@/utils/storage/jotai";
+import { ZodLocalStorage } from "@/utils/storage/typed";
 import { PythonIcon } from "../cell/code/icons";
 import {
   CompletionActions,
@@ -64,6 +67,7 @@ import { StreamingChunkTransport } from "./transport/chat-transport";
 const languageAtom = atomWithStorage<"python" | "sql">(
   "marimo:ai-language",
   "python",
+  jotaiJsonStorage,
 );
 
 const KEY = "marimo:ai-prompt-history";
@@ -79,15 +83,24 @@ export const AddCellWithAI: React.FC<{
   const store = useStore();
   const [input, setInput] = useState("");
 
-  const { deleteAllStagedCells, clearStagedCells, onStream } =
+  const { deleteAllStagedCells, clearStagedCells, onStream, addStagedCell } =
     useStagedCells(store);
   const [language, setLanguage] = useAtom(languageAtom);
   const runtimeManager = useRuntimeManager();
-  const { invokeAiTool } = useRequestClient();
+  const { invokeAiTool, sendRun } = useRequestClient();
 
   const stagedAICells = useAtomValue(stagedAICellsAtom);
   const inputRef = useRef<ReactCodeMirrorRef>(null);
 
+  const { createNewCell, prepareForRun } = useCellActions();
+  const toolContext: ToolNotebookContext = {
+    store,
+    addStagedCell,
+    createNewCell,
+    prepareForRun,
+    sendRun,
+  };
+
   const { sendMessage, stop, status, addToolResult } = useChat({
     // Throttle the messages and data updates to 100ms
     experimental_throttle: 100,
@@ -125,6 +138,7 @@ export const AddCellWithAI: React.FC<{
           toolCallId: toolCall.toolCallId,
           input: toolCall.input as Record<string, never>,
         },
+        toolContext,
       });
     },
     onError: (error) => {