@marimo-team/islands 0.17.0 → 0.17.2

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

@@ -2,7 +2,13 @@
 
 import { useCompletion } from "@ai-sdk/react";
 import { EditorView } from "@codemirror/view";
-import { AtSignIcon, Loader2Icon, SparklesIcon, XIcon } from "lucide-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";
@@ -11,14 +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";
@@ -40,15 +48,14 @@ 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;
-  triggerImmediately?: boolean;
   runCell: () => void;
   outputArea?: "above" | "below";
   /**
@@ -65,15 +72,14 @@ 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,
-  triggerImmediately,
   runCell,
   outputArea,
   children,
@@ -88,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,
@@ -178,6 +198,8 @@ export const AiCompletionEditor: React.FC<Props> = ({
 
   const showCompletionBanner =
     enabled && triggerImmediately && (completion || isLoading);
+  // Set default output area to below if not specified
+  outputArea = outputArea ?? "below";
 
   const showInput = enabled && (!triggerImmediately || showInputPrompt);
 
@@ -202,6 +224,35 @@ export const AiCompletionEditor: React.FC<Props> = ({
     </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
@@ -318,24 +369,10 @@ export const AiCompletionEditor: React.FC<Props> = ({
         )}
       </div>
       {outputArea === "above" && completionBanner}
-      {completion && enabled && (
-        <CodeMirrorMerge className="cm" theme={theme}>
-          <Original
-            onChange={onChange}
-            value={currentCode}
-            extensions={baseExtensions}
-          />
-          <Modified
-            value={completion}
-            editable={false}
-            readOnly={true}
-            extensions={baseExtensions}
-          />
-        </CodeMirrorMerge>
-      )}
-      {(!completion || !enabled) && children}
+      {renderCompletionEditor()}
+      {(!completion || !enabled) && !previousCellCode && children}
       {/* By default, show the completion banner below the code */}
-      {(outputArea === "below" || !outputArea) && completionBanner}
+      {outputArea === "below" && completionBanner}
     </div>
   );
 };
@@ -370,12 +407,22 @@ const CompletionBanner: React.FC<CompletionBannerProps> = ({
       )}
     >
       <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",
-          )}
-        />
+        {isLoading ? (
+          <Loader2Icon
+            className="animate-spin text-blue-600 mb-[1px]"
+            size={15}
+            strokeWidth={2}
+            aria-label="Generating fix"
+          />
+        ) : (
+          <CircleCheckIcon
+            className="text-green-600 mb-[1px]"
+            size={15}
+            strokeWidth={2}
+            aria-label="Fix generated"
+          />
+        )}
+
         <p className="transition-opacity duration-200 text-muted-foreground">
           {isLoading ? "Generating fix..." : "Showing fix"}
         </p>

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

@@ -98,13 +98,15 @@ export const CompletionActionsCellFooter: React.FC<{
   onDecline: () => void;
   size?: "xs" | "sm";
   multipleCompletions?: boolean;
-}> = ({ isLoading, onAccept, onDecline }) => {
+  runCell?: () => void;
+}> = ({ isLoading, onAccept, onDecline, runCell }) => {
   return (
     <>
       <AcceptCompletionButton
         isLoading={isLoading}
         onAccept={onAccept}
         size="xs"
+        runCell={runCell}
       />
       <RejectCompletionButton onDecline={onDecline} size="xs" />
     </>
@@ -114,6 +116,7 @@ export const CompletionActionsCellFooter: React.FC<{
 export const AcceptCompletionButton: React.FC<{
   isLoading: boolean;
   onAccept: () => void;
+  multipleCompletions?: boolean;
   size?: "xs" | "sm";
   buttonStyles?: string;
   playButtonStyles?: string;
@@ -122,6 +125,7 @@ export const AcceptCompletionButton: React.FC<{
 }> = ({
   isLoading,
   onAccept,
+  multipleCompletions = false,
   size = "sm",
   buttonStyles,
   acceptShortcut,
@@ -135,6 +139,8 @@ export const AcceptCompletionButton: React.FC<{
     }
   };
 
+  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)
@@ -150,12 +156,18 @@ export const AcceptCompletionButton: React.FC<{
           onClick={onAccept}
           className={`${baseClasses} rounded-r-none ${buttonStyles}`}
         >
-          Accept
+          {text}
           {acceptShortcut && (
             <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
           )}
         </Button>
-        <Tooltip content="Accept and run cell">
+        <Tooltip
+          content={
+            multipleCompletions
+              ? "Accept and run all cells"
+              : "Accept and run cell"
+          }
+        >
           <Button
             variant="text"
             size={size}
@@ -163,7 +175,7 @@ export const AcceptCompletionButton: React.FC<{
             onClick={handleAcceptAndRun}
             className={`${baseClasses} rounded-l-none px-1.5 ${playButtonStyles}`}
           >
-            <PlayIcon className="h-2.5 w-2.5 mt-0.5" />
+            <PlayIcon className="h-2.5 w-2.5" />
           </Button>
         </Tooltip>
       </div>
@@ -178,7 +190,7 @@ export const AcceptCompletionButton: React.FC<{
       onClick={onAccept}
       className={`${baseClasses} rounded px-3 ${buttonStyles}`}
     >
-      Accept
+      {text}
       {acceptShortcut && (
         <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
       )}
@@ -188,10 +200,17 @@ export const AcceptCompletionButton: React.FC<{
 
 export const RejectCompletionButton: React.FC<{
   onDecline: () => void;
+  multipleCompletions?: boolean;
   size?: "xs" | "sm";
   className?: string;
   declineShortcut?: string;
-}> = ({ onDecline, size = "sm", className, declineShortcut }) => {
+}> = ({
+  onDecline,
+  multipleCompletions = false,
+  size = "sm",
+  className,
+  declineShortcut,
+}) => {
   return (
     <Button
       variant="text"
@@ -202,7 +221,7 @@ export const RejectCompletionButton: React.FC<{
     active:bg-(--red-5) dark:active:bg-(--red-4)
     border-(--red-6) border hover:shadow-xs ${className}`}
     >
-      Reject
+      Reject{multipleCompletions && " all"}
       {declineShortcut && (
         <MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
       )}

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

@@ -407,9 +407,8 @@ const CellEditorInternal = ({
 
   return (
     <AiCompletionEditor
-      enabled={aiCompletionCell?.cellId === cellId}
-      initialPrompt={aiCompletionCell?.initialPrompt}
-      triggerImmediately={aiCompletionCell?.triggerImmediately}
+      cellId={cellId}
+      aiCompletionCell={aiCompletionCell}
       currentCode={editorViewRef.current?.state.doc.toString() ?? code}
       currentLanguageAdapter={languageAdapter}
       declineChange={useEvent(() => {

package/src/components/editor/cell/useRunCells.ts

@@ -1,15 +1,23 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { closeCompletion } from "@codemirror/autocomplete";
+import { atom, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
-import { getNotebook, useCellActions } from "@/core/cells/cells";
+import {
+  getNotebook,
+  type NotebookState,
+  useCellActions,
+} from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { enabledCellIds, staleCellIds } from "@/core/cells/utils";
 import { getCurrentLanguageAdapter } from "@/core/codemirror/language/commands";
 import { getEditorCodeAsPython } from "@/core/codemirror/language/utils";
 import { useRequestClient } from "@/core/network/requests";
+import type { RunRequest } from "@/core/network/types";
 import { Logger } from "@/utils/Logger";
 
+export const hasRunAnyCellAtom = atom<boolean>(false);
+
 /**
  * Creates a function that runs all cells that have been edited or interrupted.
  */
@@ -45,42 +53,66 @@ export function useRunAllCells() {
 export function useRunCells() {
   const { prepareForRun } = useCellActions();
   const { sendRun } = useRequestClient();
+  const setHasRunAnyCell = useSetAtom(hasRunAnyCellAtom);
 
-  const runCells = useEvent(async (cellIds: CellId[]) => {
+  const runCellsMemoized = useEvent(async (cellIds: CellId[]) => {
     if (cellIds.length === 0) {
       return;
     }
 
-    const { cellHandles, cellData } = getNotebook();
-
-    const codes: string[] = [];
-    for (const cellId of cellIds) {
-      const ref = cellHandles[cellId];
-      const ev = ref?.current?.editorView;
-      let code: string;
-      // Performs side-effects that must run whenever the cell is run, but doesn't
-      // actually run the cell.
-      if (ev) {
-        // Skip close on markdown, since we autorun, otherwise we'll close the
-        // completion each time.
-        if (getCurrentLanguageAdapter(ev) !== "markdown") {
-          closeCompletion(ev);
-        }
-        // Prefer code from editor
-        code = getEditorCodeAsPython(ev);
-      } else {
-        code = cellData[cellId]?.code || "";
-      }
+    const notebook = getNotebook();
 
-      codes.push(code);
-      prepareForRun({ cellId });
-    }
+    // Set a flag that a user has manually run at least one cell.
+    setHasRunAnyCell(true);
 
-    // Send the run request to the Kernel
-    await sendRun({ cellIds: cellIds, codes: codes }).catch((error) => {
-      Logger.error(error);
-    });
+    return runCells({ cellIds, sendRun, prepareForRun, notebook });
   });
 
-  return runCells;
+  return runCellsMemoized;
+}
+
+export async function runCells({
+  cellIds,
+  sendRun,
+  prepareForRun,
+  notebook,
+}: {
+  cellIds: CellId[];
+  sendRun: (request: RunRequest) => Promise<null>;
+  prepareForRun: (action: { cellId: CellId }) => void;
+  notebook: NotebookState;
+}) {
+  if (cellIds.length === 0) {
+    return;
+  }
+
+  const { cellHandles, cellData } = notebook;
+
+  const codes: string[] = [];
+  for (const cellId of cellIds) {
+    const ref = cellHandles[cellId];
+    const ev = ref?.current?.editorView;
+    let code: string;
+    // Performs side-effects that must run whenever the cell is run, but doesn't
+    // actually run the cell.
+    if (ev) {
+      // Skip close on markdown, since we autorun, otherwise we'll close the
+      // completion each time.
+      if (getCurrentLanguageAdapter(ev) !== "markdown") {
+        closeCompletion(ev);
+      }
+      // Prefer code from editor
+      code = getEditorCodeAsPython(ev);
+    } else {
+      code = cellData[cellId]?.code || "";
+    }
+
+    codes.push(code);
+    prepareForRun({ cellId });
+  }
+
+  // Send the run request to the Kernel
+  await sendRun({ cellIds: cellIds, codes: codes }).catch((error) => {
+    Logger.error(error);
+  });
 }

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

@@ -38,11 +38,11 @@ const CachePanel = () => {
       });
       // Request updated cache info after purge
       refetch();
-    } catch (err) {
+    } catch (error) {
       toast({
         title: "Error",
         description:
-          err instanceof Error ? err.message : "Failed to purge cache",
+          error instanceof Error ? error.message : "Failed to purge cache",
         variant: "danger",
       });
     } finally {

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

@@ -211,7 +211,7 @@ const InstallPackageForm: React.FC<{
           ) : (
             <Tooltip content="Change package manager">
               <BoxIcon
-                onClick={() => openSettings("packageManagement")}
+                onClick={() => openSettings("packageManagementAndData")}
                 className="mr-2 h-4 w-4 shrink-0 opacity-50 hover:opacity-80 cursor-pointer"
               />
             </Tooltip>

package/src/components/editor/chrome/state.ts

@@ -3,7 +3,7 @@
 import { useAtomValue } from "jotai";
 import { z } from "zod";
 import { createReducerAndAtoms } from "@/utils/createReducer";
-import { ZodLocalStorage } from "@/utils/localStorage";
+import { ZodLocalStorage } from "@/utils/storage/typed";
 import type { PanelType } from "./types";
 
 export interface ChromeState {

package/src/components/editor/chrome/wrapper/app-chrome.tsx

@@ -25,6 +25,7 @@ import { ContextAwarePanel } from "../panels/context-aware-panel/context-aware-p
 import { useChromeActions, useChromeState } from "../state";
 import { Minimap } from "./minimap";
 import { PanelsWrapper } from "./panels";
+import { PendingAICells } from "./pending-ai-cells";
 import { createStorage } from "./storage";
 import { handleDragging } from "./utils";
 
@@ -272,6 +273,7 @@ export const AppChrome: React.FC<PropsWithChildren> = ({ children }) => {
         <ContextAwarePanel />
       </PanelGroup>
       <Minimap />
+      <PendingAICells />
       <ErrorBoundary>
         <TooltipProvider>
           <Footer />

package/src/components/editor/chrome/wrapper/minimap.tsx

@@ -64,6 +64,7 @@ const MinimapCell: React.FC<MinimapCellProps> = (props) => {
 
   return (
     <button
+      type="button"
       data-node-id={cell.id}
       className={cn(
         "group bg-transparent text-left w-full flex relative justify-between items-center",
@@ -107,6 +108,7 @@ const MinimapCell: React.FC<MinimapCellProps> = (props) => {
         width="1"
         height="1"
       >
+        <title>Cell dependency connections for cell {cell.id}</title>
         {isSelected ? (
           <SelectedCell
             cell={cell}