edit notebook tool

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

@@ -36,7 +36,10 @@ import {
   type ChatId,
   chatStateAtom,
 } from "@/core/ai/state";
-import { FRONTEND_TOOL_REGISTRY } from "@/core/ai/tools/registry";
+import {
+  type CopilotMode,
+  FRONTEND_TOOL_REGISTRY,
+} from "@/core/ai/tools/registry";
 import { aiAtom, aiEnabledAtom } from "@/core/config/config";
 import { DEFAULT_AI_MODEL } from "@/core/config/config-schema";
 import { FeatureFlagged } from "@/core/config/feature-flag";
@@ -301,7 +304,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 +320,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 =

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

@@ -11,14 +11,16 @@
 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 +42,14 @@
 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 +66,14 @@
  * 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 +88,20 @@
 
   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,
@@ -134,7 +148,7 @@
       // 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
@@ -178,6 +192,8 @@
 
   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 +218,35 @@
     </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 +363,10 @@
         )}
       </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>
   );
 };

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

@@ -114,6 +114,7 @@ export const CompletionActionsCellFooter: React.FC<{
 export const AcceptCompletionButton: React.FC<{
   isLoading: boolean;
   onAccept: () => void;
+  text?: string;
   size?: "xs" | "sm";
   buttonStyles?: string;
   playButtonStyles?: string;
@@ -122,6 +123,7 @@ export const AcceptCompletionButton: React.FC<{
 }> = ({
   isLoading,
   onAccept,
+  text = "Accept",
   size = "sm",
   buttonStyles,
   acceptShortcut,
@@ -150,7 +152,7 @@ export const AcceptCompletionButton: React.FC<{
           onClick={onAccept}
           className={`${baseClasses} rounded-r-none ${buttonStyles}`}
         >
-          Accept
+          {text}
           {acceptShortcut && (
             <MinimalHotkeys className="ml-1 inline" shortcut={acceptShortcut} />
           )}
@@ -178,7 +180,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 +190,17 @@ export const AcceptCompletionButton: React.FC<{
 
 export const RejectCompletionButton: React.FC<{
   onDecline: () => void;
+  text?: string;
   size?: "xs" | "sm";
   className?: string;
   declineShortcut?: string;
-}> = ({ onDecline, size = "sm", className, declineShortcut }) => {
+}> = ({
+  onDecline,
+  text = "Reject",
+  size = "sm",
+  className,
+  declineShortcut,
+}) => {
   return (
     <Button
       variant="text"
@@ -202,7 +211,7 @@ export const RejectCompletionButton: React.FC<{
     active:bg-(--red-5) dark:active:bg-(--red-4)
     border-(--red-6) border hover:shadow-xs ${className}`}
     >
-      Reject
+      {text}
       {declineShortcut && (
         <MinimalHotkeys className="ml-1 inline" shortcut={declineShortcut} />
       )}

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

@@ -1,51 +1,117 @@
 /* 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";
 
 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 }> = ({
   cellId,
 }) => {
   const store = useStore();
   const stagedAICells = useAtomValue(stagedAICellsAtom);
+  const stagedAiCell = stagedAICells.get(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"
       />
     </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;
+  }
+}

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

@@ -230,7 +230,7 @@
     );
 
     return extensions;
   }, [
     cellId,
     userConfig.keymap,
     userConfig.completion,
@@ -383,7 +383,7 @@
   ]);
 
   // Destroy the editor when the component is unmounted
   useEffect(() => {
     const ev = editorViewRef.current;
     return () => {
       ev?.destroy();
@@ -407,9 +407,8 @@
 
   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(() => {