@marimo-team/islands 0.25.3-dev21 → 0.25.3-dev210

package/src/core/ai/tools/__tests__/run-cells-tool.test.ts

@@ -1,656 +0,0 @@
-/* Copyright 2026 Marimo. All rights reserved. */
-
-import { getDefaultStore } from "jotai";
-import { beforeEach, describe, expect, it, vi } from "vitest";
-import { MockNotebook } from "@/__mocks__/notebook";
-import { cellId } from "@/__tests__/branded";
-import { notebookAtom } from "@/core/cells/cells";
-import type { CellId } from "@/core/cells/ids";
-import { RunStaleCellsTool } from "../run-cells-tool";
-
-// Mock runCells
-vi.mock("@/components/editor/cell/useRunCells", () => ({
-  runCells: vi.fn(),
-}));
-
-// Mock getCellContextData
-vi.mock("../../context/providers/cell-output", () => ({
-  getCellContextData: vi.fn(),
-}));
-
-import { runCells } from "@/components/editor/cell/useRunCells";
-import { getCellContextData } from "../../context/providers/cell-output";
-
-describe("RunStaleCellsTool", () => {
-  let store: ReturnType<typeof getDefaultStore>;
-  let tool: RunStaleCellsTool;
-  let cellId1: CellId;
-  let cellId2: CellId;
-  let cellId3: CellId;
-  let toolContext: {
-    addStagedCell: ReturnType<typeof vi.fn>;
-    createNewCell: ReturnType<typeof vi.fn>;
-    prepareForRun: ReturnType<typeof vi.fn>;
-    sendRun: ReturnType<typeof vi.fn>;
-    store: ReturnType<typeof getDefaultStore>;
-  };
-
-  beforeEach(() => {
-    store = getDefaultStore();
-    toolContext = {
-      addStagedCell: vi.fn(),
-      createNewCell: vi.fn(),
-      prepareForRun: vi.fn(),
-      sendRun: vi.fn().mockResolvedValue(null),
-      store,
-    };
-
-    tool = new RunStaleCellsTool({ postExecutionDelay: 0 });
-
-    cellId1 = cellId("cell-1");
-    cellId2 = cellId("cell-2");
-    cellId3 = cellId("cell-3");
-
-    // Reset mocks
-    vi.clearAllMocks();
-  });
-
-  describe("no stale cells", () => {
-    it("should return success message when no stale cells exist", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: {
-            code: "x = 1",
-            edited: false,
-            lastCodeRun: "x = 1",
-            lastExecutionTime: 100,
-          },
-        },
-      });
-      // Make sure cell is not stale by setting runtime state
-      notebook.cellRuntime[cellId1] = {
-        ...notebook.cellRuntime[cellId1],
-        runElapsedTimeMs: 100 as never,
-        status: "idle",
-      };
-      store.set(notebookAtom, notebook);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.message).toBe("No stale cells found.");
-      expect(result.cellsToOutput).toBeUndefined();
-      expect(vi.mocked(runCells)).not.toHaveBeenCalled();
-    });
-  });
-
-  describe("running stale cells", () => {
-    it("should run stale cells and return their outputs", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: {
-            code: "x = 1",
-            edited: true,
-            lastCodeRun: null,
-          },
-          [cellId2]: {
-            code: "y = 2",
-            edited: false,
-            lastCodeRun: "y = 2",
-            lastExecutionTime: 100,
-          },
-          [cellId3]: {
-            code: "z = 3",
-            edited: true,
-            lastCodeRun: "z = 2",
-          },
-        },
-      });
-      // Make cellId2 not stale
-      notebook.cellRuntime[cellId2] = {
-        ...notebook.cellRuntime[cellId2],
-        runElapsedTimeMs: 100 as never,
-        status: "idle",
-      };
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells to update cell runtime status
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        updatedNotebook.cellRuntime[cellId3] = {
-          ...updatedNotebook.cellRuntime[cellId3],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return outputs
-      vi.mocked(getCellContextData).mockImplementation((cellId) => {
-        if (cellId === cellId1) {
-          return {
-            cellOutput: {
-              outputType: "text",
-              processedContent: "1",
-              imageUrl: null,
-              output: {},
-            },
-            consoleOutputs: null,
-            cellName: "cell1",
-          } as never;
-        }
-        if (cellId === cellId3) {
-          return {
-            cellOutput: {
-              outputType: "text",
-              processedContent: "3",
-              imageUrl: null,
-              output: {},
-            },
-            consoleOutputs: null,
-            cellName: "cell3",
-          } as never;
-        }
-        return {} as never;
-      });
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      expect(result.cellsToOutput).toHaveProperty(cellId1);
-      expect(result.cellsToOutput).toHaveProperty(cellId3);
-      expect(result.cellsToOutput).not.toHaveProperty(cellId2);
-
-      // Verify runCells was called with correct parameters
-      expect(vi.mocked(runCells)).toHaveBeenCalledWith({
-        cellIds: [cellId1, cellId3],
-        sendRun: toolContext.sendRun,
-        prepareForRun: toolContext.prepareForRun,
-        notebook,
-      });
-    });
-
-    it("should handle cells with no output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "x = 1", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return no outputs
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: null,
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      expect(result.cellsToOutput?.[cellId1]).toBeNull();
-    });
-
-    it("should handle cells with console output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: 'print("hello")', edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return console output
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: null,
-        consoleOutputs: [
-          {
-            outputType: "text",
-            processedContent: "hello",
-            imageUrl: null,
-            output: {},
-          },
-        ],
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      expect(result.cellsToOutput?.[cellId1]?.consoleOutput).toContain("hello");
-      expect(result.message).toContain("Console output");
-    });
-
-    it("should handle cells with media output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "plt.plot()", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return media output
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "media",
-          processedContent: null,
-          imageUrl: "https://example.com/image.png",
-          output: { mimetype: "image/png" },
-        },
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      expect(result.cellsToOutput?.[cellId1]?.cellOutput).toContain(
-        "image/png",
-      );
-      expect(result.cellsToOutput?.[cellId1]?.cellOutput).toContain(
-        "https://example.com/image.png",
-      );
-    });
-
-    it("should handle both cell output and console output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: 'x = 1; print("debug")', edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return both outputs
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: "1",
-          imageUrl: null,
-          output: {},
-        },
-        consoleOutputs: [
-          {
-            outputType: "text",
-            processedContent: "debug",
-            imageUrl: null,
-            output: {},
-          },
-        ],
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      expect(result.cellsToOutput?.[cellId1]?.cellOutput).toContain("1");
-      expect(result.cellsToOutput?.[cellId1]?.consoleOutput).toContain("debug");
-    });
-
-    it("should handle cell output with object output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "x = {'a': 1, 'b': 2}", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return object output
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: null,
-          imageUrl: null,
-          output: { data: JSON.stringify({ a: 1, b: 2 }) },
-        },
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      expect(result.cellsToOutput?.[cellId1]?.cellOutput).toEqual(
-        'Output:\n{"a":1,"b":2}',
-      );
-    });
-
-    it("should return success when all stale cells have no output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "x = 1", edited: true },
-          [cellId2]: { code: "y = 2", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        updatedNotebook.cellRuntime[cellId2] = {
-          ...updatedNotebook.cellRuntime[cellId2],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Mock getCellContextData to return no outputs for both cells
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: null,
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-      // Both cells should be in the output as null
-      expect(result.cellsToOutput?.[cellId1]).toBeNull();
-      expect(result.cellsToOutput?.[cellId2]).toBeNull();
-    });
-  });
-
-  describe("output truncation", () => {
-    it("should summarize text/html output instead of dumping raw content", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "fig.show()", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      const largeHtml = `<div>${"x".repeat(2_000_000)}</div>`;
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: null,
-          imageUrl: null,
-          output: { mimetype: "text/html", data: largeHtml },
-        },
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      const output = result.cellsToOutput?.[cellId1]?.cellOutput ?? "";
-      expect(output).toContain("HTML Output:");
-      expect(output).toContain("text/html");
-      expect(output.length).toBeLessThan(200);
-      expect(output).not.toContain(largeHtml);
-    });
-
-    it("should truncate large text output to MAX_TEXT_OUTPUT_CHARS", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "print(big_string)", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      const largeText = "a".repeat(10_000);
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: largeText,
-          imageUrl: null,
-          output: { mimetype: "text/plain", data: largeText },
-        },
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      const output = result.cellsToOutput?.[cellId1]?.cellOutput ?? "";
-      expect(output).toContain("[TRUNCATED:");
-      expect(output).toContain("Full output visible in the notebook UI.");
-      // Output should be capped (2000 chars content + "Output:\n" prefix + truncation message)
-      expect(output.length).toBeLessThan(2200);
-    });
-
-    it("should omit output for cells that exceed total output budget", async () => {
-      const cellIds = Array.from({ length: 25 }, (_, i) =>
-        cellId(`budget-cell-${i}`),
-      );
-      const cellData: Record<string, { code: string; edited: boolean }> = {};
-      for (const id of cellIds) {
-        cellData[id] = { code: "x = 1", edited: true };
-      }
-
-      const notebook = MockNotebook.notebookState({ cellData });
-      store.set(notebookAtom, notebook);
-
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        for (const id of cellIds) {
-          updatedNotebook.cellRuntime[id] = {
-            ...updatedNotebook.cellRuntime[id],
-            status: "idle",
-          };
-        }
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // Each cell produces ~2008 chars of formatted output ("Output:\n" + 2000 chars).
-      // After 20 cells the running total exceeds MAX_TOOL_OUTPUT_CHARS (40,000).
-      const content = "a".repeat(2000);
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: content,
-          imageUrl: null,
-          output: { mimetype: "text/plain", data: content },
-        },
-        consoleOutputs: null,
-        cellName: "cell",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.cellsToOutput?.[cellIds[0]]?.cellOutput).toContain(
-        "Output:",
-      );
-      expect(result.cellsToOutput?.[cellIds[24]]?.cellOutput).toBe(
-        "Cell executed (output omitted due to context limits).",
-      );
-    });
-
-    it("should use higher truncation limit for error outputs", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "raise Exception()", edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      // 2500 chars sits between MAX_TEXT_OUTPUT_CHARS (2000) and MAX_ERROR_OUTPUT_CHARS (3000)
-      const errorContent = "E".repeat(2500);
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: errorContent,
-          imageUrl: null,
-          output: {
-            mimetype: "application/vnd.marimo+error",
-            data: errorContent,
-          },
-        },
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      const output = result.cellsToOutput?.[cellId1]?.cellOutput ?? "";
-      expect(output).not.toContain("[TRUNCATED:");
-      expect(output).toContain(errorContent);
-    });
-
-    it("should truncate large console output", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: 'print("x" * 10000)', edited: true },
-        },
-      });
-      store.set(notebookAtom, notebook);
-
-      vi.mocked(runCells).mockImplementation(async () => {
-        const updatedNotebook = store.get(notebookAtom);
-        updatedNotebook.cellRuntime[cellId1] = {
-          ...updatedNotebook.cellRuntime[cellId1],
-          status: "idle",
-        };
-        store.set(notebookAtom, updatedNotebook);
-      });
-
-      const largeConsoleText = "x".repeat(10_000);
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: null,
-        consoleOutputs: [
-          {
-            outputType: "text",
-            processedContent: largeConsoleText,
-            imageUrl: null,
-            output: { mimetype: "text/plain", data: largeConsoleText },
-          },
-        ],
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      const consoleOutput =
-        result.cellsToOutput?.[cellId1]?.consoleOutput ?? "";
-      expect(consoleOutput).toContain("[TRUNCATED:");
-      expect(consoleOutput.length).toBeLessThan(2200);
-    });
-  });
-
-  describe("cell execution completion", () => {
-    it("should complete immediately if cells are already idle", async () => {
-      const notebook = MockNotebook.notebookState({
-        cellData: {
-          [cellId1]: { code: "x = 1", edited: true },
-        },
-      });
-
-      // Set cell runtime to already be idle
-      notebook.cellRuntime[cellId1] = {
-        ...notebook.cellRuntime[cellId1],
-        status: "idle",
-      };
-      store.set(notebookAtom, notebook);
-
-      // Mock runCells to do nothing (cells already idle)
-      vi.mocked(runCells).mockResolvedValue();
-
-      // Mock getCellContextData
-      vi.mocked(getCellContextData).mockReturnValue({
-        cellOutput: {
-          outputType: "text",
-          processedContent: "1",
-          imageUrl: null,
-          output: {},
-        },
-        consoleOutputs: null,
-        cellName: "cell1",
-      } as never);
-
-      const result = await tool.handler({}, toolContext as never);
-
-      expect(result.status).toBe("success");
-      expect(result.cellsToOutput).toBeDefined();
-    });
-  });
-});