@marimo-team/islands 0.17.8 → 0.18.1

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

@@ -1,9 +1,14 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { describe, expect, it } from "vitest";
-import { FrontendToolRegistry } from "../registry";
+import { COPILOT_MODES } from "@/core/config/config-schema";
+import { FRONTEND_TOOL_REGISTRY, FrontendToolRegistry } from "../registry";
 import { TestFrontendTool } from "../sample-tool";
 
+// Tools that have these keys in their parameters are likely not supported by Google
+// So we should be careful to add these types of schemas
+const INVALID_KEYS_IN_SCHEMA_PARAMS = new Set(["anyOf", "oneOf"]);
+
 describe("FrontendToolRegistry", () => {
   it("registers tools via constructor and supports has()", () => {
     const registry = new FrontendToolRegistry([new TestFrontendTool()]);
@@ -63,7 +68,7 @@ describe("FrontendToolRegistry", () => {
   it("returns tool schemas with expected shape and memoizes the result", () => {
     const registry = new FrontendToolRegistry([new TestFrontendTool()]);
 
-    const schemas1 = registry.getToolSchemas();
+    const schemas1 = registry.getToolSchemas("ask");
     expect(Array.isArray(schemas1)).toBe(true);
     expect(schemas1.length).toBe(1);
 
@@ -80,7 +85,40 @@ describe("FrontendToolRegistry", () => {
     expect(properties && typeof properties === "object").toBe(true);
     expect("name" in (properties ?? {})).toBe(true);
 
-    const schemas2 = registry.getToolSchemas();
+    const schemas2 = registry.getToolSchemas("ask");
     expect(schemas2).toBe(schemas1);
+
+    // Should not include tools for other modes
+    const schemas3 = registry.getToolSchemas("agent");
+    expect(schemas3.length).toBe(0);
+  });
+
+  it("tool schemas should not contain invalid keys", () => {
+    const registry = FRONTEND_TOOL_REGISTRY;
+
+    function hasInvalidKeys(obj: unknown): boolean {
+      if (obj && typeof obj === "object") {
+        for (const [key, value] of Object.entries(obj)) {
+          if (INVALID_KEYS_IN_SCHEMA_PARAMS.has(key)) {
+            return true;
+          }
+          if (
+            typeof value === "object" &&
+            value !== null &&
+            hasInvalidKeys(value)
+          ) {
+            return true;
+          }
+        }
+      }
+      return false;
+    }
+
+    for (const mode of COPILOT_MODES) {
+      const schemas = registry.getToolSchemas(mode);
+      for (const schema of schemas) {
+        expect(hasInvalidKeys(schema.parameters)).toBe(false);
+      }
+    }
   });
 });

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

@@ -30,13 +30,22 @@ const description: ToolDescription = {
   ],
   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 is specified by the position object with a "type" field:
-        { type: "notebook_end" } - Add at the end of the notebook
-        { type: "relative", cellId: "...", before: true } - Add before the specified cell (before: false for after)
-        { type: "column_end", columnIndex: 0 } - Add at the end of a specific column (0-based index)
-    - 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.
+    edit (object): The editing operation to perform with the following structure:
+      - type: One of "update_cell", "add_cell", or "delete_cell"
+      - code (string, optional): Required for "update_cell" and "add_cell" operations. The code content for the cell.
+      - position (object): Always present, with the following structure:
+        - type: One of "relative", "column_end", or "notebook_end". Not needed for "delete_cell" or "update_cell" operations.
+        - cellId (string, optional): Required for "update_cell", "delete_cell", and "add_cell" with "relative" position type
+        - columnIndex (number, optional): Required for "add_cell" with "column_end" position type (0-based index)
+        - before (boolean, optional): Required for "add_cell" with "relative" position type (true to add before, false to add after)
+
+    Operation types:
+    - update_cell: Update the code of an existing cell. Requires: code, position.cellId
+    - add_cell: Add a new cell to the notebook. Requires: code, and position based on position.type:
+        - position.type "notebook_end": Add at the end of the notebook (no additional position fields needed)
+        - position.type "relative": Add relative to an existing cell. Requires: position.cellId, position.before
+        - position.type "column_end": Add at the end of a specific column. Requires: position.columnIndex
+    - delete_cell: Delete an existing cell. Requires: position.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.
@@ -46,41 +55,17 @@ const description: ToolDescription = {
     - A result object containing standard tool metadata.`,
 };
 
-type CellPosition =
-  | { type: "relative"; cellId: CellId; before: boolean }
-  | { type: "column_end"; columnIndex: number }
-  | { type: "notebook_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.discriminatedUnion("type", [
-        z.object({
-          type: z.literal("relative"),
-          cellId: z.string() as unknown as z.ZodType<CellId>,
-          before: z.boolean(),
-        }),
-        z.object({
-          type: z.literal("column_end"),
-          columnIndex: z.number(),
-        }),
-        z.object({
-          type: z.literal("notebook_end"),
-        }),
-      ]) satisfies z.ZodType<CellPosition>,
-      code: z.string(),
+  edit: z.object({
+    type: z.enum(["update_cell", "add_cell", "delete_cell"]),
+    code: z.string().optional(),
+    position: z.object({
+      type: z.enum(["relative", "column_end", "notebook_end"]).optional(),
+      cellId: z.string().optional(),
+      columnIndex: z.number().optional(),
+      before: z.boolean().optional(),
     }),
-    z.object({
-      type: z.literal("delete_cell"),
-      cellId: z.string() as unknown as z.ZodType<CellId>,
-    }),
-  ]),
+  }),
 });
 
 type EditNotebookInput = z.infer<typeof editNotebookSchema>;
@@ -104,7 +89,24 @@ export class EditNotebookTool
 
     switch (edit.type) {
       case "update_cell": {
-        const { cellId, code } = edit;
+        const { position, code } = edit;
+        const cellId = position.cellId as CellId;
+
+        if (!cellId) {
+          throw new ToolExecutionError(
+            "Cell ID is required for update_cell",
+            "CELL_ID_REQUIRED_FOR_UPDATE_CELL",
+            true,
+            "Provide a cell ID to update",
+          );
+        } else if (!code) {
+          throw new ToolExecutionError(
+            "Code is required for update_cell",
+            "CODE_REQUIRED_FOR_UPDATE_CELL",
+            true,
+            "Provide the new code to update the cell",
+          );
+        }
 
         const notebook = store.get(notebookAtom);
         this.validateCellIdExists(cellId, notebook);
@@ -144,12 +146,41 @@ export class EditNotebookTool
         const notebook = store.get(notebookAtom);
 
         switch (position.type) {
-          case "relative":
-            this.validateCellIdExists(position.cellId, notebook);
-            notebookPosition = position.cellId;
+          case "relative": {
+            const cellId = position.cellId as CellId;
+            if (!cellId) {
+              throw new ToolExecutionError(
+                "Cell ID is required for add_cell with relative position",
+                "TOOL_ERROR",
+                true,
+                "Provide a cell ID to add the cell before",
+              );
+            }
+
+            if (position.before === undefined) {
+              throw new ToolExecutionError(
+                "Before is required for add_cell with relative position",
+                "TOOL_ERROR",
+                true,
+                `Provide a boolean value for before, true to add before cell ${cellId}, false to add after the cell`,
+              );
+            }
+
+            this.validateCellIdExists(cellId, notebook);
+            notebookPosition = cellId;
             before = position.before;
             break;
+          }
           case "column_end": {
+            if (!position.columnIndex) {
+              throw new ToolExecutionError(
+                "Column index is required for add_cell with column_end position",
+                "TOOL_ERROR",
+                true,
+                "Provide a column index to add the cell at the end of",
+              );
+            }
+
             const columnId = this.getColumnId(position.columnIndex, notebook);
             notebookPosition = { type: "__end__", columnId };
             break;
@@ -177,7 +208,17 @@ export class EditNotebookTool
         break;
       }
       case "delete_cell": {
-        const { cellId } = edit;
+        const { position } = edit;
+        const cellId = position.cellId as CellId;
+
+        if (!cellId) {
+          throw new ToolExecutionError(
+            "Cell ID is required for delete_cell",
+            "CELL_ID_REQUIRED_FOR_DELETE_CELL",
+            true,
+            "Provide a cell ID to delete",
+          );
+        }
 
         const notebook = store.get(notebookAtom);
         this.validateCellIdExists(cellId, notebook);

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

@@ -123,8 +123,11 @@ export class FrontendToolRegistry {
   }
 
   @Memoize()
-  getToolSchemas(): FrontendToolDefinition[] {
-    return [...this.tools.values()].map((tool) => ({
+  getToolSchemas(mode: CopilotMode): FrontendToolDefinition[] {
+    const tools = [...this.tools.values()].filter((tool) =>
+      tool.mode.includes(mode),
+    );
+    return tools.map((tool) => ({
       name: tool.name,
       description: formatToolDescription(tool.description),
       parameters: z.toJSONSchema(tool.schema),

package/src/core/cells/__tests__/logs.test.ts

@@ -0,0 +1,371 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { afterEach, beforeEach, describe, expect, test, vi } from "vitest";
+import type { CellMessage } from "../../kernel/messages";
+import { formatLogTimestamp, getCellLogsForMessage } from "../logs";
+
+describe("getCellLogsForMessage", () => {
+  beforeEach(() => {
+    // Mock console.log to avoid cluttering test output
+    vi.spyOn(console, "log").mockImplementation(() => {
+      // no-op
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  test("handles text/plain MIME type on stdout", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-1",
+      console: [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Hello, World!",
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0]).toEqual({
+      timestamp: 1_234_567_890,
+      level: "stdout",
+      message: "Hello, World!",
+      cellId: "cell-1",
+    });
+  });
+
+  test("handles text/plain MIME type on stderr", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-2",
+      console: [
+        {
+          mimetype: "text/plain",
+          channel: "stderr",
+          data: "Error occurred",
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0]).toEqual({
+      timestamp: 1_234_567_890,
+      level: "stderr",
+      message: "Error occurred",
+      cellId: "cell-2",
+    });
+  });
+
+  test("handles text/html MIME type and strips HTML tags", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-3",
+      console: [
+        {
+          mimetype: "text/html",
+          channel: "stdout",
+          data: '<span style="color: red;">Error: Something went wrong</span>',
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0]).toEqual({
+      timestamp: 1_234_567_890,
+      level: "stdout",
+      message: "Error: Something went wrong",
+      cellId: "cell-3",
+    });
+  });
+
+  test("handles text/html MIME type on stderr", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-4",
+      console: [
+        {
+          mimetype: "text/html",
+          channel: "stderr",
+          data: "<div><strong>Critical Error:</strong> System failure</div>",
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0]).toEqual({
+      timestamp: 1_234_567_890,
+      level: "stderr",
+      message: "Critical Error: System failure",
+      cellId: "cell-4",
+    });
+  });
+
+  test("handles application/vnd.marimo+traceback MIME type and strips HTML", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-5",
+      console: [
+        {
+          mimetype: "application/vnd.marimo+traceback",
+          channel: "marimo-error",
+          data: '<div class="traceback"><span style="color: red;">Traceback (most recent call last):</span><pre>  File "test.py", line 1</pre></div>',
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0].level).toBe("stderr"); // marimo-error should be treated as stderr
+    expect(logs[0].message).toContain("Traceback (most recent call last):");
+    expect(logs[0].message).toContain('File "test.py", line 1');
+    expect(logs[0].cellId).toBe("cell-5");
+  });
+
+  test("handles multiple console outputs with different MIME types", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-7",
+      console: [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "Plain text output",
+          timestamp: 1_234_567_890,
+        },
+        {
+          mimetype: "text/html",
+          channel: "stdout",
+          data: "<span>HTML output</span>",
+          timestamp: 1_234_567_891,
+        },
+        {
+          mimetype: "application/vnd.marimo+traceback",
+          channel: "stderr",
+          data: "<div>Traceback error</div>",
+          timestamp: 1_234_567_892,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(3);
+    expect(logs[0].message).toBe("Plain text output");
+    expect(logs[1].message).toBe("HTML output");
+    expect(logs[2].message).toBe("Traceback error");
+  });
+
+  test("uses Date.now() when timestamp is missing", () => {
+    const now = Date.now();
+    vi.spyOn(Date, "now").mockReturnValue(now);
+
+    const cellMessage: CellMessage = {
+      cell_id: "cell-8",
+      console: [
+        {
+          mimetype: "text/plain",
+          channel: "stdout",
+          data: "No timestamp",
+          // timestamp is undefined
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0].timestamp).toBe(now);
+  });
+
+  test("ignores unsupported MIME types", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-9",
+      console: [
+        {
+          mimetype: "application/json",
+          channel: "stdout",
+          data: '{"key": "value"}',
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(0);
+  });
+
+  test("ignores non-logging channels", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-10",
+      console: [
+        {
+          mimetype: "text/plain",
+          channel: "pdb" as unknown as "stdout", // Non-logging channel
+          data: "Should be ignored",
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(0);
+  });
+
+  test("returns empty array when console is null", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-11",
+      console: null as unknown as CellMessage["console"],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(0);
+  });
+
+  test("handles complex HTML with nested elements in text/html", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-12",
+      console: [
+        {
+          mimetype: "text/html",
+          channel: "stdout",
+          data: "<div><span>Nested</span> <strong>HTML</strong> <em>content</em></div>",
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0].message).toBe("Nested HTML content");
+  });
+
+  test("handles marimo-error channel as stderr level", () => {
+    const cellMessage: CellMessage = {
+      cell_id: "cell-13",
+      console: [
+        {
+          mimetype: "text/plain",
+          channel: "marimo-error",
+          data: "Internal error",
+          timestamp: 1_234_567_890,
+        },
+      ],
+      output: null,
+      status: "idle",
+      stale_inputs: null,
+      timestamp: 0,
+    };
+
+    const logs = getCellLogsForMessage(cellMessage);
+
+    expect(logs).toHaveLength(1);
+    expect(logs[0].level).toBe("stderr");
+  });
+});
+
+describe("formatLogTimestamp", () => {
+  test("formats unix timestamp correctly", () => {
+    // January 1, 2024, 12:00:00 PM UTC
+    const timestamp = 1_704_110_400;
+    const result = formatLogTimestamp(timestamp);
+
+    // The result depends on the timezone, so we just check it's not empty
+    // and contains expected time format elements
+    expect(result).toBeTruthy();
+    expect(result).toMatch(/\d+:\d+:\d+/); // Should contain time format
+  });
+
+  test("formats timestamp with AM/PM notation", () => {
+    const timestamp = 1_704_110_400; // Noon
+    const result = formatLogTimestamp(timestamp);
+
+    // Should contain AM or PM
+    expect(result).toMatch(/AM|PM/);
+  });
+
+  test("returns 'Invalid Date' for invalid timestamp", () => {
+    const invalidTimestamp = Number.NaN;
+    const result = formatLogTimestamp(invalidTimestamp);
+
+    expect(result).toBe("Invalid Date");
+  });
+
+  test("handles edge case: zero timestamp", () => {
+    const timestamp = 0; // Unix epoch
+    const result = formatLogTimestamp(timestamp);
+
+    // Should format successfully (Jan 1, 1970)
+    expect(result).toBeTruthy();
+    expect(result).toMatch(/\d+:\d+:\d+/);
+  });
+
+  test("handles recent timestamp", () => {
+    // Use a recent timestamp (seconds since epoch)
+    const timestamp = Math.floor(Date.now() / 1000);
+    const result = formatLogTimestamp(timestamp);
+
+    expect(result).toBeTruthy();
+    expect(result).toMatch(/\d+:\d+:\d+/);
+    expect(result).toMatch(/AM|PM/);
+  });
+});

package/src/core/cells/__tests__/session.test.ts

@@ -446,9 +446,6 @@ describe("notebookStateFromSession", () => {
       const session = createSession([]);
       const result = notebookStateFromSession(session, null);
 
-      expect(Logger.warn).toHaveBeenCalledWith(
-        "Session and notebook must have at least one cell",
-      );
       expect(result).toBeNull();
     });
 
@@ -456,9 +453,6 @@ describe("notebookStateFromSession", () => {
       const notebook = createNotebook([]);
       const result = notebookStateFromSession(null, notebook);
 
-      expect(Logger.warn).toHaveBeenCalledWith(
-        "Session and notebook must have at least one cell",
-      );
       expect(result).toBeNull();
     });
 
@@ -467,9 +461,6 @@ describe("notebookStateFromSession", () => {
       const notebook = createNotebook([]);
       const result = notebookStateFromSession(session, notebook);
 
-      expect(Logger.warn).toHaveBeenCalledWith(
-        "Session and notebook must have at least one cell",
-      );
       expect(result).toBeNull();
     });
   });