@marimo-team/islands 0.17.0 → 0.17.2

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

@@ -0,0 +1,450 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { getDefaultStore } from "jotai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+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();
+
+    cellId1 = "cell-1" as CellId;
+    cellId2 = "cell-2" as CellId;
+    cellId3 = "cell-3" as CellId;
+
+    // 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("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();
+    });
+  });
+});

package/src/core/ai/tools/__tests__/utils.test.ts

@@ -0,0 +1,87 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import type { ToolDescription } from "../base";
+import { formatToolDescription } from "../utils";
+
+describe("formatToolDescription", () => {
+  it("formats a basic description only", () => {
+    const description: ToolDescription = {
+      baseDescription: "This is a simple tool",
+    };
+
+    expect(formatToolDescription(description)).toMatchInlineSnapshot(
+      `"This is a simple tool"`,
+    );
+  });
+
+  it("formats description with whenToUse", () => {
+    const description: ToolDescription = {
+      baseDescription: "Edit notebook cells",
+      whenToUse: [
+        "When user requests code changes",
+        "When refactoring is needed",
+      ],
+    };
+
+    expect(formatToolDescription(description)).toMatchInlineSnapshot(`
+      "Edit notebook cells
+
+      ## When to use:
+      - When user requests code changes
+      - When refactoring is needed"
+    `);
+  });
+
+  it("formats description with all fields", () => {
+    const description: ToolDescription = {
+      baseDescription: "A comprehensive tool",
+      whenToUse: ["Use case 1", "Use case 2"],
+      avoidIf: ["Avoid case 1", "Avoid case 2"],
+      prerequisites: ["Prerequisite 1", "Prerequisite 2"],
+      sideEffects: ["Side effect 1", "Side effect 2"],
+      additionalInfo: "Some extra context and notes.",
+    };
+
+    expect(formatToolDescription(description)).toMatchInlineSnapshot(`
+      "A comprehensive tool
+
+      ## When to use:
+      - Use case 1
+      - Use case 2
+
+      ## Avoid if:
+      - Avoid case 1
+      - Avoid case 2
+
+      ## Prerequisites:
+      - Prerequisite 1
+      - Prerequisite 2
+
+      ## Side effects:
+      - Side effect 1
+      - Side effect 2
+
+      ## Additional info:
+      - Some extra context and notes."
+    `);
+  });
+
+  it("formats description with only some fields", () => {
+    const description: ToolDescription = {
+      baseDescription: "A selective tool",
+      avoidIf: ["Don't use in production"],
+      sideEffects: ["Modifies state"],
+    };
+
+    expect(formatToolDescription(description)).toMatchInlineSnapshot(`
+      "A selective tool
+
+      ## Avoid if:
+      - Don't use in production
+
+      ## Side effects:
+      - Modifies state"
+    `);
+  });
+});

package/src/core/ai/tools/base.ts

@@ -1,6 +1,11 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { z } from "zod";
+import type { CreateNewCellAction } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import type { RunRequest } from "@/core/network/types";
+import type { JotaiStore } from "@/core/state/jotai";
+import type { Edit } from "../staged-cells";
 import type { CopilotMode } from "./registry";
 
 /**
@@ -53,6 +58,9 @@ export class ToolExecutionError extends Error {
   }
 }
 
+// Can be any input
+export type EmptyToolInput = unknown;
+
 /**
  * Base interface for tool output responses.
  * Mirrors the SuccessResult dataclass from marimo/_ai/_tools/types.py
@@ -83,17 +91,32 @@ export const toolOutputBaseSchema = z.object({
   meta: z.record(z.string(), z.unknown()).optional(),
 });
 
-/**
- * Contract for a frontend tool.
- *
- * Implementations can be plain objects or classes. The registry consumes this
- * interface without caring about the underlying implementation.
- */
+export interface ToolDescription {
+  baseDescription: string;
+  whenToUse?: string[];
+  avoidIf?: string[];
+  prerequisites?: string[];
+  sideEffects?: string[];
+  additionalInfo?: string;
+}
+
+/** Utility functions for tools to interact with the notebook */
+export interface ToolNotebookContext {
+  store: JotaiStore;
+  addStagedCell: (payload: { cellId: CellId; edit: Edit }) => void;
+  createNewCell: (payload: CreateNewCellAction) => void;
+  prepareForRun: (payload: { cellId: CellId }) => void;
+  sendRun: (request: RunRequest) => Promise<null>;
+}
+
 export interface AiTool<TIn, TOut> {
   name: string;
-  description: string;
+  description: ToolDescription;
   schema: z.ZodType<TIn>;
   outputSchema: z.ZodType<TOut>;
   mode: CopilotMode[];
-  handler: (args: TIn) => TOut | Promise<TOut>;
+  handler: (
+    args: TIn,
+    toolContext: ToolNotebookContext,
+  ) => TOut | Promise<TOut>;
 }