@marimo-team/islands 0.16.4 → 0.17.0

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

@@ -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;

package/src/components/editor/__tests__/data-attributes.test.tsx

@@ -12,7 +12,7 @@ import type { UserConfig } from "@/core/config/config-schema";
 import type { OutputMessage } from "@/core/kernel/messages";
 import type { AppMode } from "@/core/mode";
 import { requestClientAtom } from "@/core/network/requests";
-import { Cell } from "../Cell";
+import { Cell } from "../notebook-cell";
 import { OutputArea } from "../Output";
 
 function createTestWrapper() {

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

@@ -102,7 +102,7 @@ export function useNotebookActions() {
     updateCellConfig,
     undoDeleteCell,
     clearAllCellOutputs,
-    upsertSetupCell,
+    addSetupCellIfDoesntExist,
     collapseAllCells,
     expandAllCells,
   } = useCellActions();
@@ -401,9 +401,7 @@ export function useNotebookActions() {
       icon: <DiamondPlusIcon size={14} strokeWidth={1.5} />,
       label: "Add setup cell",
       handle: () => {
-        upsertSetupCell({
-          code: "# Initialization code that runs before all other cells",
-        });
+        addSetupCellIfDoesntExist({});
       },
     },
     {

package/src/components/editor/ai/__tests__/completion-utils.test.ts

@@ -55,15 +55,13 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="dataset1" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>
-
-      <data name="dataset2" source="unknown">
-      Columns:
-        - col3: boolean
-        - col4: date</data>",
+          "plainText": "<data name="dataset1" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>
+
+      <data name="dataset2" source="unknown">Columns:
+        col3 (boolean)
+        col4 (date)</data>",
           "schema": [],
           "variables": [],
         },
@@ -108,10 +106,9 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="existingDataset" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>",
+          "plainText": "<data name="existingDataset" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>",
           "schema": [],
           "variables": [],
         },
@@ -144,14 +141,12 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="dataset.with.dots" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>
-
-      <data name="regular_dataset" source="unknown">
-      Columns:
-        - col3: boolean</data>",
+          "plainText": "<data name="dataset.with.dots" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>
+
+      <data name="regular_dataset" source="unknown">Columns:
+        col3 (boolean)</data>",
           "schema": [],
           "variables": [],
         },
@@ -198,9 +193,8 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="table1" source="unknown">
-      Columns:
-        - col1: number</data>",
+          "plainText": "<data name="table1" source="unknown">Columns:
+        col1 (number)</data>",
           "schema": [],
           "variables": [],
         },
@@ -276,10 +270,9 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="dataset1" source="unknown">
-      Columns:
-        - col1: number
-        - col2: string</data>
+          "plainText": "<data name="dataset1" source="unknown">Columns:
+        col1 (number)
+        col2 (string)</data>
 
       <variable name="var1" dataType="string">"string value"</variable>",
           "schema": [],
@@ -346,9 +339,8 @@ describe("getAICompletionBody", () => {
     expect(result).toMatchInlineSnapshot(`
       {
         "context": {
-          "plainText": "<data name="conflict" source="unknown">
-      Columns:
-        - col1: number</data>",
+          "plainText": "<data name="conflict" source="unknown">Columns:
+        col1 (number)</data>",
           "schema": [],
           "variables": [],
         },

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -3,7 +3,7 @@
 import { useCompletion } from "@ai-sdk/react";
 import { EditorView } from "@codemirror/view";
 import { AtSignIcon, Loader2Icon, SparklesIcon, XIcon } from "lucide-react";
-import React, { useEffect, useId, useState } from "react";
+import React, { useCallback, useEffect, useId, useState } from "react";
 import CodeMirrorMerge from "react-codemirror-merge";
 import { Button } from "@/components/ui/button";
 import { customPythonLanguageSupport } from "@/core/codemirror/language/languages/python";
@@ -15,6 +15,7 @@ import { useAtom } from "jotai";
 import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
 import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
+import { Switch } from "@/components/ui/switch";
 import { Tooltip } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
 import { includeOtherCellsAtom } from "@/core/ai/state";
@@ -28,8 +29,10 @@ import { prettyError } from "@/utils/errors";
 import { retryWithTimeout } from "@/utils/timeout";
 import { PromptInput } from "./add-cell-with-ai";
 import {
+  AcceptCompletionButton,
   CompletionActions,
   createAiCompletionOnKeydown,
+  RejectCompletionButton,
 } from "./completion-handlers";
 import { addContextCompletion, getAICompletionBody } from "./completion-utils";
 
@@ -45,6 +48,9 @@ interface Props {
   declineChange: () => void;
   acceptChange: (rightHandCode: string) => void;
   enabled: boolean;
+  triggerImmediately?: boolean;
+  runCell: () => void;
+  outputArea?: "above" | "below";
   /**
    * Children shown when there is no completion
    */
@@ -67,8 +73,12 @@ export const AiCompletionEditor: React.FC<Props> = ({
   declineChange,
   acceptChange,
   enabled,
+  triggerImmediately,
+  runCell,
+  outputArea,
   children,
 }) => {
+  const [showInputPrompt, setShowInputPrompt] = useState(false);
   const [completionBody, setCompletionBody] = useState<object>({});
 
   const [includeOtherCells, setIncludeOtherCells] = useAtom(
@@ -86,6 +96,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
     setCompletion,
     setInput,
     handleSubmit,
+    complete,
   } = useCompletion({
     api: runtimeManager.getAiURL("completion").toString(),
     headers: runtimeManager.headers(),
@@ -94,7 +105,11 @@ export const AiCompletionEditor: React.FC<Props> = ({
     // Throttle the messages and data updates to 100ms
     experimental_throttle: 100,
     body: {
-      ...completionBody,
+      ...(Object.keys(completionBody).length > 0
+        ? completionBody
+        : initialPrompt
+          ? getAICompletionBody({ input: initialPrompt })
+          : {}),
       includeOtherCode: includeOtherCells ? getCodes(currentCode) : "",
       code: currentCode,
       language: currentLanguageAdapter,
@@ -114,6 +129,14 @@ export const AiCompletionEditor: React.FC<Props> = ({
   const inputRef = React.useRef<ReactCodeMirrorRef>(null);
   const completion = untrimmedCompletion.trimEnd();
 
+  const initialSubmit = useCallback(() => {
+    if (triggerImmediately && !isLoading && initialPrompt) {
+      // Use complete to pass the prompt directly, else input might be empty
+      complete(initialPrompt);
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [triggerImmediately]);
+
   // Focus the input
   useEffect(() => {
     if (enabled) {
@@ -122,6 +145,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
           const input = inputRef.current;
           if (input?.view) {
             input.view.focus();
+            initialSubmit();
             return true;
           }
           return false;
@@ -131,7 +155,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
 
       selectAllText(inputRef.current?.view);
     }
-  }, [enabled]);
+  }, [enabled, initialSubmit]);
 
   // Reset the input when the prompt changes
   useEffect(() => {
@@ -148,16 +172,43 @@ export const AiCompletionEditor: React.FC<Props> = ({
   };
 
   const handleDeclineCompletion = () => {
+    declineChange();
     setCompletion("");
   };
 
+  const showCompletionBanner =
+    enabled && triggerImmediately && (completion || isLoading);
+
+  const showInput = enabled && (!triggerImmediately || showInputPrompt);
+
+  const completionBanner = (
+    <div
+      className={cn(
+        "w-full bg-(--cm-background) flex justify-center transition-all duration-300 ease-in-out overflow-hidden",
+        showCompletionBanner
+          ? "max-h-20 opacity-100 translate-y-0"
+          : "max-h-0 opacity-0 -translate-y-2",
+      )}
+    >
+      <CompletionBanner
+        status={isLoading ? "loading" : "generated"}
+        onAccept={handleAcceptCompletion}
+        onReject={handleDeclineCompletion}
+        showInputPrompt={showInputPrompt}
+        setShowInputPrompt={setShowInputPrompt}
+        runCell={runCell}
+        className="mt-4 mb-3 w-128"
+      />
+    </div>
+  );
+
   return (
     <div className={cn("flex flex-col w-full rounded-[inherit]", className)}>
       <div
         className={cn(
           "flex items-center gap-2 border-b px-3 transition-all rounded-[inherit] rounded-b-none duration-300",
-          enabled && "max-h-[400px] min-h-11 visible",
-          !enabled && "max-h-0 min-h-0 invisible",
+          showInput && "max-h-[400px] min-h-11 visible",
+          !showInput && "max-h-0 min-h-0 invisible",
         )}
       >
         {enabled && (
@@ -266,6 +317,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
           </>
         )}
       </div>
+      {outputArea === "above" && completionBanner}
       {completion && enabled && (
         <CodeMirrorMerge className="cm" theme={theme}>
           <Original
@@ -282,6 +334,84 @@ export const AiCompletionEditor: React.FC<Props> = ({
         </CodeMirrorMerge>
       )}
       {(!completion || !enabled) && children}
+      {/* By default, show the completion banner below the code */}
+      {(outputArea === "below" || !outputArea) && completionBanner}
+    </div>
+  );
+};
+
+interface CompletionBannerProps {
+  status: "loading" | "generated";
+  onAccept: () => void;
+  onReject: () => void;
+  showInputPrompt: boolean;
+  setShowInputPrompt: (show: boolean) => void;
+  runCell: () => void;
+  className?: string;
+}
+
+const CompletionBanner: React.FC<CompletionBannerProps> = ({
+  status,
+  onAccept,
+  onReject,
+  className,
+  showInputPrompt,
+  setShowInputPrompt,
+  runCell,
+}) => {
+  const isLoading = status === "loading";
+
+  return (
+    <div
+      className={cn(
+        "flex flex-row items-center gap-6 rounded-md py-2 px-2.5 text-sm border border-border",
+        "shadow-[0_0_6px_1px_rgba(34,197,94,0.15)]",
+        className,
+      )}
+    >
+      <div className="flex flex-row items-center gap-2">
+        <div
+          className={cn(
+            "w-2 h-2 rounded-full",
+            status === "loading" ? "bg-blue-500 animate-pulse" : "bg-green-500",
+          )}
+        />
+        <p className="transition-opacity duration-200 text-muted-foreground">
+          {isLoading ? "Generating fix..." : "Showing fix"}
+        </p>
+      </div>
+
+      <div className="flex flex-row items-center gap-1">
+        <Label
+          htmlFor="show-input-prompt"
+          className="text-muted-foreground text-xs whitespace-nowrap ellipsis"
+        >
+          Show prompt
+        </Label>
+        <Switch
+          checked={showInputPrompt}
+          onCheckedChange={setShowInputPrompt}
+          size="xs"
+        />
+      </div>
+
+      <div className="flex flex-row items-center gap-2 ml-auto">
+        <AcceptCompletionButton
+          isLoading={isLoading}
+          onAccept={onAccept}
+          size="xs"
+          buttonStyles="border-none rounded-md rounded-r-none"
+          playButtonStyles="border-0 border-l-1 rounded-md rounded-l-none"
+          runCell={runCell}
+          // acceptShortcut="Mod-↵"
+        />
+        <RejectCompletionButton
+          onDecline={onReject}
+          size="xs"
+          className="border-none rounded-md"
+          // declineShortcut="Shift-Mod-Delete"
+        />
+      </div>
     </div>
   );
 };

package/src/components/editor/ai/completion-handlers.tsx

@@ -1,8 +1,10 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { PlayIcon } from "lucide-react";
 import React from "react";
 import { MinimalHotkeys } from "@/components/shortcuts/renderShortcut";
 import { Button } from "@/components/ui/button";
+import { Tooltip } from "@/components/ui/tooltip";
 import { isPlatformMac } from "@/core/hotkeys/shortcuts";
 
 /**
@@ -99,29 +101,111 @@ export const CompletionActionsCellFooter: React.FC<{
 }> = ({ isLoading, onAccept, onDecline }) => {
   return (
     <>
-      <Button
-        variant="text"
-        size="xs"
-        disabled={isLoading}
-        onClick={onAccept}
-        className="h-6 text-(--grass-11) border-(--grass-7) bg-(--grass-3)/60
-        hover:bg-(--grass-3) dark:bg-(--grass-4)/80 dark:hover:bg-(--grass-3) rounded px-3 font-semibold
-        active:bg-(--grass-5) dark:active:bg-(--grass-4)
-        border-(--green-6) border hover:shadow-xs"
-      >
-        Accept
-      </Button>
-      <Button
-        variant="text"
+      <AcceptCompletionButton
+        isLoading={isLoading}
+        onAccept={onAccept}
         size="xs"
-        onClick={onDecline}
-        className="h-6 text-(--red-10) border-(--red-7) bg-(--red-3)/60 hover:bg-(--red-3)
-        dark:bg-(--red-4)/80 dark:hover:bg-(--red-3) rounded px-3 font-semibold
-        active:bg-(--red-5) dark:active:bg-(--red-4)
-        border-(--red-6) border hover:shadow-xs"
-      >
-        Reject
-      </Button>
+      />
+      <RejectCompletionButton onDecline={onDecline} size="xs" />
     </>
   );
 };
+
+export const AcceptCompletionButton: React.FC<{
+  isLoading: boolean;
+  onAccept: () => void;
+  size?: "xs" | "sm";
+  buttonStyles?: string;
+  playButtonStyles?: string;
+  acceptShortcut?: string;
+  runCell?: () => void;
+}> = ({
+  isLoading,
+  onAccept,
+  size = "sm",
+  buttonStyles,
+  acceptShortcut,
+  runCell,
+  playButtonStyles,
+}) => {
+  const handleAcceptAndRun = () => {
+    onAccept();
+    if (runCell) {
+      runCell();
+    }
+  };
+
+  const baseClasses = `h-6 text-(--grass-11) bg-(--grass-3)/60
+    hover:bg-(--grass-3) dark:bg-(--grass-4)/80 dark:hover:bg-(--grass-3) font-semibold
+    active:bg-(--grass-5) dark:active:bg-(--grass-4)
+    border-(--green-6) border hover:shadow-xs`;
+
+  if (runCell) {
+    return (
+      <div className="flex">
+        <Button
+          variant="text"
+          size={size}
+          disabled={isLoading}
+          onClick={onAccept}
+          className={`${baseClasses} rounded-r-none ${buttonStyles}`}
+        >
+          Accept
+          {acceptShortcut && (
+            <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
+          )}
+        </Button>
+        <Tooltip content="Accept and run cell">
+          <Button
+            variant="text"
+            size={size}
+            disabled={isLoading}
+            onClick={handleAcceptAndRun}
+            className={`${baseClasses} rounded-l-none px-1.5 ${playButtonStyles}`}
+          >
+            <PlayIcon className="h-2.5 w-2.5 mt-0.5" />
+          </Button>
+        </Tooltip>
+      </div>
+    );
+  }
+
+  return (
+    <Button
+      variant="text"
+      size={size}
+      disabled={isLoading}
+      onClick={onAccept}
+      className={`${baseClasses} rounded px-3 ${buttonStyles}`}
+    >
+      Accept
+      {acceptShortcut && (
+        <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
+      )}
+    </Button>
+  );
+};
+
+export const RejectCompletionButton: React.FC<{
+  onDecline: () => void;
+  size?: "xs" | "sm";
+  className?: string;
+  declineShortcut?: string;
+}> = ({ onDecline, size = "sm", className, declineShortcut }) => {
+  return (
+    <Button
+      variant="text"
+      size={size}
+      onClick={onDecline}
+      className={`h-6 text-(--red-10) bg-(--red-3)/60 hover:bg-(--red-3)
+    dark:bg-(--red-4)/80 dark:hover:bg-(--red-3) rounded px-3 font-semibold
+    active:bg-(--red-5) dark:active:bg-(--red-4)
+    border-(--red-6) border hover:shadow-xs ${className}`}
+    >
+      Reject
+      {declineShortcut && (
+        <MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
+      )}
+    </Button>
+  );
+};