@marimo-team/frontend 0.15.5 → 0.16.0

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

@@ -0,0 +1,356 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { renderHook } from "@testing-library/react";
+import { getDefaultStore } from "jotai";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { CellId } from "@/core/cells/ids";
+import { updateEditorCodeFromPython } from "../../codemirror/language/utils";
+import {
+  stagedAICellsAtom,
+  useStagedCells,
+  visibleForTesting,
+} from "../staged-cells";
+
+const { createActions, reducer, initialState } = visibleForTesting;
+
+// Mock the dependencies
+const mockCreateNewCell = vi.fn();
+const mockUpdateCellEditor = vi.fn();
+const mockDeleteCellCallback = vi.fn();
+
+// Mock cell handle with editor view
+const mockCellHandle = {
+  current: {
+    editorViewOrNull: {
+      dispatch: vi.fn(),
+    },
+  },
+};
+
+vi.mock("../../cells/cells", () => ({
+  useCellActions: () => ({
+    createNewCell: mockCreateNewCell,
+    updateCellEditor: mockUpdateCellEditor,
+  }),
+  cellHandleAtom: vi.fn(() => ({
+    read: vi.fn(() => mockCellHandle),
+  })),
+}));
+
+vi.mock("@/components/editor/cell/useDeleteCell", () => ({
+  useDeleteCellCallback: () => mockDeleteCellCallback,
+}));
+
+vi.mock("../../codemirror/language/utils", () => ({
+  updateEditorCodeFromPython: vi.fn(),
+}));
+
+// Mock CellId.create
+vi.mock("@/core/cells/ids", () => ({
+  CellId: {
+    create: vi.fn(),
+  },
+}));
+
+describe("staged-cells", () => {
+  let store: ReturnType<typeof getDefaultStore>;
+  let cellId1: CellId;
+  let cellId2: CellId;
+
+  beforeEach(() => {
+    store = getDefaultStore();
+    cellId1 = "cell-1" as CellId;
+    cellId2 = "cell-2" as CellId;
+
+    // Reset mocks
+    vi.clearAllMocks();
+
+    // Reset the atom state
+    store.set(stagedAICellsAtom, new Set<CellId>());
+  });
+
+  describe("reducer and actions", () => {
+    it("should initialize with empty map", () => {
+      const state = initialState();
+      expect(state).toEqual(new Set());
+    });
+
+    it("should add cell IDs", () => {
+      let state = initialState();
+      state = reducer(state, {
+        type: "addStagedCell",
+        payload: { cellId: cellId1 },
+      });
+      state = reducer(state, {
+        type: "addStagedCell",
+        payload: { cellId: cellId2 },
+      });
+
+      expect(state.has(cellId1)).toBe(true);
+      expect(state.has(cellId2)).toBe(true);
+    });
+
+    it("should remove cell IDs", () => {
+      const state = new Set([cellId1, cellId2]);
+      const newState = reducer(state, {
+        type: "removeStagedCell",
+        payload: cellId1,
+      });
+
+      expect(newState.has(cellId1)).toBe(false);
+      expect(newState.has(cellId2)).toBe(true);
+    });
+
+    it("should clear all cell IDs", () => {
+      const state = new Set([cellId1, cellId2]);
+      const newState = reducer(state, {
+        type: "clearStagedCells",
+        payload: undefined,
+      });
+
+      expect(newState).toEqual(new Set());
+    });
+
+    it("should not mutate original state", () => {
+      const state = new Set([cellId1]);
+      const originalSize = state.size;
+
+      reducer(state, {
+        type: "addStagedCell",
+        payload: { cellId: cellId2 },
+      });
+
+      expect(state.size).toBe(originalSize);
+      expect(state.has(cellId1)).toBe(true);
+      expect(state.has(cellId2)).toBe(false);
+    });
+
+    it("should create action functions", () => {
+      const mockDispatch = vi.fn();
+      const actions = createActions(mockDispatch);
+
+      expect(typeof actions.addStagedCell).toBe("function");
+      expect(typeof actions.removeStagedCell).toBe("function");
+      expect(typeof actions.clearStagedCells).toBe("function");
+    });
+
+    it("should initialize atom with empty map", () => {
+      const state = store.get(stagedAICellsAtom);
+      expect(state).toEqual(new Set());
+    });
+  });
+
+  describe("useStagedCells hook", () => {
+    it("should create a staged cell with code", () => {
+      const { result } = renderHook(() => useStagedCells(store));
+      const testCode = "print('hello world')";
+
+      // Mock CellId.create to return a predictable ID
+      const mockCellId = "mock-cell-id" as CellId;
+      vi.mocked(CellId.create).mockReturnValue(mockCellId);
+
+      const returnedCellId = result.current.createStagedCell(testCode);
+
+      expect(returnedCellId).toBe(mockCellId);
+      expect(mockCreateNewCell).toHaveBeenCalledWith({
+        cellId: "__end__",
+        code: testCode,
+        before: false,
+        newCellId: mockCellId,
+      });
+    });
+
+    it("should delete a staged cell", () => {
+      const { result } = renderHook(() => useStagedCells(store));
+      const testCellId = "test-cell-id" as CellId;
+
+      result.current.deleteStagedCell(testCellId);
+
+      expect(mockDeleteCellCallback).toHaveBeenCalledWith({
+        cellId: testCellId,
+      });
+    });
+
+    it("should delete all staged cells when none exist", () => {
+      const { result } = renderHook(() => useStagedCells(store));
+
+      // Should not throw when no cells exist
+      expect(() => result.current.deleteAllStagedCells()).not.toThrow();
+      expect(mockDeleteCellCallback).not.toHaveBeenCalled();
+    });
+
+    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 { result } = renderHook(() => useStagedCells(store));
+      result.current.deleteAllStagedCells();
+
+      expect(mockDeleteCellCallback).toHaveBeenCalledTimes(2);
+      expect(mockDeleteCellCallback).toHaveBeenCalledWith({ cellId: cellId1 });
+      expect(mockDeleteCellCallback).toHaveBeenCalledWith({ cellId: cellId2 });
+
+      // Verify cells were cleared from the atom
+      const state = store.get(stagedAICellsAtom);
+      expect(state).toEqual(new Set());
+    });
+  });
+
+  it("should add staged cell", () => {
+    const { result } = renderHook(() => useStagedCells(store));
+
+    result.current.addStagedCell({ cellId: cellId1 });
+
+    // Check that the cell was added to the atom
+    const state = store.get(stagedAICellsAtom);
+    expect(state.has(cellId1)).toBe(true);
+  });
+
+  it("should remove staged cell", () => {
+    const { result } = renderHook(() => useStagedCells(store));
+
+    // First add cells
+    result.current.addStagedCell({ cellId: cellId1 });
+    result.current.addStagedCell({ cellId: cellId2 });
+
+    // 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);
+  });
+
+  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 });
+
+    // Then clear all
+    result.current.clearStagedCells();
+
+    // Check that no cells remain
+    const state = store.get(stagedAICellsAtom);
+    expect(state).toEqual(new Set());
+  });
+
+  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);
+
+    const createdCellId = result.current.createStagedCell("test code");
+
+    // Verify it was created and added
+    expect(createdCellId).toBe(mockCellId);
+    expect(mockCreateNewCell).toHaveBeenCalled();
+
+    let state = store.get(stagedAICellsAtom);
+    expect(state.has(mockCellId)).toBe(true);
+
+    // 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);
+  });
+});
+
+describe("onStream", () => {
+  let store: ReturnType<typeof getDefaultStore>;
+  beforeEach(() => {
+    store = getDefaultStore();
+  });
+
+  it("should create a cell creation stream", () => {
+    const { result } = renderHook(() => useStagedCells(store));
+    result.current.onStream({ type: "text-start", id: "test-id" });
+
+    // No cell or cell update should have been called
+    expect(mockCreateNewCell).not.toHaveBeenCalled();
+    expect(mockUpdateCellEditor).not.toHaveBeenCalled();
+  });
+
+  it("should not create cells when text-delta is received and no stream has been created", () => {
+    const { result } = renderHook(() => useStagedCells(store));
+    result.current.onStream({
+      type: "text-delta",
+      id: "test-id",
+      delta: "test-delta",
+    });
+
+    // No cell or cell update should have been called
+    expect(mockCreateNewCell).not.toHaveBeenCalled();
+    expect(mockUpdateCellEditor).not.toHaveBeenCalled();
+  });
+
+  it("should create cells when text-delta is received and a stream has been created", () => {
+    const { result } = renderHook(() => useStagedCells(store));
+    result.current.onStream({ type: "text-start", id: "test-id" });
+
+    // Mock CellId.create to return a predictable ID
+    const mockCellId = "mock-cell-id" as CellId;
+    vi.mocked(CellId.create).mockReturnValue(mockCellId);
+
+    result.current.onStream({
+      type: "text-delta",
+      id: "test-id",
+      delta: "some code",
+    });
+
+    expect(mockCreateNewCell).toHaveBeenCalledWith({
+      cellId: "__end__",
+      code: "some code",
+      before: false,
+      newCellId: "mock-cell-id",
+    });
+  });
+
+  it("should handle delta chunks", () => {
+    const { result } = renderHook(() => useStagedCells(store));
+    result.current.onStream({ type: "text-start", id: "test-id" });
+
+    const mockCellId = "mock-cell-id" as CellId;
+    vi.mocked(CellId.create).mockReturnValue(mockCellId);
+
+    result.current.onStream({
+      type: "text-delta",
+      id: "test-id",
+      delta: "``",
+    });
+
+    expect(mockCreateNewCell).toHaveBeenCalledWith({
+      cellId: "__end__",
+      code: "``",
+      before: false,
+      newCellId: "mock-cell-id",
+    });
+
+    result.current.onStream({
+      type: "text-delta",
+      id: "test-id",
+      delta: "```python\nsome code",
+    });
+
+    // Now the cell is recognized and only some code is seen
+    expect(vi.mocked(updateEditorCodeFromPython)).toHaveBeenCalledWith(
+      mockCellHandle.current.editorViewOrNull,
+      "some code",
+    );
+
+    result.current.onStream({
+      type: "text-delta",
+      id: "test-id",
+      delta: "\n```",
+    });
+  });
+});

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

@@ -0,0 +1,208 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { UIMessageChunk } from "ai";
+import { useRef } from "react";
+import {
+  type AiCompletion,
+  codeToCells,
+} from "@/components/editor/ai/completion-utils";
+import { useDeleteCellCallback } from "@/components/editor/cell/useDeleteCell";
+import { CellId } from "@/core/cells/ids";
+import { createReducerAndAtoms } from "@/utils/createReducer";
+import { Logger } from "@/utils/Logger";
+import { cellHandleAtom, useCellActions } from "../cells/cells";
+import type { LanguageAdapterType } from "../codemirror/language/types";
+import { updateEditorCodeFromPython } from "../codemirror/language/utils";
+import type { JotaiStore } from "../state/jotai";
+
+/**
+ * Cells that are staged for AI completion
+ * They function similarly to cells in the notebook, but they can be deleted or accepted by the user.
+ * We only track one set of staged cells at a time.
+ */
+
+const initialState = (): Set<CellId> => {
+  return new Set();
+};
+
+const {
+  valueAtom: stagedAICellsAtom,
+  useActions: useStagedAICellsActions,
+  createActions,
+  reducer,
+} = createReducerAndAtoms(initialState, {
+  addStagedCell: (state, action: { cellId: CellId }) => {
+    const { cellId } = action;
+    return new Set([...state, cellId]);
+  },
+  removeStagedCell: (state, cellId: CellId) => {
+    return new Set([...state].filter((id) => id !== cellId));
+  },
+  clearStagedCells: () => {
+    return initialState();
+  },
+});
+
+interface UpdateStagedCellAction {
+  cellId: CellId;
+  code: string;
+  language?: LanguageAdapterType;
+}
+
+/**
+ * Helper functions to create and delete staged cells.
+ */
+export function useStagedCells(store: JotaiStore) {
+  const { addStagedCell, removeStagedCell, clearStagedCells } =
+    useStagedAICellsActions();
+  const { createNewCell } = useCellActions();
+  const deleteCellCallback = useDeleteCellCallback();
+
+  const cellCreationStream = useRef<CellCreationStream | null>(null);
+
+  const createStagedCell = (code: string): CellId => {
+    const newCellId = CellId.create();
+    addStagedCell({ cellId: newCellId });
+    createNewCell({
+      cellId: "__end__",
+      code,
+      before: false,
+      newCellId: newCellId,
+    });
+    return newCellId;
+  };
+
+  const updateStagedCell = (opts: UpdateStagedCellAction) => {
+    const { cellId, code } = opts;
+    const stagedAICells = store.get(stagedAICellsAtom);
+
+    if (!stagedAICells.has(cellId)) {
+      Logger.error("Staged cell not found", { cellId });
+      return;
+    }
+
+    const cellHandle = store.get(cellHandleAtom(cellId));
+    const editorView = cellHandle?.current?.editorViewOrNull;
+    if (!editorView) {
+      Logger.error("Editor for this cell not found", { cellId });
+      return;
+    }
+    // TODO: Update the language
+    updateEditorCodeFromPython(editorView, code);
+  };
+
+  // Delete a staged cell and the corresponding cell in the notebook.
+  const deleteStagedCell = (cellId: CellId) => {
+    removeStagedCell(cellId);
+    deleteCellCallback({ cellId });
+  };
+
+  // Delete all staged cells and the corresponding cells in the notebook.
+  const deleteAllStagedCells = () => {
+    const stagedAICells = store.get(stagedAICellsAtom);
+    for (const cellId of stagedAICells) {
+      deleteCellCallback({ cellId });
+    }
+    clearStagedCells();
+  };
+
+  const onStream = (chunk: UIMessageChunk) => {
+    switch (chunk.type) {
+      case "text-start":
+        // Create stream
+        cellCreationStream.current = new CellCreationStream(
+          createStagedCell,
+          updateStagedCell,
+        );
+        break;
+      case "text-delta":
+        if (!cellCreationStream.current) {
+          Logger.error("Cell creation stream not found");
+          return;
+        }
+        cellCreationStream.current.stream(chunk);
+        break;
+      case "text-end":
+      case "finish":
+        if (!cellCreationStream.current) {
+          Logger.error("Cell creation stream not found");
+          return;
+        }
+        cellCreationStream.current.stop();
+        break;
+      default:
+        Logger.error("Unknown chunk type", { chunk });
+    }
+  };
+
+  return {
+    createStagedCell,
+    updateStagedCell,
+    addStagedCell,
+    removeStagedCell,
+    clearStagedCells,
+    deleteStagedCell,
+    deleteAllStagedCells,
+    onStream,
+  };
+}
+
+export { stagedAICellsAtom };
+export const visibleForTesting = {
+  createActions,
+  reducer,
+  initialState,
+  useStagedAICellsActions,
+};
+
+type TextDeltaChunk = Extract<UIMessageChunk, { type: "text-delta" }>;
+
+interface CreatedCell {
+  cellId: CellId;
+  cell: AiCompletion;
+}
+
+class CellCreationStream {
+  private createdCells: CreatedCell[] = [];
+  private buffer = "";
+
+  private onCreateCell: (code: string) => CellId;
+  private onUpdateCell: (opts: UpdateStagedCellAction) => void;
+
+  constructor(
+    onCreateCell: (code: string) => CellId,
+    onUpdateCell: (opts: UpdateStagedCellAction) => void,
+  ) {
+    this.onCreateCell = onCreateCell;
+    this.onUpdateCell = onUpdateCell;
+  }
+
+  stream(chunk: TextDeltaChunk) {
+    const delta = chunk.delta;
+    this.buffer += delta;
+    const completionCells = codeToCells(this.buffer);
+
+    // As incoming chunks are appended to the buffer,
+    // we parse the buffer into cells and determine which parts correspond to which cell.
+    // For each parsed cell, we either update an existing staged cell or create a new one.
+    for (const [idx, cell] of completionCells.entries()) {
+      if (idx < this.createdCells.length) {
+        const existingCell = this.createdCells[idx];
+        this.createdCells[idx] = { ...existingCell, cell };
+        this.onUpdateCell({
+          cellId: existingCell.cellId,
+          code: cell.code,
+          language: cell.language,
+        });
+      } else {
+        const newCellId = this.onCreateCell(cell.code);
+        this.createdCells.push({ cellId: newCellId, cell });
+      }
+    }
+  }
+
+  stop() {
+    // Clear all state
+    this.buffer = "";
+  }
+}

package/src/core/cells/cells.ts

@@ -1623,7 +1623,7 @@ const cellDataAtom = atomFamily((cellId: CellId) =>
 const cellRuntimeAtom = atomFamily((cellId: CellId) =>
   atom((get) => get(notebookAtom).cellRuntime[cellId]),
 );
-const cellHandleAtom = atomFamily((cellId: CellId) =>
+export const cellHandleAtom = atomFamily((cellId: CellId) =>
   atom((get) => get(notebookAtom).cellHandles[cellId]),
 );
 /**

package/src/core/codemirror/lsp/federated-lsp.ts

@@ -4,7 +4,7 @@ import type { LanguageServerPlugin } from "@marimo-team/codemirror-languageserve
 import type * as LSP from "vscode-languageserver-protocol";
 import { Objects } from "@/utils/objects";
 import type { ILanguageServerClient } from "./types";
-import { getLSPDocument } from "./utils.ts";
+import { getLSPDocument } from "./utils";
 
 function removeFalseyValues<T extends object>(obj: T): T {
   return Objects.filter(obj, (value) => value !== false && value !== null) as T;

package/src/core/islands/main.ts

@@ -103,7 +103,7 @@ export async function initialize() {
   // Consume messages from the kernel
   IslandsPyodideBridge.INSTANCE.consumeMessages((message) => {
     const msg = jsonParseWithSpecialChar(message);
-    switch (msg.op) {
+    switch (msg.data.op) {
       case "banner":
       case "missing-package-alert":
       case "installing-package-alert":
@@ -185,7 +185,7 @@ export async function initialize() {
       case "reconnected":
         return;
       default:
-        logNever(msg);
+        logNever(msg.data);
     }
   });
 

package/src/core/kernel/messages.ts

@@ -40,21 +40,17 @@ export type SQLTableListPreview =
   OperationMessageData<"sql-table-list-preview">;
 export type SecretKeysResult = OperationMessageData<"secret-keys-result">;
 export type StartupLogs = OperationMessageData<"startup-logs">;
-export type MessageOperation = schemas["KnownUnions"]["operation"];
+export type CellMessage = OperationMessageData<"cell-op">;
+export type Capabilities = OperationMessageData<"kernel-ready">["capabilities"];
 
-export type OperationMessageType = MessageOperation["op"];
-export type OperationMessage = {
-  [Type in OperationMessageType]: {
-    op: Type;
-    data: Omit<Extract<MessageOperation, { op: Type }>, "op">;
-  };
-}[OperationMessageType];
+export type MessageOperationUnion = schemas["KnownUnions"]["operation"];
 
-export type CellMessage = OperationMessageData<"cell-op">;
+export type OperationMessageType = MessageOperationUnion["op"];
+export type OperationMessage = {
+  data: MessageOperationUnion;
+};
 
 export type OperationMessageData<T extends OperationMessageType> = Omit<
-  Extract<MessageOperation, { op: T }>,
+  Extract<MessageOperationUnion, { op: T }>,
   "op"
 >;
-
-export type Capabilities = OperationMessageData<"kernel-ready">["capabilities"];

package/src/core/saving/__tests__/filename.test.ts

@@ -0,0 +1,37 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { EDGE_CASE_FILENAMES } from "../../../__tests__/mocks";
+import { Paths } from "../../../utils/paths";
+
+describe("filename handling logic", () => {
+  it.each(EDGE_CASE_FILENAMES)(
+    "should extract basename correctly for document title: %s",
+    (filename) => {
+      const basename = Paths.basename(filename);
+      expect(basename).toBe(filename); // Since no path separator
+    },
+  );
+
+  it("should handle full paths with unicode filenames", () => {
+    EDGE_CASE_FILENAMES.forEach((filename) => {
+      const fullPath = `/path/to/${filename}`;
+
+      const basename = Paths.basename(fullPath);
+      expect(basename).toBe(filename);
+    });
+  });
+
+  it("should handle document title setting with unicode", () => {
+    EDGE_CASE_FILENAMES.forEach((filename) => {
+      const originalTitle = document.title;
+
+      // In case this does any conversions, we want to simulate reading/writing the title
+      document.title = filename;
+      expect(document.title).toBe(filename);
+
+      // Restore
+      document.title = originalTitle;
+    });
+  });
+});