@marimo-team/islands 0.16.5 → 0.17.1

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

@@ -6,6 +6,7 @@ import { beforeEach, describe, expect, it, vi } from "vitest";
 import { CellId } from "@/core/cells/ids";
 import { updateEditorCodeFromPython } from "../../codemirror/language/utils";
 import {
+  type StagedAICells,
   stagedAICellsAtom,
   useStagedCells,
   visibleForTesting,
@@ -35,6 +36,7 @@ vi.mock("../../cells/cells", () => ({
   cellHandleAtom: vi.fn(() => ({
     read: vi.fn(() => mockCellHandle),
   })),
+  getCellEditorView: vi.fn(() => mockCellHandle.current.editorViewOrNull),
 }));
 
 vi.mock("@/components/editor/cell/useDeleteCell", () => ({
@@ -66,32 +68,77 @@ describe("staged-cells", () => {
     vi.clearAllMocks();
 
     // Reset the atom state
-    store.set(stagedAICellsAtom, new Set<CellId>());
+    store.set(stagedAICellsAtom, new Map());
   });
 
   describe("reducer and actions", () => {
     it("should initialize with empty map", () => {
       const state = initialState();
-      expect(state).toEqual(new Set());
+      expect(state).toEqual(new Map());
     });
 
-    it("should add cell IDs", () => {
+    it("should add cells with update_cell edit", () => {
       let state = initialState();
       state = reducer(state, {
         type: "addStagedCell",
-        payload: { cellId: cellId1 },
+        payload: {
+          cellId: cellId1,
+          edit: { type: "update_cell", previousCode: "old code 1" },
+        },
       });
       state = reducer(state, {
         type: "addStagedCell",
-        payload: { cellId: cellId2 },
+        payload: {
+          cellId: cellId2,
+          edit: { type: "update_cell", previousCode: "old code 2" },
+        },
       });
 
       expect(state.has(cellId1)).toBe(true);
       expect(state.has(cellId2)).toBe(true);
+      expect(state.get(cellId1)).toEqual({
+        type: "update_cell",
+        previousCode: "old code 1",
+      });
+      expect(state.get(cellId2)).toEqual({
+        type: "update_cell",
+        previousCode: "old code 2",
+      });
+    });
+
+    it("should add cells with add_cell edit", () => {
+      let state = initialState();
+      state = reducer(state, {
+        type: "addStagedCell",
+        payload: { cellId: cellId1, edit: { type: "add_cell" } },
+      });
+
+      expect(state.has(cellId1)).toBe(true);
+      expect(state.get(cellId1)).toEqual({ type: "add_cell" });
+    });
+
+    it("should add cells with delete_cell edit", () => {
+      let state = initialState();
+      state = reducer(state, {
+        type: "addStagedCell",
+        payload: {
+          cellId: cellId1,
+          edit: { type: "delete_cell", previousCode: "deleted code" },
+        },
+      });
+
+      expect(state.has(cellId1)).toBe(true);
+      expect(state.get(cellId1)).toEqual({
+        type: "delete_cell",
+        previousCode: "deleted code",
+      });
     });
 
     it("should remove cell IDs", () => {
-      const state = new Set([cellId1, cellId2]);
+      const state = new Map([
+        [cellId1, { type: "add_cell" as const }],
+        [cellId2, { type: "add_cell" as const }],
+      ]);
       const newState = reducer(state, {
         type: "removeStagedCell",
         payload: cellId1,
@@ -101,23 +148,26 @@ describe("staged-cells", () => {
       expect(newState.has(cellId2)).toBe(true);
     });
 
-    it("should clear all cell IDs", () => {
-      const state = new Set([cellId1, cellId2]);
+    it("should clear all cells", () => {
+      const state = new Map([
+        [cellId1, { type: "add_cell" as const }],
+        [cellId2, { type: "add_cell" as const }],
+      ]);
       const newState = reducer(state, {
         type: "clearStagedCells",
         payload: undefined,
       });
 
-      expect(newState).toEqual(new Set());
+      expect(newState).toEqual(new Map());
     });
 
-    it("should not mutate original state", () => {
-      const state = new Set([cellId1]);
+    it("should not mutate original state when adding", () => {
+      const state = new Map([[cellId1, { type: "add_cell" as const }]]);
       const originalSize = state.size;
 
       reducer(state, {
         type: "addStagedCell",
-        payload: { cellId: cellId2 },
+        payload: { cellId: cellId2, edit: { type: "add_cell" } },
       });
 
       expect(state.size).toBe(originalSize);
@@ -125,6 +175,23 @@ describe("staged-cells", () => {
       expect(state.has(cellId2)).toBe(false);
     });
 
+    it("should not mutate original state when removing", () => {
+      const state = new Map([
+        [cellId1, { type: "add_cell" as const }],
+        [cellId2, { type: "add_cell" as const }],
+      ]);
+      const originalSize = state.size;
+
+      reducer(state, {
+        type: "removeStagedCell",
+        payload: cellId1,
+      });
+
+      expect(state.size).toBe(originalSize);
+      expect(state.has(cellId1)).toBe(true);
+      expect(state.has(cellId2)).toBe(true);
+    });
+
     it("should create action functions", () => {
       const mockDispatch = vi.fn();
       const actions = createActions(mockDispatch);
@@ -136,7 +203,7 @@ describe("staged-cells", () => {
 
     it("should initialize atom with empty map", () => {
       const state = store.get(stagedAICellsAtom);
-      expect(state).toEqual(new Set());
+      expect(state).toEqual(new Map());
     });
   });
 
@@ -181,7 +248,11 @@ describe("staged-cells", () => {
 
     it("should delete all staged cells when cells exist", () => {
       // First set the atom state before rendering the hook
-      store.set(stagedAICellsAtom, new Set([cellId1, cellId2]));
+      const initialState: StagedAICells = new Map([
+        [cellId1, { type: "add_cell" }],
+        [cellId2, { type: "add_cell" }],
+      ]);
+      store.set(stagedAICellsAtom, initialState);
 
       const { result } = renderHook(() => useStagedCells(store));
       result.current.deleteAllStagedCells();
@@ -192,76 +263,130 @@ describe("staged-cells", () => {
 
       // Verify cells were cleared from the atom
       const state = store.get(stagedAICellsAtom);
-      expect(state).toEqual(new Set());
+      expect(state).toEqual(new Map());
     });
-  });
 
-  it("should add staged cell", () => {
-    const { result } = renderHook(() => useStagedCells(store));
+    it("should add staged cell with edit info", () => {
+      const { result } = renderHook(() => useStagedCells(store));
 
-    result.current.addStagedCell({ cellId: cellId1 });
+      result.current.addStagedCell({
+        cellId: cellId1,
+        edit: { type: "update_cell", previousCode: "old code" },
+      });
 
-    // Check that the cell was added to the atom
-    const state = store.get(stagedAICellsAtom);
-    expect(state.has(cellId1)).toBe(true);
-  });
+      // Check that the cell was added to the atom with edit info
+      const state = store.get(stagedAICellsAtom);
+      expect(state.has(cellId1)).toBe(true);
+      expect(state.get(cellId1)).toEqual({
+        type: "update_cell",
+        previousCode: "old code",
+      });
+    });
 
-  it("should remove staged cell", () => {
-    const { result } = renderHook(() => useStagedCells(store));
+    it("should remove staged cell", () => {
+      const { result } = renderHook(() => useStagedCells(store));
 
-    // First add cells
-    result.current.addStagedCell({ cellId: cellId1 });
-    result.current.addStagedCell({ cellId: cellId2 });
+      // First add cells
+      result.current.addStagedCell({
+        cellId: cellId1,
+        edit: { type: "add_cell" },
+      });
+      result.current.addStagedCell({
+        cellId: cellId2,
+        edit: { type: "add_cell" },
+      });
 
-    // Then remove one
-    result.current.removeStagedCell(cellId1);
+      // Then remove one
+      result.current.removeStagedCell(cellId1);
 
-    // Check that only the remaining cell is in the map
-    const state = store.get(stagedAICellsAtom);
-    expect(state.has(cellId1)).toBe(false);
-    expect(state.has(cellId2)).toBe(true);
-  });
+      // Check that only the remaining cell is in the map
+      const state = store.get(stagedAICellsAtom);
+      expect(state.has(cellId1)).toBe(false);
+      expect(state.has(cellId2)).toBe(true);
+    });
 
-  it("should clear all staged cells", () => {
-    const { result } = renderHook(() => useStagedCells(store));
+    it("should clear all staged cells", () => {
+      const { result } = renderHook(() => useStagedCells(store));
 
-    // First add some cells
-    result.current.addStagedCell({ cellId: cellId1 });
-    result.current.addStagedCell({ cellId: cellId2 });
+      // First add some cells
+      result.current.addStagedCell({
+        cellId: cellId1,
+        edit: { type: "add_cell" },
+      });
+      result.current.addStagedCell({
+        cellId: cellId2,
+        edit: { type: "add_cell" },
+      });
 
-    // Then clear all
-    result.current.clearStagedCells();
+      // Then clear all
+      result.current.clearStagedCells();
 
-    // Check that no cells remain
-    const state = store.get(stagedAICellsAtom);
-    expect(state).toEqual(new Set());
-  });
+      // Check that no cells remain
+      const state = store.get(stagedAICellsAtom);
+      expect(state).toEqual(new Map());
+    });
 
-  it("should handle multiple operations correctly", () => {
-    const { result } = renderHook(() => useStagedCells(store));
+    it("should handle multiple operations correctly", () => {
+      const { result } = renderHook(() => useStagedCells(store));
 
-    // Create a staged cell
-    const mockCellId = "mock-cell-id" as CellId;
-    vi.mocked(CellId.create).mockReturnValue(mockCellId);
+      // Create a staged cell
+      const mockCellId = "mock-cell-id" as CellId;
+      vi.mocked(CellId.create).mockReturnValue(mockCellId);
 
-    const createdCellId = result.current.createStagedCell("test code");
+      const createdCellId = result.current.createStagedCell("test code");
 
-    // Verify it was created and added
-    expect(createdCellId).toBe(mockCellId);
-    expect(mockCreateNewCell).toHaveBeenCalled();
+      // Verify it was created and added
+      expect(createdCellId).toBe(mockCellId);
+      expect(mockCreateNewCell).toHaveBeenCalled();
 
-    let state = store.get(stagedAICellsAtom);
-    expect(state.has(mockCellId)).toBe(true);
+      let state = store.get(stagedAICellsAtom);
+      expect(state.has(mockCellId)).toBe(true);
+      expect(state.get(mockCellId)).toEqual({ type: "add_cell" });
 
-    // Delete the staged cell
-    result.current.deleteStagedCell(mockCellId);
-    expect(mockDeleteCellCallback).toHaveBeenCalledWith({
-      cellId: mockCellId,
+      // Delete the staged cell
+      result.current.deleteStagedCell(mockCellId);
+      expect(mockDeleteCellCallback).toHaveBeenCalledWith({
+        cellId: mockCellId,
+      });
+
+      // Verify it was removed from staged cells
+      state = store.get(stagedAICellsAtom);
+      expect(state.has(mockCellId)).toBe(false);
     });
 
-    // Verify it was removed from staged cells
-    state = store.get(stagedAICellsAtom);
-    expect(state.has(mockCellId)).toBe(false);
+    it("should track edit history for updated cells", () => {
+      const { result } = renderHook(() => useStagedCells(store));
+
+      // Add a cell with update_cell edit type
+      result.current.addStagedCell({
+        cellId: cellId1,
+        edit: { type: "update_cell", previousCode: "previous code" },
+      });
+
+      const state = store.get(stagedAICellsAtom);
+      const edit = state.get(cellId1);
+      expect(edit).toEqual({
+        type: "update_cell",
+        previousCode: "previous code",
+      });
+    });
+
+    it("should track edit history for deleted cells", () => {
+      const { result } = renderHook(() => useStagedCells(store));
+
+      // Add a cell with delete_cell edit type
+      result.current.addStagedCell({
+        cellId: cellId1,
+        edit: { type: "delete_cell", previousCode: "deleted content" },
+      });
+
+      const state = store.get(stagedAICellsAtom);
+      const edit = state.get(cellId1);
+      expect(edit).toEqual({
+        type: "delete_cell",
+        previousCode: "deleted content",
+      });
+    });
   });
 });
 

package/src/core/ai/config.ts

@@ -4,7 +4,11 @@ import type { Role } from "@marimo-team/llm-info";
 import { useAtom } from "jotai";
 import type { QualifiedModelId } from "@/core/ai/ids/ids";
 import { userConfigAtom } from "@/core/config/config";
-import type { AIModelKey, UserConfig } from "@/core/config/config-schema";
+import type {
+  AIModelKey,
+  CopilotMode,
+  UserConfig,
+} from "@/core/config/config-schema";
 import { useRequestClient } from "@/core/network/requests";
 
 // Extract only the supported roles from the Role type
@@ -60,7 +64,7 @@ export const useModelChange = () => {
     saveConfig(newConfig);
   };
 
-  const saveModeChange = async (newMode: "ask" | "manual") => {
+  const saveModeChange = async (newMode: CopilotMode) => {
     const newConfig: UserConfig = {
       ...userConfig,
       ai: {

package/src/core/ai/context/providers/__tests__/__snapshots__/cell-output.test.ts.snap

@@ -0,0 +1,17 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`CellOutputContextProvider > formatContext > should format media output context correctly > media-output-context 1`] = `
+"<cell-output name="cell-1" cellId="cell2" outputType="media" mimetype="image/png">Cell Code:
+import matplotlib.pyplot as plt
+plt.plot([1,2,3])
+
+Media Output: Contains image/png content</cell-output>"
+`;
+
+exports[`CellOutputContextProvider > formatContext > should format text output context correctly > text-output-context 1`] = `
+"<cell-output name="My Cell" cellId="cell1" outputType="text" mimetype="text/plain">Cell Code:
+print('hello world')
+
+Output:
+hello world</cell-output>"
+`;

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

@@ -16,7 +16,7 @@ import type { NotebookState } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import type { OutputMessage } from "@/core/kernel/messages";
 import type { JotaiStore } from "@/core/state/jotai";
-import { CellOutputContextProvider } from "../cell-output";
+import { CellOutputContextProvider, getCellContextData } from "../cell-output";
 
 // Test helper to create mock store
 function createMockStore(notebook: NotebookState): JotaiStore {
@@ -104,23 +104,23 @@ describe("CellOutputContextProvider", () => {
       const textItem = items.find((item) => item.data.cellId === "cell1");
       const mediaItem = items.find((item) => item.data.cellId === "cell2");
 
-      expect(textItem?.data.outputType).toBe("text");
-      expect(mediaItem?.data.outputType).toBe("media");
+      expect(textItem?.data.cellOutput.outputType).toBe("text");
+      expect(mediaItem?.data.cellOutput.outputType).toBe("media");
     });
 
     it("should process text content correctly", () => {
       const items = provider.getItems();
       const textItem = items.find((item) => item.data.cellId === "cell1");
 
-      expect(textItem?.data.processedContent).toBe("hello world");
+      expect(textItem?.data.cellOutput.processedContent).toBe("hello world");
     });
 
     it("should mark media items for download when no direct URL available", () => {
       const items = provider.getItems();
       const mediaItem = items.find((item) => item.data.cellId === "cell2");
 
-      expect(mediaItem?.data.shouldDownloadImage).toBe(true);
-      expect(mediaItem?.data.imageUrl).toBeUndefined();
+      expect(mediaItem?.data.cellOutput.shouldDownloadImage).toBe(true);
+      expect(mediaItem?.data.cellOutput.imageUrl).toBeUndefined();
     });
 
     it("should include cell code and names", () => {
@@ -156,6 +156,8 @@ describe("CellOutputContextProvider", () => {
       expect(context).toContain("print('hello world')");
       expect(context).toContain("Output:");
       expect(context).toContain("hello world");
+
+      expect(context).toMatchSnapshot("text-output-context");
     });
 
     it("should format media output context correctly", () => {
@@ -173,6 +175,8 @@ describe("CellOutputContextProvider", () => {
       expect(context).toContain("Cell Code:");
       expect(context).toContain("import matplotlib.pyplot as plt");
       expect(context).toContain("Media Output: Contains image/png content");
+
+      expect(context).toMatchSnapshot("media-output-context");
     });
   });
 
@@ -239,7 +243,7 @@ describe("Cell output utility functions", () => {
         provider = new CellOutputContextProvider(mockStore);
         const items = provider.getItems();
 
-        expect(items[0]?.data.outputType).toBe(testCase.expected);
+        expect(items[0]?.data.cellOutput.outputType).toBe(testCase.expected);
       }
     });
 
@@ -274,7 +278,7 @@ describe("Cell output utility functions", () => {
         provider = new CellOutputContextProvider(mockStore);
         const items = provider.getItems();
 
-        expect(items[0]?.data.outputType).toBe(testCase.expected);
+        expect(items[0]?.data.cellOutput.outputType).toBe(testCase.expected);
       }
     });
 
@@ -308,7 +312,7 @@ describe("Cell output utility functions", () => {
         provider = new CellOutputContextProvider(mockStore);
         const items = provider.getItems();
 
-        expect(items[0]?.data.outputType).toBe(testCase.expected);
+        expect(items[0]?.data.cellOutput.outputType).toBe(testCase.expected);
       }
     });
   });
@@ -371,7 +375,187 @@ describe("Cell output utility functions", () => {
       provider = new CellOutputContextProvider(mockStore);
       const items = provider.getItems();
 
-      expect(items[0]?.data.processedContent).toBe("Hello world!");
+      expect(items[0]?.data.cellOutput.processedContent).toBe("Hello world!");
+    });
+  });
+
+  describe("getCellContextData", () => {
+    let mockNotebook: NotebookState;
+
+    beforeEach(() => {
+      mockNotebook = {
+        cellIds: {
+          inOrderIds: ["cell1" as CellId, "cell2" as CellId],
+        },
+        cellData: {
+          cell1: {
+            id: "cell1" as CellId,
+            name: "My Named Cell",
+            code: "x = 42",
+          },
+          cell2: {
+            id: "cell2" as CellId,
+            name: "",
+            code: "y = x * 2",
+          },
+        },
+        cellRuntime: {
+          cell1: {
+            output: {
+              mimetype: "text/plain",
+              data: "42",
+            } as OutputMessage,
+            consoleOutputs: [],
+          },
+          cell2: {
+            output: null,
+            consoleOutputs: [
+              {
+                mimetype: "text/plain",
+                data: "console log message",
+              } as OutputMessage,
+            ],
+          },
+        },
+      } as unknown as NotebookState;
+    });
+
+    it("should extract basic cell data", () => {
+      const result = getCellContextData("cell1" as CellId, mockNotebook);
+
+      expect(result.cellId).toBe("cell1");
+      expect(result.cellName).toBe("My Named Cell");
+      expect(result.cellCode).toBe("x = 42");
+    });
+
+    it("should generate cell name for unnamed cells", () => {
+      const result = getCellContextData("cell2" as CellId, mockNotebook);
+
+      expect(result.cellName).toBe("cell-1"); // 0-indexed, so cell2 is at index 1
+    });
+
+    it("should include cell output when present", () => {
+      const result = getCellContextData("cell1" as CellId, mockNotebook);
+
+      expect(result.cellOutput).toBeDefined();
+      expect(result.cellOutput?.outputType).toBe("text");
+      expect(result.cellOutput?.processedContent).toBe("42");
+    });
+
+    it("should not include cell output when not present", () => {
+      const result = getCellContextData("cell2" as CellId, mockNotebook);
+
+      expect(result.cellOutput).toBeUndefined();
+    });
+
+    it("should not include console outputs by default", () => {
+      const result = getCellContextData("cell2" as CellId, mockNotebook);
+
+      expect(result.consoleOutputs).toBeUndefined();
+    });
+
+    it("should include console outputs when opted in", () => {
+      const result = getCellContextData("cell2" as CellId, mockNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      expect(result.consoleOutputs).toBeDefined();
+      expect(result.consoleOutputs).toHaveLength(1);
+      expect(result.consoleOutputs?.[0]?.processedContent).toBe(
+        "console log message",
+      );
+    });
+
+    it("should filter out empty console outputs", () => {
+      const cell2Runtime = mockNotebook.cellRuntime["cell2" as CellId];
+      if (cell2Runtime) {
+        (cell2Runtime as { consoleOutputs: OutputMessage[] }).consoleOutputs = [
+          {
+            mimetype: "text/plain",
+            data: "valid output",
+          } as OutputMessage,
+          {
+            mimetype: "text/plain",
+            data: "",
+          } as OutputMessage,
+        ];
+      }
+
+      const result = getCellContextData("cell2" as CellId, mockNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      expect(result.consoleOutputs).toHaveLength(1);
+      expect(result.consoleOutputs?.[0]?.processedContent).toBe("valid output");
+    });
+
+    it("should handle cells with both cell output and console outputs", () => {
+      const cell1Runtime = mockNotebook.cellRuntime["cell1" as CellId];
+      if (cell1Runtime) {
+        (cell1Runtime as { consoleOutputs: OutputMessage[] }).consoleOutputs = [
+          {
+            mimetype: "text/plain",
+            data: "debug message",
+          } as OutputMessage,
+        ];
+      }
+
+      const result = getCellContextData("cell1" as CellId, mockNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      expect(result.cellOutput).toBeDefined();
+      expect(result.cellOutput?.processedContent).toBe("42");
+      expect(result.consoleOutputs).toHaveLength(1);
+      expect(result.consoleOutputs?.[0]?.processedContent).toBe(
+        "debug message",
+      );
+    });
+
+    it("should handle empty console outputs array", () => {
+      const result = getCellContextData("cell1" as CellId, mockNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      expect(result.consoleOutputs).toEqual([]);
+    });
+
+    it("should handle media outputs in cell output", () => {
+      const cell1Runtime = mockNotebook.cellRuntime["cell1" as CellId];
+      if (cell1Runtime) {
+        (cell1Runtime as { output: OutputMessage | null }).output = {
+          mimetype: "image/png",
+          data: "data:image/png;base64,test",
+        } as OutputMessage;
+      }
+
+      const result = getCellContextData("cell1" as CellId, mockNotebook);
+
+      expect(result.cellOutput).toBeDefined();
+      expect(result.cellOutput?.outputType).toBe("media");
+      expect(result.cellOutput?.imageUrl).toBe("data:image/png;base64,test");
+    });
+
+    it("should handle media outputs in console outputs", () => {
+      const cell2Runtime = mockNotebook.cellRuntime["cell2" as CellId];
+      if (cell2Runtime) {
+        (cell2Runtime as { consoleOutputs: OutputMessage[] }).consoleOutputs = [
+          {
+            mimetype: "image/png",
+            data: "https://example.com/image.png",
+          } as OutputMessage,
+        ];
+      }
+
+      const result = getCellContextData("cell2" as CellId, mockNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      expect(result.consoleOutputs).toHaveLength(1);
+      expect(result.consoleOutputs?.[0]?.outputType).toBe("media");
+      expect(result.consoleOutputs?.[0]?.imageUrl).toBe(
+        "https://example.com/image.png",
+      );
     });
   });
 });