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 === undefined ? "below" : outputArea;
 
   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/cell/StagedAICell.tsx

@@ -2,8 +2,11 @@
 
 import { useAtomValue, useStore } from "jotai";
 import { 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<{
@@ -24,9 +27,11 @@ export const StagedAICellFooter: React.FC<{ 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;
   }
 
@@ -35,7 +40,28 @@ export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
   };
 
   const handleDeclineCompletion = () => {
-    deleteStagedCell(cellId);
+    switch (stagedAiCell.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, stagedAiCell.previousCode);
+        removeStagedCell(cellId);
+        break;
+      }
+      case "add_cell":
+        // Delete the cell since it's newly created
+        deleteStagedCell(cellId);
+        break;
+      case "delete_cell":
+        // TODO: Revert delete
+        removeStagedCell(cellId);
+        break;
+    }
   };
 
   return (

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(() => {