@marimo-team/islands 0.16.5 → 0.17.1

package/src/core/ai/tools/edit-notebook-tool.ts

@@ -0,0 +1,239 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { EditorView } from "@codemirror/view";
+import { z } from "zod";
+import { scrollAndHighlightCell } from "@/components/editor/links/cell-link";
+import {
+  type CellPosition as NotebookCellPosition,
+  type NotebookState,
+  notebookAtom,
+} from "@/core/cells/cells";
+import { CellId } from "@/core/cells/ids";
+import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
+import type { CellColumnId } from "@/utils/id-tree";
+import { stagedAICellsAtom } from "../staged-cells";
+import {
+  type AiTool,
+  type ToolDescription,
+  ToolExecutionError,
+  type ToolNotebookContext,
+  type ToolOutputBase,
+  toolOutputBaseSchema,
+} from "./base";
+import type { CopilotMode } from "./registry";
+
+const description: ToolDescription = {
+  baseDescription:
+    "Perform editing operations on the current notebook. You should prefer to create new cells unless you need to edit existing cells. Call this tool multiple times to perform multiple edits. Separate code into logical individual cells to take advantage of the notebook's reactive execution model.",
+  prerequisites: [
+    "If you are updating existing cells, you need the cellIds. If they are not known, call the lightweight_cell_map_tool to find out.",
+  ],
+  additionalInfo: `
+  Args:
+    edit (object): The editing operation to perform. Must be one of:
+    - update_cell: Update the code of an existing cell, pass CellId and the new code.
+    - add_cell: Add a new cell to the notebook. The position of the new cell is specified by the position argument.
+        Pass "end" to add the new cell at the end of the notebook. 
+        Pass { cellId: cellId, before: true } to add the new cell before the specified cell. And before: false if after the specified cell.
+        Pass { type: "end", columnIndex: number } to add the new cell at the end of a specified column index. The column index is 0-based.
+    - delete_cell: Delete an existing cell, pass CellId. For deleting cells, the user needs to accept the deletion to actually delete the cell, so you may still see the cell in the notebook on subsequent edits which is fine.
+
+    For adding code, use the following guidelines:
+    - Markdown cells: use mo.md(f"""{content}""") function to insert content.
+    - SQL cells: use mo.sql(f"""{content}""") function to insert content. If a database engine is specified, use mo.sql(f"""{content}""", engine=engine) instead.
+
+    Returns:
+    - A result object containing standard tool metadata.`,
+};
+
+type CellPosition =
+  | { cellId: CellId; before: boolean }
+  | { type: "end"; columnIndex: number }
+  | "end";
+
+const editNotebookSchema = z.object({
+  edit: z.discriminatedUnion("type", [
+    z.object({
+      type: z.literal("update_cell"),
+      cellId: z.string() as unknown as z.ZodType<CellId>,
+      code: z.string(),
+    }),
+    z.object({
+      type: z.literal("add_cell"),
+      position: z.union([
+        z.object({
+          cellId: z.string() as unknown as z.ZodType<CellId>,
+          before: z.boolean(),
+        }),
+        z.object({
+          type: z.literal("end"),
+          columnIndex: z.number(),
+        }),
+        z.literal("end"),
+      ]) satisfies z.ZodType<CellPosition>,
+      code: z.string(),
+    }),
+    z.object({
+      type: z.literal("delete_cell"),
+      cellId: z.string() as unknown as z.ZodType<CellId>,
+    }),
+  ]),
+});
+
+type EditNotebookInput = z.infer<typeof editNotebookSchema>;
+type EditOperation = EditNotebookInput["edit"];
+export type EditType = EditOperation["type"];
+
+export class EditNotebookTool
+  implements AiTool<EditNotebookInput, ToolOutputBase>
+{
+  readonly name = "edit_notebook_tool";
+  readonly description = description;
+  readonly schema = editNotebookSchema;
+  readonly outputSchema = toolOutputBaseSchema;
+  readonly mode: CopilotMode[] = ["agent"];
+
+  handler = async (
+    { edit }: EditNotebookInput,
+    toolContext: ToolNotebookContext,
+  ): Promise<ToolOutputBase> => {
+    const { addStagedCell, createNewCell, store } = toolContext;
+
+    switch (edit.type) {
+      case "update_cell": {
+        const { cellId, code } = edit;
+
+        const notebook = store.get(notebookAtom);
+        this.validateCellIdExists(cellId, notebook);
+        const editorView = this.getCellEditorView(cellId, notebook);
+
+        scrollAndHighlightCell(cellId);
+
+        // If previous code exists, we don't want to replace it, it means there is a new edit on top of the previous edit
+        // Keep the original code
+        const stagedCell = store.get(stagedAICellsAtom).get(cellId);
+        const currentCellCode = editorView.state.doc.toString();
+        const previousCode =
+          stagedCell?.type === "update_cell" ||
+          stagedCell?.type === "delete_cell"
+            ? stagedCell.previousCode
+            : currentCellCode;
+
+        addStagedCell({
+          cellId,
+          edit: { type: "update_cell", previousCode: previousCode },
+        });
+
+        updateEditorCodeFromPython(editorView, code);
+
+        break;
+      }
+      case "add_cell": {
+        const { position, code } = edit;
+
+        // By default, add the new cell to the end of the notebook
+        let notebookPosition: NotebookCellPosition = "__end__";
+        let before = false;
+        const newCellId = CellId.create();
+
+        if (typeof position === "object") {
+          const notebook = store.get(notebookAtom);
+          if ("cellId" in position) {
+            this.validateCellIdExists(position.cellId, notebook);
+            notebookPosition = position.cellId;
+            before = position.before;
+          } else if ("columnIndex" in position) {
+            const columnId = this.getColumnId(position.columnIndex, notebook);
+            notebookPosition = { type: "__end__", columnId };
+          }
+        }
+
+        createNewCell({
+          cellId: notebookPosition,
+          before,
+          code,
+          newCellId,
+        });
+
+        addStagedCell({
+          cellId: newCellId,
+          edit: { type: "add_cell" },
+        });
+
+        // Scroll into view
+        scrollAndHighlightCell(newCellId);
+
+        break;
+      }
+      case "delete_cell": {
+        const { cellId } = edit;
+
+        const notebook = store.get(notebookAtom);
+        this.validateCellIdExists(cellId, notebook);
+
+        const editorView = this.getCellEditorView(cellId, notebook);
+        const currentCellCode = editorView.state.doc.toString();
+
+        // Add to staged AICells - don't actually delete the cell yet
+        addStagedCell({
+          cellId,
+          edit: { type: "delete_cell", previousCode: currentCellCode },
+        });
+
+        scrollAndHighlightCell(cellId);
+        break;
+      }
+    }
+    return {
+      status: "success",
+      next_steps: ["If you need to perform more edits, call this tool again."],
+    };
+  };
+
+  private validateCellIdExists(cellId: CellId, notebook: NotebookState) {
+    const cellIds = notebook.cellIds;
+    if (!cellIds.getColumns().some((column) => column.idSet.has(cellId))) {
+      throw new ToolExecutionError(
+        "Cell not found",
+        "CELL_NOT_FOUND",
+        false,
+        "Check which cells exist in the notebook",
+      );
+    }
+  }
+
+  private getColumnId(
+    columnIndex: number,
+    notebook: NotebookState,
+  ): CellColumnId {
+    const cellIds = notebook.cellIds;
+    const columns = cellIds.getColumns();
+
+    if (columnIndex < 0 || columnIndex >= columns.length) {
+      throw new ToolExecutionError(
+        "Column index is out of range",
+        "COLUMN_INDEX_OUT_OF_RANGE",
+        true,
+        "Choose a column index between 0 and the number of columns in the notebook (0-based)",
+      );
+    }
+    return columns[columnIndex].id;
+  }
+
+  private getCellEditorView(
+    cellId: CellId,
+    notebook: NotebookState,
+  ): EditorView {
+    const cellHandles = notebook.cellHandles;
+    const cellHandle = cellHandles[cellId].current;
+    if (!cellHandle?.editorView) {
+      throw new ToolExecutionError(
+        "Cell editor not found",
+        "CELL_EDITOR_NOT_FOUND",
+        false,
+        "Internal error, ask the user to report this error",
+      );
+    }
+    return cellHandle.editorView;
+  }
+}

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

@@ -3,14 +3,26 @@
 import type { components } from "@marimo-team/marimo-api";
 import { Memoize } from "typescript-memoize";
 import { type ZodObject, z } from "zod";
-import type { AiTool } from "./base";
-import { TestFrontendTool } from "./sample-tool";
+import {
+  type AiTool,
+  ToolExecutionError,
+  type ToolNotebookContext,
+} from "./base";
+import { EditNotebookTool } from "./edit-notebook-tool";
+import { RunStaleCellsTool } from "./run-cells-tool";
+import { formatToolDescription } from "./utils";
 
 export type AnyZodObject = ZodObject<z.ZodRawShape>;
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type StoredTool = AiTool<any, any>;
 
+interface InvokeResult<TName> {
+  tool_name: TName;
+  result: unknown;
+  error: string | null;
+}
+
 /** should be the same as marimo/_config/config.py > CopilotMode */
 
 type ToolDefinition = components["schemas"]["ToolDefinition"];
@@ -43,7 +55,8 @@ export class FrontendToolRegistry {
   async invoke<TName extends string>(
     toolName: TName,
     rawArgs: unknown,
-  ): Promise<unknown> {
+    toolContext: ToolNotebookContext,
+  ): Promise<InvokeResult<TName>> {
     const tool = this.getToolOrThrow(toolName);
     const handler = tool.handler;
     const inputSchema = tool.schema;
@@ -59,7 +72,7 @@ export class FrontendToolRegistry {
       const args = inputResponse.data;
 
       // Call the handler
-      const rawOutput = await handler(args);
+      const rawOutput = await handler(args, toolContext);
 
       // Parse output
       const response = await outputSchema.safeParseAsync(rawOutput);
@@ -69,17 +82,37 @@ export class FrontendToolRegistry {
           `Tool ${toolName} returned invalid output: ${strError}`,
         );
       }
-      const output = response.data;
-      return output;
-    } catch (error) {
+      const result = response.data;
       return {
-        status: "error",
-        code: "TOOL_ERROR",
-        message: error instanceof Error ? error.message : String(error),
-        suggestedFix: "Try again with valid arguments.",
-        meta: {
+        tool_name: toolName,
+        result,
+        error: null,
+      };
+    } catch (error) {
+      // Handle structured errors
+      if (error instanceof ToolExecutionError) {
+        return {
+          tool_name: toolName,
+          result: null,
+          error: error.toStructuredString(),
+        };
+      }
+
+      // Handle unknown/generic errors
+      const genericError = new ToolExecutionError(
+        error instanceof Error ? error.message : String(error),
+        "TOOL_ERROR",
+        false,
+        "Check the error message and try again with valid arguments.",
+        {
           args: rawArgs,
+          errorType: error?.constructor?.name ?? typeof error,
         },
+      );
+      return {
+        tool_name: toolName,
+        result: null,
+        error: genericError.toStructuredString(),
       };
     }
   }
@@ -88,7 +121,7 @@ export class FrontendToolRegistry {
   getToolSchemas(): FrontendToolDefinition[] {
     return [...this.tools.values()].map((tool) => ({
       name: tool.name,
-      description: tool.description,
+      description: formatToolDescription(tool.description),
       parameters: z.toJSONSchema(tool.schema),
       source: "frontend",
       mode: tool.mode,
@@ -97,6 +130,6 @@ export class FrontendToolRegistry {
 }
 
 export const FRONTEND_TOOL_REGISTRY = new FrontendToolRegistry([
-  ...(import.meta.env.DEV ? [new TestFrontendTool()] : []),
-  // ADD MORE TOOLS HERE
+  new EditNotebookTool(),
+  new RunStaleCellsTool(),
 ]);

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

@@ -0,0 +1,233 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { FileUIPart } from "ai";
+import { z } from "zod";
+import { runCells } from "@/components/editor/cell/useRunCells";
+import { notebookAtom } from "@/core/cells/cells";
+import type { CellId } from "@/core/cells/ids";
+import { staleCellIds } from "@/core/cells/utils";
+import { type JotaiStore, waitFor } from "@/core/state/jotai";
+import {
+  type BaseOutput,
+  getCellContextData,
+} from "../context/providers/cell-output";
+import {
+  type AiTool,
+  type EmptyToolInput,
+  type ToolDescription,
+  type ToolNotebookContext,
+  type ToolOutputBase,
+  toolOutputBaseSchema,
+} from "./base";
+import type { CopilotMode } from "./registry";
+
+interface CellOutput {
+  consoleOutput?: string;
+  // consoleAttachments?: FileUIPart[];
+  cellOutput?: string;
+  // cellAttachments?: FileUIPart[];
+}
+
+// Must use Record instead of Map because Map serializes to JSON as {}
+interface RunStaleCellsOutput extends ToolOutputBase {
+  cellsToOutput?: Record<CellId, CellOutput | null>;
+}
+
+const description: ToolDescription = {
+  baseDescription:
+    "Run cells in the current notebook that are stale. Stale cells are cells that have been edited or newly added. You can run this tool after editing the notebook or when requested. The output of the ran cells will be returned alongside metadata.",
+};
+
+const filePartSchema = z.object({
+  type: z.literal("file"),
+  url: z.string(),
+  mediaType: z.string(),
+}) satisfies z.ZodType<FileUIPart>;
+
+export class RunStaleCellsTool
+  implements AiTool<EmptyToolInput, RunStaleCellsOutput>
+{
+  readonly name = "run_stale_cells_tool";
+  readonly description = description;
+  readonly schema = z.object({});
+  readonly outputSchema = toolOutputBaseSchema.extend({
+    cellsToOutput: z
+      .record(
+        z.string(),
+        z
+          .object({
+            consoleOutput: z.string().optional(),
+            cellOutput: z.string().optional(),
+            consoleAttachments: z.array(filePartSchema).optional(),
+            cellAttachments: z.array(filePartSchema).optional(),
+          })
+          .nullable(),
+      )
+      .optional(),
+  }) satisfies z.ZodType<RunStaleCellsOutput>;
+  readonly mode: CopilotMode[] = ["agent"];
+
+  handler = async (
+    _args: EmptyToolInput,
+    toolContext: ToolNotebookContext,
+  ): Promise<RunStaleCellsOutput> => {
+    const { prepareForRun, sendRun, store } = toolContext;
+
+    const notebook = store.get(notebookAtom);
+    const staleCells = staleCellIds(notebook);
+
+    if (staleCells.length === 0) {
+      return {
+        status: "success",
+        message: "No stale cells found.",
+      };
+    }
+
+    await runCells({
+      cellIds: staleCells,
+      sendRun: sendRun,
+      prepareForRun,
+      notebook,
+    });
+
+    // Wait for all cells to finish executing
+    const allCellsFinished = await this.waitForCellsToFinish(store, staleCells);
+    if (!allCellsFinished) {
+      return {
+        status: "success",
+        message:
+          "No output was returned because some cells have not finished executing",
+      };
+    }
+
+    // Get notebook state after cells have finished
+    const updatedNotebook = store.get(notebookAtom);
+
+    const cellsToOutput = new Map<CellId, CellOutput | null>();
+    let resultMessage = "";
+
+    for (const cellId of staleCells) {
+      const cellContextData = getCellContextData(cellId, updatedNotebook, {
+        includeConsoleOutput: true,
+      });
+
+      let cellOutputString: string | undefined;
+      // let cellAttachments: FileUIPart[] | undefined;
+      let consoleOutputString: string | undefined;
+      // let consoleAttachments: FileUIPart[] | undefined;
+
+      const cellOutput = cellContextData.cellOutput;
+      const consoleOutputs = cellContextData.consoleOutputs;
+      if (!cellOutput && !consoleOutputs) {
+        // Set null to show no output
+        cellsToOutput.set(cellId, null);
+        continue;
+      }
+
+      if (cellOutput) {
+        cellOutputString = this.formatOutputString(cellOutput);
+        // cellAttachments = await getAttachmentsForOutputs(
+        //   [cellOutput],
+        //   cellId,
+        //   cellContextData.cellName,
+        // );
+      }
+
+      if (consoleOutputs) {
+        // consoleAttachments = await getAttachmentsForOutputs(
+        //   consoleOutputs,
+        //   cellId,
+        //   cellContextData.cellName,
+        // );
+        consoleOutputString = consoleOutputs
+          .map((output) => this.formatOutputString(output))
+          .join("\n");
+        resultMessage +=
+          "Console output represents the stdout or stderr of the cell (eg. print statements).";
+      }
+
+      cellsToOutput.set(cellId, {
+        cellOutput: cellOutputString,
+        consoleOutput: consoleOutputString,
+        // cellAttachments: cellAttachments,
+        // consoleAttachments: consoleAttachments,
+      });
+    }
+
+    if (cellsToOutput.size === 0) {
+      return {
+        status: "success",
+        message: "Stale cells have been run. No output was returned.",
+      };
+    }
+
+    return {
+      status: "success",
+      cellsToOutput: Object.fromEntries(cellsToOutput),
+      message: resultMessage === "" ? undefined : resultMessage,
+      next_steps: [
+        "Review the output of the cells, if you need to make any changes, you can run this tool again, the CellId is the key of the result object if you want to edit the cell.",
+      ],
+    };
+  };
+
+  private formatOutputString(cellOutput: BaseOutput): string {
+    let outputString = "";
+    const { outputType, processedContent, imageUrl, output } = cellOutput;
+    if (outputType === "text") {
+      outputString += "Output:\n";
+      if (processedContent) {
+        outputString += processedContent;
+      } else if (typeof output.data === "string") {
+        outputString += output.data;
+      } else {
+        outputString += JSON.stringify(output.data);
+      }
+    } else if (outputType === "media") {
+      outputString += `Media Output: Contains ${output.mimetype} content`;
+      if (imageUrl) {
+        outputString += `\nImage URL: ${imageUrl}`;
+      }
+    }
+    return outputString;
+  }
+
+  /**
+   * Wait for cells to finish executing (status becomes "idle")
+   * Returns true if all cells finished executing, false if the timeout was reached
+   */
+  private async waitForCellsToFinish(
+    store: JotaiStore,
+    cellIds: CellId[],
+    timeout = 30_000,
+  ): Promise<boolean> {
+    const checkAllFinished = (
+      notebook: ReturnType<typeof notebookAtom.read>,
+    ) => {
+      return cellIds.every((cellId) => {
+        const cellRuntime = notebook.cellRuntime[cellId];
+        return (
+          cellRuntime?.status !== "running" && cellRuntime?.status !== "queued"
+        );
+      });
+    };
+
+    // If already finished, return immediately
+    if (checkAllFinished(store.get(notebookAtom))) {
+      return true;
+    }
+
+    // Wait for notebook state changes with timeout
+    try {
+      await Promise.race([
+        waitFor(notebookAtom, checkAllFinished),
+        new Promise<never>((_, reject) =>
+          setTimeout(() => reject(new Error("timeout")), timeout),
+        ),
+      ]);
+      return true;
+    } catch {
+      return false;
+    }
+  }
+}

package/src/core/ai/tools/sample-tool.ts

@@ -1,36 +1,74 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { z } from "zod";
-import type { AiTool } from "./base";
+import {
+  type AiTool,
+  type ToolDescription,
+  ToolExecutionError,
+  type ToolOutputBase,
+  toolOutputBaseSchema,
+} from "./base";
 import type { CopilotMode } from "./registry";
 
-const description = `
-Test frontend tool that returns a greeting message.
+const description: ToolDescription = {
+  baseDescription: "Test frontend tool that returns a greeting message.",
+  additionalInfo: `
+  Args:
+    - name (string): The name to include in the greeting.
 
-Args:
-- name (string): The name to include in the greeting.
-
-Returns:
-- { message: string } — The greeting message, e.g., "Hello: Alice".
-`;
+  Returns:
+    - Output with data containing the greeting message.
+  `,
+};
 
 interface Input {
   name: string;
 }
 
-interface Output {
-  message: string;
+interface GreetingData {
+  greeting: string;
+  timestamp: string;
+}
+
+interface Output extends ToolOutputBase {
+  data: GreetingData;
 }
 
-/** A sample frontend tool that returns "hello world" */
+/** A sample frontend tool that demonstrates real tool output structure */
 export class TestFrontendTool implements AiTool<Input, Output> {
   readonly name = "test_frontend_tool";
   readonly description = description;
   readonly schema = z.object({ name: z.string() });
-  readonly outputSchema = z.object({ message: z.string() });
+  readonly outputSchema = toolOutputBaseSchema.extend({
+    data: z.object({
+      greeting: z.string(),
+      timestamp: z.string(),
+    }),
+  });
   readonly mode: CopilotMode[] = ["ask"];
 
-  async handler({ name }: Input) {
-    return { message: `Hello: ${name}` };
+  async handler({ name }: Input): Promise<Output> {
+    // Example: Validate input and throw ToolExecutionError on invalid data
+    if (!name.trim()) {
+      throw new ToolExecutionError(
+        "Name cannot be empty",
+        "INVALID_INPUT",
+        false,
+        "Please provide a non-empty name",
+        { field: "name", received: name },
+      );
+    }
+
+    return {
+      status: "success",
+      data: {
+        greeting: `Hello: ${name}`,
+        timestamp: new Date().toISOString(),
+      },
+      next_steps: [
+        "You can now proceed with your next task",
+        "Try calling another tool if needed",
+      ],
+    };
   }
 }

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

@@ -0,0 +1,23 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { ToolDescription } from "./base";
+
+export function formatToolDescription(description: ToolDescription): string {
+  let result = description.baseDescription;
+  if (description.whenToUse) {
+    result += `\n\n## When to use:\n- ${description.whenToUse.join("\n- ")}`;
+  }
+  if (description.avoidIf) {
+    result += `\n\n## Avoid if:\n- ${description.avoidIf.join("\n- ")}`;
+  }
+  if (description.prerequisites) {
+    result += `\n\n## Prerequisites:\n- ${description.prerequisites.join("\n- ")}`;
+  }
+  if (description.sideEffects) {
+    result += `\n\n## Side effects:\n- ${description.sideEffects.join("\n- ")}`;
+  }
+  if (description.additionalInfo) {
+    result += `\n\n## Additional info:\n- ${description.additionalInfo}`;
+  }
+  return result;
+}

package/src/core/cache/requests.ts

@@ -0,0 +1,5 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { atom } from "jotai";
+import type { CacheInfoFetched } from "../kernel/messages";
+
+export const cacheInfoAtom = atom<CacheInfoFetched | null>(null);