@marimo-team/frontend 0.15.5 → 0.16.0-dev96986

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/context/__tests__/registry.test.ts

@@ -101,11 +101,13 @@ class AttachmentContextProvider extends AIContextProvider<MockContextItem> {
   readonly mentionPrefix = "@";
   readonly contextType = "attachment";
 
-  constructor(
-    private items: MockContextItem[] = [],
-    private attachments: FileUIPart[] = [],
-  ) {
+  private items: MockContextItem[];
+  private attachments: FileUIPart[];
+
+  constructor(items: MockContextItem[] = [], attachments: FileUIPart[] = []) {
     super();
+    this.items = items;
+    this.attachments = attachments;
   }
 
   getItems(): MockContextItem[] {

package/src/core/ai/context/providers/cell-output.ts

@@ -86,9 +86,10 @@ export class CellOutputContextProvider extends AIContextProvider<CellOutputConte
   readonly title = "Cell Outputs";
   readonly mentionPrefix = "@";
   readonly contextType = "cell-output";
-
-  constructor(private store: JotaiStore) {
+  private store: JotaiStore;
+  constructor(store: JotaiStore) {
     super();
+    this.store = store;
   }
 
   getItems(): CellOutputContextItem[] {

package/src/core/ai/context/providers/error.ts

@@ -64,9 +64,11 @@ export class ErrorContextProvider extends AIContextProvider<ErrorContextItem> {
   readonly title = "Errors";
   readonly mentionPrefix = "@";
   readonly contextType = "error";
+  private store: JotaiStore;
 
-  constructor(private store: JotaiStore) {
+  constructor(store: JotaiStore) {
     super();
+    this.store = store;
   }
 
   getItems(): ErrorContextItem[] {

package/src/core/ai/context/providers/file.ts

@@ -43,11 +43,16 @@ export class FileContextProvider extends AIContextProvider<FileContextItem> {
   readonly mentionPrefix = "#";
   readonly contextType = "file";
 
+  private apiRequests: EditRequests & RunRequests;
+  private config: FileSearchConfig;
+
   constructor(
-    private apiRequests: EditRequests & RunRequests,
-    private config: FileSearchConfig = DEFAULT_FILE_SEARCH_CONFIG,
+    apiRequests: EditRequests & RunRequests,
+    config: FileSearchConfig = DEFAULT_FILE_SEARCH_CONFIG,
   ) {
     super();
+    this.apiRequests = apiRequests;
+    this.config = config;
   }
 
   /**

package/src/core/ai/context/providers/tables.ts

@@ -17,9 +17,10 @@ export class TableContextProvider extends AIContextProvider<TableContextItem> {
   readonly title = "Tables";
   readonly mentionPrefix = "@";
   readonly contextType = "data";
-
-  constructor(private tablesMap: DatasetTablesMap) {
+  private tablesMap: DatasetTablesMap;
+  constructor(tablesMap: DatasetTablesMap) {
     super();
+    this.tablesMap = tablesMap;
   }
 
   getItems(): TableContextItem[] {

package/src/core/ai/context/providers/variable.ts

@@ -20,11 +20,13 @@ export class VariableContextProvider extends AIContextProvider<VariableContextIt
   readonly mentionPrefix = "@";
   readonly contextType = "variable";
 
-  constructor(
-    private variables: Variables,
-    private tablesMap: DatasetTablesMap,
-  ) {
+  private variables: Variables;
+  private tablesMap: DatasetTablesMap;
+
+  constructor(variables: Variables, tablesMap: DatasetTablesMap) {
     super();
+    this.variables = variables;
+    this.tablesMap = tablesMap;
   }
 
   getItems(): VariableContextItem[] {

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/cells/logs.ts

@@ -102,7 +102,7 @@ export function formatLogTimestamp(timestamp: number): string {
   try {
     // parse from UTC
     const date = fromUnixTime(timestamp);
-    return date.toLocaleTimeString("en-US", {
+    return date.toLocaleTimeString(undefined, {
       hour12: true,
       hour: "numeric",
       minute: "numeric",

package/src/core/codemirror/find-replace/search-highlight.ts

@@ -91,8 +91,10 @@ const HighlightMargin = 250;
 export const searchHighlighter = ViewPlugin.fromClass(
   class {
     decorations: DecorationSet;
+    readonly view: EditorView;
 
-    constructor(readonly view: EditorView) {
+    constructor(view: EditorView) {
+      this.view = view;
       this.decorations = this.highlight(view.state.field(searchState));
     }