@marimo-team/islands 0.17.0 → 0.17.2

package/src/core/ai/context/providers/cell-output.ts

@@ -3,7 +3,7 @@
 import type { Completion } from "@codemirror/autocomplete";
 import type { FileUIPart } from "ai";
 import { toPng } from "html-to-image";
-import { notebookAtom } from "@/core/cells/cells";
+import { type NotebookState, notebookAtom } from "@/core/cells/cells";
 import { type CellId, CellOutputId } from "@/core/cells/ids";
 import { displayCellName } from "@/core/cells/names";
 import { isOutputEmpty } from "@/core/cells/outputs";
@@ -15,18 +15,33 @@ import { type AIContextItem, AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
 import { Boosts, Sections } from "./common";
 
+export interface BaseOutput {
+  processedContent?: string;
+  imageUrl?: string;
+  shouldDownloadImage?: boolean;
+  output: OutputMessage;
+  outputType: "text" | "media";
+}
+
+export interface CellOutputData {
+  cellId: CellId;
+  cellName: string;
+  cellCode: string;
+  cellOutput?: BaseOutput;
+  consoleOutputs?: BaseOutput[];
+}
+
+// For the context provider
+// Currently, we enforce that cellOutput is present.
+interface CellOutputContextData
+  extends CellOutputData,
+    Record<string, unknown> {
+  cellOutput: BaseOutput;
+}
+
 export interface CellOutputContextItem extends AIContextItem {
   type: "cell-output";
-  data: {
-    cellId: CellId;
-    cellName: string;
-    cellCode: string;
-    output: OutputMessage;
-    outputType: "text" | "media";
-    processedContent?: string;
-    imageUrl?: string;
-    shouldDownloadImage?: boolean;
-  };
+  data: CellOutputContextData;
 }
 
 function isMediaMimetype(
@@ -79,69 +94,23 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
     const notebook = this.store.get(notebookAtom);
     const items: CellOutputContextItem[] = [];
 
-    for (const [cellIndex, cellId] of notebook.cellIds.inOrderIds.entries()) {
-      const cellRuntime = notebook.cellRuntime[cellId];
+    for (const cellId of notebook.cellIds.inOrderIds) {
+      const cellContextData = getCellContextData(cellId, notebook);
+      const cellOutput = cellContextData.cellOutput;
 
-      // Filter to only cells with output
-      if (!cellRuntime?.output || isOutputEmpty(cellRuntime.output)) {
+      // Skip cells with no output
+      if (!cellOutput) {
         continue;
       }
 
-      const cellData = notebook.cellData[cellId];
-      const output = cellRuntime.output;
-      const mimetype = output.mimetype;
-
-      // Determine output type
-      const isMedia = isMediaMimetype(mimetype, String(output.data));
-      const outputType = isMedia ? "media" : "text";
-
-      const cellName = displayCellName(cellData.name, cellIndex);
-
-      let processedContent: string | undefined;
-      let imageUrl: string | undefined;
-      let shouldDownloadImage = false;
-
-      // Process text content
-      if (outputType === "text" && typeof output.data === "string") {
-        processedContent =
-          mimetype === "text/html"
-            ? parseHtmlContent(output.data)
-            : output.data;
-      }
-
-      // Process media content - for now, we'll just note that it's media
-      if (outputType === "media") {
-        if (
-          typeof output.data === "string" &&
-          output.data.startsWith("data:")
-        ) {
-          imageUrl = output.data; // Data URL
-        } else if (
-          typeof output.data === "string" &&
-          output.data.startsWith("http")
-        ) {
-          imageUrl = output.data; // External URL
-        } else {
-          // Download the image from the DOM element
-          shouldDownloadImage = true;
-        }
-      }
+      const mimetype = cellOutput.output.mimetype || "unknown";
 
       items.push({
         uri: this.asURI(cellId),
-        name: cellName,
+        name: cellContextData.cellName,
         type: this.contextType,
-        description: `Cell output (${mimetype || "unknown"})`,
-        data: {
-          cellId,
-          cellName,
-          cellCode: cellData?.code || "",
-          output,
-          outputType,
-          processedContent,
-          imageUrl,
-          shouldDownloadImage,
-        },
+        description: `Cell output (${mimetype})`,
+        data: { ...cellContextData, cellOutput },
       });
     }
 
@@ -149,16 +118,16 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
   }
 
   formatCompletion(item: CellOutputContextItem): Completion {
-    const { data } = item;
+    const { cellOutput: data, cellName, cellCode } = item.data;
 
     return {
-      label: `@${data.cellName}`,
-      displayLabel: data.cellName,
+      label: `@${cellName}`,
+      displayLabel: cellName,
       detail: `${data.outputType} output`,
       boost: data.outputType === "media" ? Boosts.HIGH : Boosts.MEDIUM,
       type: this.contextType,
       section: Sections.CELL_OUTPUT,
-      apply: `@${data.cellName}`,
+      apply: `@${cellName}`,
       info: () => {
         const infoContainer = document.createElement("div");
         infoContainer.classList.add(
@@ -176,7 +145,7 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
 
         const nameDiv = document.createElement("div");
         nameDiv.classList.add("font-bold", "text-base");
-        nameDiv.textContent = data.cellName;
+        nameDiv.textContent = cellName;
         headerDiv.append(nameDiv);
 
         const descriptionDiv = document.createElement("div");
@@ -186,7 +155,7 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
         infoContainer.append(headerDiv);
 
         // Show cell code preview
-        if (data.cellCode) {
+        if (cellCode) {
           const codeHeaderDiv = document.createElement("div");
           codeHeaderDiv.classList.add(
             "text-xs",
@@ -207,8 +176,7 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
             "overflow-y-auto",
           );
           codeDiv.textContent =
-            data.cellCode.slice(0, 200) +
-            (data.cellCode.length > 200 ? "..." : "");
+            cellCode.slice(0, 200) + (cellCode.length > 200 ? "..." : "");
           infoContainer.append(codeDiv);
         }
 
@@ -258,16 +226,16 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
   }
 
   formatContext(item: CellOutputContextItem): string {
-    const { data } = item;
+    const { cellOutput: data, cellName, cellId, cellCode } = item.data;
 
     const contextData = {
-      name: data.cellName,
-      cellId: data.cellId,
+      name: cellName,
+      cellId: cellId,
       outputType: data.outputType,
       mimetype: data.output.mimetype,
     } as const;
 
-    let details = `Cell Code:\n${data.cellCode}\n\n`;
+    let details = `Cell Code:\n${cellCode}\n\n`;
 
     if (data.outputType === "text" && data.processedContent) {
       details += `Output:\n${data.processedContent}`;
@@ -289,45 +257,135 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
   override async getAttachments(
     items: CellOutputContextItem[],
   ): Promise<FileUIPart[]> {
-    // Filter items that need image downloading
-    const itemsNeedingDownload = items.filter(
-      (item) =>
-        item.data.shouldDownloadImage && item.data.outputType === "media",
-    );
+    const cellId = items[0].data.cellId;
+    const cellName = items[0].data.cellName;
+    const cellOutputs = items.map((item) => item.data.cellOutput);
 
-    if (itemsNeedingDownload.length === 0) {
+    return getAttachmentsForOutputs(cellOutputs, cellId, cellName);
+  }
+}
+
+export async function getAttachmentsForOutputs(
+  data: BaseOutput[],
+  cellId: CellId,
+  cellName: string,
+): Promise<FileUIPart[]> {
+  // Filter items that need image downloading
+  const itemsNeedingDownload = data.filter(
+    (item) => item.shouldDownloadImage && item.outputType === "media",
+  );
+
+  if (itemsNeedingDownload.length === 0) {
+    return [];
+  }
+
+  // Prepare download requests
+  const downloadRequests = itemsNeedingDownload.flatMap((item) => {
+    const outputElement = document.getElementById(CellOutputId.create(cellId));
+    if (!outputElement) {
+      Logger.warn(`Output element not found for cell ${cellId}`);
       return [];
     }
+    return {
+      cellId: cellId,
+      cellName: cellName,
+      mimetype: item.output.mimetype,
+      element: outputElement,
+    };
+  });
+
+  try {
+    return await Promise.all(
+      downloadRequests.map(async (item) => ({
+        type: "file",
+        filename: `${cellName}-output-screenshot`,
+        mediaType: "image/png",
+        url: await toPng(item.element),
+      })),
+    );
+  } catch (error) {
+    Logger.error("Error downloading cell output images:", error);
+    return [];
+  }
+}
 
-    // Prepare download requests
-    const downloadRequests = itemsNeedingDownload.flatMap((item) => {
-      const outputElement = document.getElementById(
-        CellOutputId.create(item.data.cellId),
-      );
-      if (!outputElement) {
-        Logger.warn(`Output element not found for cell ${item.data.cellId}`);
-        return [];
-      }
-      return {
-        cellId: item.data.cellId,
-        cellName: item.data.cellName,
-        mimetype: item.data.output.mimetype,
-        element: outputElement,
-      };
-    });
+function getBaseOutput(output: OutputMessage): BaseOutput | null {
+  if (isOutputEmpty(output)) {
+    return null;
+  }
 
-    try {
-      return await Promise.all(
-        downloadRequests.map(async (item) => ({
-          type: "file",
-          filename: `${item.cellName}-output-screenshot`,
-          mediaType: "image/png",
-          url: await toPng(item.element),
-        })),
-      );
-    } catch (error) {
-      Logger.error("Error downloading cell output images:", error);
-      return [];
+  const mimetype = output.mimetype;
+
+  // Determine output type
+  const isMedia = isMediaMimetype(mimetype, String(output.data));
+  const outputType = isMedia ? "media" : "text";
+
+  let processedContent: string | undefined;
+  let imageUrl: string | undefined;
+  let shouldDownloadImage = false;
+
+  // Process text content
+  if (outputType === "text" && typeof output.data === "string") {
+    processedContent =
+      mimetype === "text/html" ? parseHtmlContent(output.data) : output.data;
+  }
+
+  // Process media content - for now, we'll just note that it's media
+  if (outputType === "media") {
+    if (typeof output.data === "string" && output.data.startsWith("data:")) {
+      imageUrl = output.data; // Data URL
+    } else if (
+      typeof output.data === "string" &&
+      output.data.startsWith("http")
+    ) {
+      imageUrl = output.data; // External URL
+    } else {
+      // Download the image from the DOM element
+      shouldDownloadImage = true;
     }
   }
+
+  return {
+    processedContent,
+    imageUrl,
+    shouldDownloadImage,
+    outputType,
+    output,
+  };
+}
+
+export function getCellContextData(
+  cellId: CellId,
+  notebook: NotebookState,
+  opts?: {
+    includeConsoleOutput: boolean;
+  },
+): CellOutputData {
+  const { includeConsoleOutput = false } = opts || {};
+  const cellRuntime = notebook.cellRuntime[cellId];
+
+  const cellData = notebook.cellData[cellId];
+  const cellIndex = notebook.cellIds.inOrderIds.indexOf(cellId);
+  const cellName = displayCellName(cellData.name, cellIndex);
+
+  let consoleOutputs: BaseOutput[] | undefined;
+  if (includeConsoleOutput) {
+    consoleOutputs = cellRuntime.consoleOutputs
+      .map((output) => getBaseOutput(output))
+      .filter((output) => output !== null);
+  }
+
+  let cellOutput: BaseOutput | undefined;
+  const rawCellOutput = cellRuntime.output;
+  if (rawCellOutput && !isOutputEmpty(rawCellOutput)) {
+    cellOutput = getBaseOutput(rawCellOutput) ?? undefined;
+  }
+
+  return {
+    cellId,
+    cellName,
+    cellCode: cellData.code,
+    cellOutput,
+    consoleOutputs,
+  };
 }

package/src/core/ai/context/providers/tables.ts

@@ -43,8 +43,8 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
 
     // Build shape information
     const shape = [
-      num_rows != null ? `${num_rows} rows` : undefined,
-      num_columns != null ? `${num_columns} columns` : undefined,
+      num_rows == null ? undefined : `${num_rows} rows`,
+      num_columns == null ? undefined : `${num_columns} columns`,
     ]
       .filter(Boolean)
       .join(", ");

package/src/core/ai/staged-cells.ts

@@ -13,21 +13,30 @@ import { Logger } from "@/utils/Logger";
 import { maybeAddMarimoImport } from "../cells/add-missing-import";
 import {
   type CreateNewCellAction,
-  cellHandleAtom,
+  getCellEditorView,
   useCellActions,
 } from "../cells/cells";
 import type { LanguageAdapterType } from "../codemirror/language/types";
 import { updateEditorCodeFromPython } from "../codemirror/language/utils";
 import type { JotaiStore } from "../state/jotai";
+import type { EditType } from "./tools/edit-notebook-tool";
 
 /**
  * Cells that are staged for AI completion
- * They function similarly to cells in the notebook, but they can be deleted or accepted by the user.
- * We only track one set of staged cells at a time.
+ * They function similarly to cells in the notebook, but they can be accepted or rejected by the user.
+ * We track edited, new and deleted cells.
+ * And we only track one set of staged cells at a time.
  */
 
-const initialState = (): Set<CellId> => {
-  return new Set();
+export type Edit =
+  | { type: Extract<EditType, "update_cell">; previousCode: string }
+  | { type: Extract<EditType, "add_cell"> }
+  | { type: Extract<EditType, "delete_cell">; previousCode: string };
+
+export type StagedAICells = Map<CellId, Edit>;
+
+const initialState = (): StagedAICells => {
+  return new Map();
 };
 
 const {
@@ -36,18 +45,26 @@ const {
   createActions,
   reducer,
 } = createReducerAndAtoms(initialState, {
-  addStagedCell: (state, action: { cellId: CellId }) => {
-    const { cellId } = action;
-    return new Set([...state, cellId]);
+  addStagedCell: (state, action: { cellId: CellId; edit: Edit }) => {
+    const { cellId, edit } = action;
+    return new Map([...state, [cellId, edit]]);
   },
   removeStagedCell: (state, cellId: CellId) => {
-    return new Set([...state].filter((id) => id !== cellId));
+    const newState = new Map(state);
+    newState.delete(cellId);
+    return newState;
   },
   clearStagedCells: () => {
     return initialState();
   },
 });
 
+export {
+  useStagedAICellsActions,
+  createActions as createStagedAICellsActions,
+  reducer as stagedAICellsReducer,
+};
+
 interface UpdateStagedCellAction {
   cellId: CellId;
   code: string;
@@ -67,7 +84,7 @@ export function useStagedCells(store: JotaiStore) {
 
   const createStagedCell = (code: string): CellId => {
     const newCellId = CellId.create();
-    addStagedCell({ cellId: newCellId });
+    addStagedCell({ cellId: newCellId, edit: { type: "add_cell" } });
     createNewCell({
       cellId: "__end__",
       code,
@@ -86,8 +103,7 @@ export function useStagedCells(store: JotaiStore) {
       return;
     }
 
-    const cellHandle = store.get(cellHandleAtom(cellId));
-    const editorView = cellHandle?.current?.editorViewOrNull;
+    const editorView = getCellEditorView(cellId);
     if (!editorView) {
       Logger.error("Editor for this cell not found", { cellId });
       return;
@@ -105,7 +121,7 @@ export function useStagedCells(store: JotaiStore) {
   // Delete all staged cells and the corresponding cells in the notebook.
   const deleteAllStagedCells = () => {
     const stagedAICells = store.get(stagedAICellsAtom);
-    for (const cellId of stagedAICells) {
+    for (const cellId of stagedAICells.keys()) {
       deleteCellCallback({ cellId });
     }
     clearStagedCells();
@@ -175,14 +191,14 @@ class CellCreationStream {
 
   private onCreateCell: (code: string) => CellId;
   private onUpdateCell: (opts: UpdateStagedCellAction) => void;
-  private addStagedCell: (payload: { cellId: CellId }) => void;
+  private addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
   private createNewCell: (opts: CreateNewCellAction) => void;
   private hasMarimoImport = false;
 
   constructor(
     onCreateCell: (code: string) => CellId,
     onUpdateCell: (opts: UpdateStagedCellAction) => void,
-    addStagedCell: (payload: { cellId: CellId }) => void,
+    addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void,
     createNewCell: (opts: CreateNewCellAction) => void,
   ) {
     this.onCreateCell = onCreateCell;
@@ -229,7 +245,7 @@ class CellCreationStream {
       before: true,
     });
     if (cellId) {
-      this.addStagedCell({ cellId });
+      this.addStagedCell({ cellId, edit: { type: "add_cell" } });
     }
     this.hasMarimoImport = true;
   }

package/src/core/ai/state.ts

@@ -5,7 +5,7 @@ import type { FileUIPart } from "ai";
 import { atom } from "jotai";
 import { atomWithStorage } from "jotai/utils";
 import { uniqueBy } from "@/utils/arrays";
-import { adaptForLocalStorage } from "@/utils/storage";
+import { adaptForLocalStorage, jotaiJsonStorage } from "@/utils/storage/jotai";
 import type { TypedString } from "@/utils/typed";
 import type { CellId } from "../cells/ids";
 
@@ -25,6 +25,7 @@ const INCLUDE_OTHER_CELLS_KEY = "marimo:ai:includeOtherCells";
 export const includeOtherCellsAtom = atomWithStorage<boolean>(
   INCLUDE_OTHER_CELLS_KEY,
   true,
+  jotaiJsonStorage,
 );
 
 export interface Message {