@marimo-team/islands 0.16.5 → 0.17.1

package/src/core/ai/context/providers/__tests__/datasource.test.ts

@@ -2,6 +2,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import { beforeEach, describe, expect, it } from "vitest";
+import type { CellId } from "@/core/cells/ids";
 import type {
   ConnectionsMap,
   DatasetTablesMap,
@@ -9,7 +10,7 @@ import type {
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
 import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
 import { Boosts, Sections } from "../common";
-import { DatasourceContextProvider } from "../datasource";
+import { DatasourceContextProvider, getDatasourceContext } from "../datasource";
 
 // Mock data for testing
 const createMockDataSourceConnection = (
@@ -630,3 +631,10 @@ describe("DatasourceContextProvider", () => {
     });
   });
 });
+
+describe("getDatasourceContext", () => {
+  it("should return null if no cell ID is found", () => {
+    const context = getDatasourceContext("1" as CellId);
+    expect(context).toBeNull();
+  });
+});

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/datasource.ts

@@ -3,10 +3,14 @@
 import type { Completion } from "@codemirror/autocomplete";
 import { createRoot } from "react-dom/client";
 import { dbDisplayName } from "@/components/databases/display";
+import { cellDataAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { LanguageAdapters } from "@/core/codemirror/language/LanguageAdapters";
 import { renderDatasourceInfo } from "@/core/codemirror/language/languages/sql/renderers";
 import {
   type ConnectionsMap,
   type DatasetTablesMap,
+  dataSourceConnectionsAtom,
   getTableType,
 } from "@/core/datasets/data-source-connections";
 import {
@@ -14,6 +18,7 @@ import {
   INTERNAL_SQL_ENGINES,
 } from "@/core/datasets/engines";
 import type { DataSourceConnection, DataTable } from "@/core/kernel/messages";
+import { store } from "@/core/state/jotai";
 import type { AIContextItem } from "../registry";
 import { AIContextProvider } from "../registry";
 import { contextToXml } from "../utils";
@@ -37,10 +42,12 @@ export interface DatasourceContextItem extends AIContextItem {
   };
 }
 
+const CONTEXT_TYPE = "datasource";
+
 export class DatasourceContextProvider extends AIContextProvider<DatasourceContextItem> {
   readonly title = "Datasource";
   readonly mentionPrefix = "@";
-  readonly contextType = "datasource";
+  readonly contextType = CONTEXT_TYPE;
   private connectionsMap: ConnectionsMap;
   private dataframes: DataTable[];
 
@@ -140,3 +147,20 @@ export class DatasourceContextProvider extends AIContextProvider<DatasourceConte
     };
   }
 }
+
+export function getDatasourceContext(cellId: CellId): string | null {
+  const cellData = store.get(cellDataAtom(cellId));
+  const code = cellData?.code;
+  if (!code || code.trim() === "") {
+    return null;
+  }
+
+  const [_sqlStatement, _, metadata] = LanguageAdapters.sql.transformIn(code);
+  const datasourceSchema = store
+    .get(dataSourceConnectionsAtom)
+    .connectionsMap.get(metadata.engine);
+  if (datasourceSchema) {
+    return `@${CONTEXT_TYPE}://${datasourceSchema.name}`;
+  }
+  return null;
+}

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";
 
@@ -13,15 +13,19 @@ const KEY = "marimo:ai:chatState:v5";
 
 export type ChatId = TypedString<"ChatId">;
 
-export const aiCompletionCellAtom = atom<{
+export interface AiCompletionCell {
   cellId: CellId;
   initialPrompt?: string;
-} | null>(null);
+  triggerImmediately?: boolean;
+}
+
+export const aiCompletionCellAtom = atom<AiCompletionCell | null>(null);
 
 const INCLUDE_OTHER_CELLS_KEY = "marimo:ai:includeOtherCells";
 export const includeOtherCellsAtom = atomWithStorage<boolean>(
   INCLUDE_OTHER_CELLS_KEY,
   true,
+  jotaiJsonStorage,
 );
 
 export interface Message {
@@ -86,7 +90,7 @@ export const activeChatAtom = atom(
     }
     return state.chats.get(state.activeChatId);
   },
-  (get, set, chatId: ChatId | null) => {
+  (_get, set, chatId: ChatId | null) => {
     set(chatStateAtom, (prev) => ({
       ...prev,
       activeChatId: chatId,