@marimo-team/islands 0.16.5 → 0.17.1

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/dependency-graph/custom-node.tsx

@@ -1,13 +1,18 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { useAtomValue } from "jotai";
-import React, { memo, use, useId } from "react";
+import React, { memo, use } from "react";
 import { Handle, Position, useStore } from "reactflow";
 import { TinyCode } from "@/components/editor/cell/TinyCode";
 import { useCellIds } from "@/core/cells/cells";
 import { displayCellName } from "@/core/cells/names";
 import { cn } from "@/utils/cn";
-import { type CustomNodeProps, getNodeHeight } from "./elements";
+import {
+  type CustomNodeProps,
+  getNodeHeight,
+  INPUTS_HANDLE_ID,
+  OUTPUTS_HANDLE_ID,
+} from "./elements";
 import type { LayoutDirection } from "./types";
 
 function getWidth(canvasWidth: number) {
@@ -37,24 +42,19 @@ export const CustomNode = memo((props: CustomNodeProps) => {
   const reactFlowWidth = useStore(({ width }) => width);
   const edgeMarkers = use(EdgeMarkerContext);
 
-  const inputOneId = useId();
-  const inputTwoId = useId();
-  const outputOneId = useId();
-  const outputTwoId = useId();
-
   const linesOfCode = cell.code.split("\n").length;
   return (
     <div>
       <Handle
         type="target"
-        id={inputOneId}
+        id={INPUTS_HANDLE_ID}
         data-testid="input-one"
         position={edgeMarkers === "LR" ? Position.Left : Position.Top}
         style={{ background: color }}
       />
       <Handle
         type="source"
-        id={inputTwoId}
+        id={INPUTS_HANDLE_ID}
         data-testid="input-two"
         position={edgeMarkers === "LR" ? Position.Left : Position.Top}
         style={{ background: color }}
@@ -76,14 +76,14 @@ export const CustomNode = memo((props: CustomNodeProps) => {
       </div>
       <Handle
         type="source"
-        id={outputOneId}
+        id={OUTPUTS_HANDLE_ID}
         data-testid="output-one"
         position={edgeMarkers === "LR" ? Position.Right : Position.Bottom}
         style={{ background: color }}
       />
       <Handle
         type="target"
-        id={outputTwoId}
+        id={OUTPUTS_HANDLE_ID}
         data-testid="output-two"
         position={edgeMarkers === "LR" ? Position.Right : Position.Bottom}
         style={{ background: color }}

package/src/components/dependency-graph/elements.ts

@@ -19,6 +19,10 @@ export function getNodeHeight(linesOfCode: number) {
   return Math.min(linesOfCode * LINE_HEIGHT + 35, 200);
 }
 
+// The nodes must have the same handle IDs to ensure edges connect correctly
+export const OUTPUTS_HANDLE_ID = "outputs";
+export const INPUTS_HANDLE_ID = "inputs";
+
 interface ElementsBuilder {
   createElements: (
     cellIds: CellId[],
@@ -93,8 +97,8 @@ export class VerticalElementsBuilder implements ElementsBuilder {
           }
           visited.add(key);
           edges.push(
-            this.createEdge(fromId, toId, "inputs"),
-            this.createEdge(fromId, toId, "outputs"),
+            this.createEdge(fromId, toId, INPUTS_HANDLE_ID),
+            this.createEdge(fromId, toId, OUTPUTS_HANDLE_ID),
           );
         }
       }
@@ -114,8 +118,9 @@ export class TreeElementsBuilder implements ElementsBuilder {
       // Make thicker
       style: { strokeWidth: 2 },
       source: source,
-      sourceHandle: "outputs",
-      targetHandle: "inputs",
+      // Use the same handle ids as the custom node
+      sourceHandle: OUTPUTS_HANDLE_ID,
+      targetHandle: INPUTS_HANDLE_ID,
       target: target,
     };
   }

package/src/components/dependency-graph/panels.tsx

@@ -12,7 +12,7 @@ import {
   WorkflowIcon,
   XIcon,
 } from "lucide-react";
-import React, { memo } from "react";
+import React, { memo, useId } from "react";
 import { type Edge, Panel } from "reactflow";
 import { getCellEditorView } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
@@ -44,6 +44,8 @@ export const GraphToolbar: React.FC<Props> = memo(
       onSettingsChange({ ...settings, [key]: value });
     };
 
+    const checkboxId = useId();
+
     const settingsButton = (
       <Popover>
         <PopoverTrigger asChild={true}>
@@ -56,13 +58,13 @@ export const GraphToolbar: React.FC<Props> = memo(
           <div className="flex items-center gap-2">
             <Checkbox
               data-testid="hide-pure-markdown-checkbox"
-              id="hide-pure-markdown"
+              id={checkboxId}
               checked={settings.hidePureMarkdown}
               onCheckedChange={(checked) =>
                 handleSettingChange("hidePureMarkdown", Boolean(checked))
               }
             />
-            <Label htmlFor="hide-pure-markdown">Hide pure markdown</Label>
+            <Label htmlFor={checkboxId}>Hide pure markdown</Label>
           </div>
         </PopoverContent>
       </Popover>
@@ -114,7 +116,7 @@ export const GraphSelectionPanel: React.FC<{
   onClearSelection: () => void;
   edges: Edge[];
   variables: Variables;
-}> = memo(({ selection, edges, variables, onClearSelection }) => {
+}> = memo(({ selection, variables, onClearSelection }) => {
   if (!selection) {
     return null;
   }

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";
@@ -49,13 +49,18 @@ const LazyVegaLite = React.lazy(() =>
 type MimeBundle = Record<OutputMessage["mimetype"], { [key: string]: unknown }>;
 type MimeBundleOrTuple = MimeBundle | [MimeBundle, { [key: string]: unknown }];
 
+export type OnRefactorWithAI = (opts: {
+  prompt: string;
+  triggerImmediately: boolean;
+}) => void;
+
 /**
  * Renders an output based on an OutputMessage.
  */
 export const OutputRenderer: React.FC<{
   message: Pick<OutputMessage, "channel" | "data" | "mimetype">;
   cellId?: CellId;
-  onRefactorWithAI?: (opts: { prompt: string }) => void;
+  onRefactorWithAI?: OnRefactorWithAI;
   wrapText?: boolean;
 }> = memo((props) => {
   const { message, onRefactorWithAI, cellId, wrapText } = props;
@@ -88,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(
@@ -119,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":
@@ -159,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":
@@ -373,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(() => {
@@ -398,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) => {