@marimo-team/islands 0.16.5 → 0.17.1

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

@@ -2,8 +2,14 @@
 
 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 {
+  AtSignIcon,
+  CircleCheckIcon,
+  Loader2Icon,
+  SparklesIcon,
+  XIcon,
+} from "lucide-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";
@@ -11,13 +17,16 @@ import { customPythonLanguageSupport } from "@/core/codemirror/language/language
 import "./merge-editor.css";
 import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
-import { useAtom } from "jotai";
+import { useAtom, useAtomValue } 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";
+import { stagedAICellsAtom } from "@/core/ai/staged-cells";
+import { type AiCompletionCell, includeOtherCellsAtom } from "@/core/ai/state";
+import type { CellId } from "@/core/cells/ids";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import type { LanguageAdapterType } from "@/core/codemirror/language/types";
 import { selectAllText } from "@/core/codemirror/utils";
@@ -28,8 +37,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";
 
@@ -37,14 +48,16 @@ const Original = CodeMirrorMerge.Original;
 const Modified = CodeMirrorMerge.Modified;
 
 interface Props {
+  cellId: CellId;
+  aiCompletionCell: AiCompletionCell | null;
   className?: string;
   currentCode: string;
   currentLanguageAdapter: LanguageAdapterType | undefined;
-  initialPrompt: string | undefined;
   onChange: (code: string) => void;
   declineChange: () => void;
   acceptChange: (rightHandCode: string) => void;
-  enabled: boolean;
+  runCell: () => void;
+  outputArea?: "above" | "below";
   /**
    * Children shown when there is no completion
    */
@@ -59,16 +72,19 @@ const baseExtensions = [customPythonLanguageSupport(), EditorView.lineWrapping];
  * This shows a left/right split with the original and modified code.
  */
 export const AiCompletionEditor: React.FC<Props> = ({
+  cellId,
+  aiCompletionCell,
   className,
   onChange,
-  initialPrompt,
   currentLanguageAdapter,
   currentCode,
   declineChange,
   acceptChange,
-  enabled,
+  runCell,
+  outputArea,
   children,
 }) => {
+  const [showInputPrompt, setShowInputPrompt] = useState(false);
   const [completionBody, setCompletionBody] = useState<object>({});
 
   const [includeOtherCells, setIncludeOtherCells] = useAtom(
@@ -78,6 +94,20 @@ export const AiCompletionEditor: React.FC<Props> = ({
 
   const runtimeManager = useRuntimeManager();
 
+  const {
+    initialPrompt,
+    triggerImmediately,
+    cellId: aiCellId,
+  } = aiCompletionCell ?? {};
+  const enabled = aiCellId === cellId;
+
+  const stagedAICells = useAtomValue(stagedAICellsAtom);
+  const updatedCell = stagedAICells.get(cellId);
+  let previousCellCode: string | undefined;
+  if (updatedCell?.type === "update_cell") {
+    previousCellCode = updatedCell.previousCode;
+  }
+
   const {
     completion: untrimmedCompletion,
     input,
@@ -86,6 +116,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
     setCompletion,
     setInput,
     handleSubmit,
+    complete,
   } = useCompletion({
     api: runtimeManager.getAiURL("completion").toString(),
     headers: runtimeManager.headers(),
@@ -94,7 +125,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 +149,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 +165,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
           const input = inputRef.current;
           if (input?.view) {
             input.view.focus();
+            initialSubmit();
             return true;
           }
           return false;
@@ -131,7 +175,7 @@ export const AiCompletionEditor: React.FC<Props> = ({
 
       selectAllText(inputRef.current?.view);
     }
-  }, [enabled]);
+  }, [enabled, initialSubmit]);
 
   // Reset the input when the prompt changes
   useEffect(() => {
@@ -148,16 +192,74 @@ export const AiCompletionEditor: React.FC<Props> = ({
   };
 
   const handleDeclineCompletion = () => {
+    declineChange();
     setCompletion("");
   };
 
+  const showCompletionBanner =
+    enabled && triggerImmediately && (completion || isLoading);
+  // Set default output area to below if not specified
+  outputArea = outputArea ?? "below";
+
+  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>
+  );
+
+  const renderMergeEditor = (originalCode: string, modifiedCode: string) => {
+    return (
+      <CodeMirrorMerge className="cm" theme={theme}>
+        <Original
+          onChange={onChange}
+          value={originalCode}
+          extensions={baseExtensions}
+        />
+        <Modified
+          value={modifiedCode}
+          editable={false}
+          readOnly={true}
+          extensions={baseExtensions}
+        />
+      </CodeMirrorMerge>
+    );
+  };
+
+  const renderCompletionEditor = () => {
+    if (completion && enabled) {
+      return renderMergeEditor(currentCode, completion);
+    }
+    // If there is no completion and there is previous cell code, it means there is an AI change to the cell.
+    // And we want to render the previous cell code as the original
+    if (!completion && previousCellCode) {
+      return renderMergeEditor(previousCellCode, currentCode);
+    }
+  };
+
   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,22 +368,97 @@ export const AiCompletionEditor: React.FC<Props> = ({
           </>
         )}
       </div>
-      {completion && enabled && (
-        <CodeMirrorMerge className="cm" theme={theme}>
-          <Original
-            onChange={onChange}
-            value={currentCode}
-            extensions={baseExtensions}
+      {outputArea === "above" && completionBanner}
+      {renderCompletionEditor()}
+      {(!completion || !enabled) && !previousCellCode && children}
+      {/* By default, show the completion banner below the code */}
+      {outputArea === "below" && 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">
+        {isLoading ? (
+          <Loader2Icon
+            className="animate-spin text-blue-600 mb-[1px]"
+            size={15}
+            strokeWidth={2}
+            aria-label="Generating fix"
           />
-          <Modified
-            value={completion}
-            editable={false}
-            readOnly={true}
-            extensions={baseExtensions}
+        ) : (
+          <CircleCheckIcon
+            className="text-green-600 mb-[1px]"
+            size={15}
+            strokeWidth={2}
+            aria-label="Fix generated"
           />
-        </CodeMirrorMerge>
-      )}
-      {(!completion || !enabled) && children}
+        )}
+
+        <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";
 
 /**
@@ -96,32 +98,133 @@ export const CompletionActionsCellFooter: React.FC<{
   onDecline: () => void;
   size?: "xs" | "sm";
   multipleCompletions?: boolean;
-}> = ({ isLoading, onAccept, onDecline }) => {
+  runCell?: () => void;
+}> = ({ isLoading, onAccept, onDecline, runCell }) => {
   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>
+        runCell={runCell}
+      />
+      <RejectCompletionButton onDecline={onDecline} size="xs" />
     </>
   );
 };
+
+export const AcceptCompletionButton: React.FC<{
+  isLoading: boolean;
+  onAccept: () => void;
+  multipleCompletions?: boolean;
+  size?: "xs" | "sm";
+  buttonStyles?: string;
+  playButtonStyles?: string;
+  acceptShortcut?: string;
+  runCell?: () => void;
+}> = ({
+  isLoading,
+  onAccept,
+  multipleCompletions = false,
+  size = "sm",
+  buttonStyles,
+  acceptShortcut,
+  runCell,
+  playButtonStyles,
+}) => {
+  const handleAcceptAndRun = () => {
+    onAccept();
+    if (runCell) {
+      runCell();
+    }
+  };
+
+  const text = multipleCompletions ? "Accept all" : "Accept";
+
+  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}`}
+        >
+          {text}
+          {acceptShortcut && (
+            <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
+          )}
+        </Button>
+        <Tooltip
+          content={
+            multipleCompletions
+              ? "Accept and run all cells"
+              : "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" />
+          </Button>
+        </Tooltip>
+      </div>
+    );
+  }
+
+  return (
+    <Button
+      variant="text"
+      size={size}
+      disabled={isLoading}
+      onClick={onAccept}
+      className={`${baseClasses} rounded px-3 ${buttonStyles}`}
+    >
+      {text}
+      {acceptShortcut && (
+        <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
+      )}
+    </Button>
+  );
+};
+
+export const RejectCompletionButton: React.FC<{
+  onDecline: () => void;
+  multipleCompletions?: boolean;
+  size?: "xs" | "sm";
+  className?: string;
+  declineShortcut?: string;
+}> = ({
+  onDecline,
+  multipleCompletions = false,
+  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{multipleCompletions && " all"}
+      {declineShortcut && (
+        <MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
+      )}
+    </Button>
+  );
+};

package/src/components/editor/cell/StagedAICell.tsx

@@ -1,22 +1,36 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { useAtomValue, useStore } from "jotai";
-import { stagedAICellsAtom, useStagedCells } from "@/core/ai/staged-cells";
+import {
+  type Edit,
+  stagedAICellsAtom,
+  useStagedCells,
+} from "@/core/ai/staged-cells";
+import { getCellEditorView } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
+import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
 import { cn } from "@/utils/cn";
+import { Logger } from "@/utils/Logger";
 import { CompletionActionsCellFooter } from "../ai/completion-handlers";
+import { useRunCell } from "./useRunCells";
 
 export const StagedAICellBackground: React.FC<{
   cellId: CellId;
   className?: string;
 }> = ({ cellId, className }) => {
   const stagedAICells = useAtomValue(stagedAICellsAtom);
+  const stagedCell = stagedAICells.get(cellId);
 
-  if (!stagedAICells.has(cellId)) {
+  if (!stagedCell) {
     return null;
   }
 
-  return <div className={cn("mo-ai-generated-cell", className)} />;
+  const cellClass =
+    stagedCell.type === "delete_cell"
+      ? "mo-ai-deleted-cell"
+      : "mo-ai-generated-cell";
+
+  return <div className={cn(cellClass, className)} />;
 };
 
 export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
@@ -24,28 +38,83 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
 }) => {
   const store = useStore();
   const stagedAICells = useAtomValue(stagedAICellsAtom);
+  const stagedAiCell = stagedAICells.get(cellId);
+  const runCell = useRunCell(cellId);
+
   const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
 
-  if (!stagedAICells.has(cellId)) {
+  if (!stagedAiCell) {
     return null;
   }
 
-  const handleAcceptCompletion = () => {
-    removeStagedCell(cellId);
-  };
-
-  const handleDeclineCompletion = () => {
-    deleteStagedCell(cellId);
+  const handleCompletion = (type: "accept" | "reject") => {
+    const completionFunc =
+      type === "accept" ? acceptStagedCell : rejectStagedCell;
+    completionFunc(cellId, stagedAiCell, removeStagedCell, deleteStagedCell);
   };
 
   return (
     <div className="flex items-center justify-end gap-1.5 w-full pb-1 pt-2">
       <CompletionActionsCellFooter
         isLoading={false}
-        onAccept={handleAcceptCompletion}
-        onDecline={handleDeclineCompletion}
+        onAccept={() => handleCompletion("accept")}
+        onDecline={() => handleCompletion("reject")}
         size="xs"
+        runCell={runCell}
       />
     </div>
   );
 };
+
+/**
+ * Accept a staged cell and apply the changes.
+ */
+export function acceptStagedCell(
+  cellId: CellId,
+  edit: Edit,
+  removeStagedCell: (cellId: CellId) => void,
+  deleteStagedCell: (cellId: CellId) => void,
+): void {
+  switch (edit.type) {
+    case "delete_cell":
+      // For delete cells, the cell is deleted when the completion is accepted
+      deleteStagedCell(cellId);
+      break;
+    default:
+      removeStagedCell(cellId);
+      break;
+  }
+}
+
+/**
+ * Reject a staged cell and revert the changes.
+ */
+export function rejectStagedCell(
+  cellId: CellId,
+  edit: Edit,
+  removeStagedCell: (cellId: CellId) => void,
+  deleteStagedCell: (cellId: CellId) => void,
+): void {
+  switch (edit.type) {
+    case "update_cell": {
+      // Revert cell code
+      const editorView = getCellEditorView(cellId);
+      if (!editorView) {
+        Logger.error("Editor for this cell not found", { cellId });
+        break;
+      }
+
+      updateEditorCodeFromPython(editorView, edit.previousCode);
+      removeStagedCell(cellId);
+      break;
+    }
+    case "add_cell":
+      // Delete the cell since it's newly created
+      deleteStagedCell(cellId);
+      break;
+    case "delete_cell":
+      // Just remove the deletion marker - cell stays in notebook
+      removeStagedCell(cellId);
+      break;
+  }
+}

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

@@ -67,6 +67,7 @@ export interface CellEditorProps
   // DOM node where the editorView will be mounted
   editorViewParentRef?: React.RefObject<HTMLDivElement | null>;
   showHiddenCode: (opts?: { focus?: boolean }) => void;
+  outputArea?: "above" | "below";
 }
 
 const CellEditorInternal = ({
@@ -88,6 +89,7 @@ const CellEditorInternal = ({
   languageAdapter,
   setLanguageAdapter,
   showLanguageToggles = true,
+  outputArea,
 }: CellEditorProps) => {
   const [aiCompletionCell, setAiCompletionCell] = useAtom(aiCompletionCellAtom);
   const deleteCell = useDeleteCellCallback();
@@ -405,8 +407,8 @@ const CellEditorInternal = ({
 
   return (
     <AiCompletionEditor
-      enabled={aiCompletionCell?.cellId === cellId}
-      initialPrompt={aiCompletionCell?.initialPrompt}
+      cellId={cellId}
+      aiCompletionCell={aiCompletionCell}
       currentCode={editorViewRef.current?.state.doc.toString() ?? code}
       currentLanguageAdapter={languageAdapter}
       declineChange={useEvent(() => {
@@ -433,6 +435,8 @@ const CellEditorInternal = ({
         editorViewRef.current?.focus();
         setAiCompletionCell(null);
       })}
+      runCell={handleRunCell}
+      outputArea={outputArea}
     >
       <div className="relative w-full" {...navigationProps}>
         {showHideButton && (