@marimo-team/frontend 0.16.0 → 0.16.1

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

@@ -10,7 +10,12 @@ 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 { maybeAddMarimoImport } from "../cells/add-missing-import";
+import {
+  type CreateNewCellAction,
+  cellHandleAtom,
+  useCellActions,
+} from "../cells/cells";
 import type { LanguageAdapterType } from "../codemirror/language/types";
 import { updateEditorCodeFromPython } from "../codemirror/language/utils";
 import type { JotaiStore } from "../state/jotai";
@@ -113,6 +118,8 @@ export function useStagedCells(store: JotaiStore) {
         cellCreationStream.current = new CellCreationStream(
           createStagedCell,
           updateStagedCell,
+          addStagedCell,
+          createNewCell,
         );
         break;
       case "text-delta":
@@ -168,13 +175,20 @@ class CellCreationStream {
 
   private onCreateCell: (code: string) => CellId;
   private onUpdateCell: (opts: UpdateStagedCellAction) => void;
+  private addStagedCell: (payload: { cellId: CellId }) => void;
+  private createNewCell: (opts: CreateNewCellAction) => void;
+  private hasMarimoImport = false;
 
   constructor(
     onCreateCell: (code: string) => CellId,
     onUpdateCell: (opts: UpdateStagedCellAction) => void,
+    addStagedCell: (payload: { cellId: CellId }) => void,
+    createNewCell: (opts: CreateNewCellAction) => void,
   ) {
     this.onCreateCell = onCreateCell;
     this.onUpdateCell = onUpdateCell;
+    this.addStagedCell = addStagedCell;
+    this.createNewCell = createNewCell;
   }
 
   stream(chunk: TextDeltaChunk) {
@@ -187,6 +201,7 @@ class CellCreationStream {
     // 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) {
+        this.addMarimoImport(cell.language);
         const existingCell = this.createdCells[idx];
         this.createdCells[idx] = { ...existingCell, cell };
         this.onUpdateCell({
@@ -201,6 +216,24 @@ class CellCreationStream {
     }
   }
 
+  /** Add a marimo import if the cell is SQL or Markdown and we haven't added it yet. */
+  private addMarimoImport(language: LanguageAdapterType) {
+    if (this.hasMarimoImport || language === "python") {
+      return;
+    }
+
+    const cellId = maybeAddMarimoImport({
+      autoInstantiate: false,
+      createNewCell: this.createNewCell,
+      fromCellId: this.createdCells[0]?.cellId,
+      before: true,
+    });
+    if (cellId) {
+      this.addStagedCell({ cellId });
+    }
+    this.hasMarimoImport = true;
+  }
+
   stop() {
     // Clear all state
     this.buffer = "";

package/src/core/cells/__tests__/add-missing-import.test.ts

@@ -1,18 +1,35 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import { createStore } from "jotai";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { MockRequestClient } from "@/__mocks__/requests";
+import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import type { Variables } from "@/core/variables/types";
-import { CollapsibleTree, MultiColumn } from "@/utils/id-tree";
-import { maybeAddMissingImport } from "../add-missing-import";
-import { type NotebookState, notebookAtom } from "../cells";
+import {
+  maybeAddAltairImport,
+  maybeAddMarimoImport,
+  maybeAddMissingImport,
+} from "../add-missing-import";
+import { notebookAtom } from "../cells";
 import type { CellId } from "../ids";
-import type { CellData } from "../types";
 
-const CELL_IDS = new MultiColumn([CollapsibleTree.from(["1" as CellId])]);
+// Mock the getRequestClient function
+const mockRequestClient = MockRequestClient.create();
+vi.mock("@/core/network/requests", () => ({
+  getRequestClient: () => mockRequestClient,
+}));
+
+const Cell1 = "1" as CellId;
+const Cell2 = "2" as CellId;
 
 describe("maybeAddMissingImport", () => {
+  beforeEach(() => {
+    store.set(variablesAtom, {} as Variables);
+    store.set(notebookAtom, MockNotebook.notebookState({ cellData: {} }));
+  });
+
   it("should not add an import if the variable is already in the variables state", () => {
     const appStore = createStore();
     appStore.set(variablesAtom, { mo: {} } as Variables);
@@ -36,14 +53,14 @@ describe("maybeAddMissingImport", () => {
     (code) => {
       const appStore = createStore();
       appStore.set(variablesAtom, {} as Variables);
-      appStore.set(notebookAtom, {
-        cellData: {
-          ["1" as CellId]: {
-            code: code,
-          } as CellData,
-        },
-        cellIds: CELL_IDS,
-      } as NotebookState);
+      appStore.set(
+        notebookAtom,
+        MockNotebook.notebookState({
+          cellData: {
+            [Cell1]: { code: code },
+          },
+        }),
+      );
       const onAddImport = vi.fn();
       maybeAddMissingImport({
         moduleName: "marimo",
@@ -58,14 +75,14 @@ describe("maybeAddMissingImport", () => {
   it("should add an import if the variable is not in the variables state and the import statement does not exist in the notebook", () => {
     const appStore = createStore();
     appStore.set(variablesAtom, {} as Variables);
-    appStore.set(notebookAtom, {
-      cellData: {
-        ["1" as CellId]: {
-          code: "mo.md('hello')",
-        } as CellData,
-      },
-      cellIds: CELL_IDS,
-    } as NotebookState);
+    appStore.set(
+      notebookAtom,
+      MockNotebook.notebookState({
+        cellData: {
+          [Cell1]: { code: "mo.md('hello')" },
+        },
+      }),
+    );
     const onAddImport = vi.fn();
     maybeAddMissingImport({
       moduleName: "marimo",
@@ -75,4 +92,32 @@ describe("maybeAddMissingImport", () => {
     });
     expect(onAddImport).toHaveBeenCalledWith("import marimo as mo");
   });
+
+  it("should not create a new cell if import already exists due to skipIfCodeExists", () => {
+    const addImports = [maybeAddMarimoImport, maybeAddAltairImport];
+    for (const addImport of addImports) {
+      store.set(variablesAtom, {} as Variables);
+      store.set(
+        notebookAtom,
+        MockNotebook.notebookState({
+          cellData: {
+            [Cell1]: { code: "import marimo as mo" },
+            [Cell2]: { code: "import altair as alt" },
+          },
+        }),
+      );
+
+      const createNewCell = vi.fn();
+      const result = addImport({
+        autoInstantiate: false,
+        createNewCell,
+        fromCellId: Cell1,
+        before: false,
+      });
+
+      // Should not create a new cell since the import already exists
+      expect(createNewCell).not.toHaveBeenCalled();
+      expect(result).toBeNull();
+    }
+  });
 });

package/src/core/cells/add-missing-import.ts

@@ -47,27 +47,40 @@ export function maybeAddMissingImport({
   return true;
 }
 
+/**
+ * Adds a marimo import to the notebook if not already present.
+ * @param autoInstantiate Whether to automatically run the cell.
+ * @param createNewCell The function to create a new cell.
+ * @param fromCellId The cell to add the import to.
+ * @param before Whether to add the import before or after the cell.
+ *
+ * Returns the ID of the new cell if added, otherwise null.
+ */
 export function maybeAddMarimoImport({
   autoInstantiate,
   createNewCell,
   fromCellId,
+  before,
 }: {
   autoInstantiate: boolean;
   createNewCell: CellActions["createNewCell"];
   fromCellId?: CellId | null;
-}): boolean {
+  before?: boolean;
+}): CellId | null {
   const client = getRequestClient();
-  return maybeAddMissingImport({
+  let newCellId: CellId | null = null;
+  const added = maybeAddMissingImport({
     moduleName: "marimo",
     variableName: "mo",
     onAddImport: (importStatement) => {
-      const newCellId = CellId.create();
+      newCellId = CellId.create();
       createNewCell({
         cellId: fromCellId ?? "__end__",
-        before: false,
+        before: before ?? false,
         code: importStatement,
         lastCodeRun: autoInstantiate ? importStatement : undefined,
         newCellId: newCellId,
+        skipIfCodeExists: true,
         autoFocus: false,
       });
       if (autoInstantiate) {
@@ -78,6 +91,7 @@ export function maybeAddMarimoImport({
       }
     },
   });
+  return added ? newCellId : null;
 }
 
 export function maybeAddAltairImport({
@@ -88,19 +102,21 @@ export function maybeAddAltairImport({
   autoInstantiate: boolean;
   createNewCell: CellActions["createNewCell"];
   fromCellId?: CellId | null;
-}): boolean {
+}): CellId | null {
   const client = getRequestClient();
-  return maybeAddMissingImport({
+  let newCellId: CellId | null = null;
+  const added = maybeAddMissingImport({
     moduleName: "altair",
     variableName: "alt",
     onAddImport: (importStatement) => {
-      const newCellId = CellId.create();
+      newCellId = CellId.create();
       createNewCell({
         cellId: fromCellId ?? "__end__",
         before: false,
         code: importStatement,
         lastCodeRun: autoInstantiate ? importStatement : undefined,
         newCellId: newCellId,
+        skipIfCodeExists: true,
         autoFocus: false,
       });
       if (autoInstantiate) {
@@ -111,4 +127,5 @@ export function maybeAddAltairImport({
       }
     },
   });
+  return added ? newCellId : null;
 }

package/src/core/cells/cells.ts

@@ -145,6 +145,31 @@ export function initialNotebookState(): NotebookState {
   });
 }
 
+export interface CreateNewCellAction {
+  /** The target cell ID to create a new cell relative to. Can be:
+   * - A CellId string for an existing cell
+   * - "__end__" to append at the end of the first column
+   * - {type: "__end__", columnId} to append at the end of a specific column
+   */
+  cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
+  /** Whether to insert before (true) or after (false) the target cell */
+  before: boolean;
+  /** Initial code content for the new cell */
+  code?: string;
+  /** The last executed code for the new cell */
+  lastCodeRun?: string;
+  /** Timestamp of the last execution */
+  lastExecutionTime?: number;
+  /** Optional custom ID for the new cell. Auto-generated if not provided */
+  newCellId?: CellId;
+  /** Whether to focus the new cell after creation */
+  autoFocus?: boolean;
+  /** If true, skip creation if code already exists */
+  skipIfCodeExists?: boolean;
+  /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
+  hideCode?: boolean;
+}
+
 /**
  * Actions and reducer for the notebook state.
  */
@@ -154,33 +179,7 @@ const {
   useActions,
   valueAtom: notebookAtom,
 } = createReducerAndAtoms(initialNotebookState, {
-  createNewCell: (
-    state,
-    action: {
-      /** The target cell ID to create a new cell relative to. Can be:
-       * - A CellId string for an existing cell
-       * - "__end__" to append at the end of the first column
-       * - {type: "__end__", columnId} to append at the end of a specific column
-       */
-      cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
-      /** Whether to insert before (true) or after (false) the target cell */
-      before: boolean;
-      /** Initial code content for the new cell */
-      code?: string;
-      /** The last executed code for the new cell */
-      lastCodeRun?: string;
-      /** Timestamp of the last execution */
-      lastExecutionTime?: number;
-      /** Optional custom ID for the new cell. Auto-generated if not provided */
-      newCellId?: CellId;
-      /** Whether to focus the new cell after creation */
-      autoFocus?: boolean;
-      /** If true, skip creation if code already exists */
-      skipIfCodeExists?: boolean;
-      /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
-      hideCode?: boolean;
-    },
-  ) => {
+  createNewCell: (state, action: CreateNewCellAction) => {
     const {
       cellId,
       before,

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));
     }
 

package/src/core/codemirror/language/LanguageAdapters.ts

@@ -1,20 +1,26 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { once } from "@/utils/once";
 import { MarkdownLanguageAdapter } from "./languages/markdown";
 import { PythonLanguageAdapter } from "./languages/python";
 import { SQLLanguageAdapter } from "./languages/sql/sql";
 import type { LanguageAdapter, LanguageAdapterType } from "./types";
 
+// Create cached instances
+const createPythonAdapter = once(() => new PythonLanguageAdapter());
+const createMarkdownAdapter = once(() => new MarkdownLanguageAdapter());
+const createSqlAdapter = once(() => new SQLLanguageAdapter());
+
 export const LanguageAdapters: Record<LanguageAdapterType, LanguageAdapter> = {
   // Getters to prevent circular dependencies
   get python() {
-    return new PythonLanguageAdapter();
+    return createPythonAdapter();
   },
   get markdown() {
-    return new MarkdownLanguageAdapter();
+    return createMarkdownAdapter();
   },
   get sql() {
-    return new SQLLanguageAdapter();
+    return createSqlAdapter();
   },
 };
 

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

@@ -18,13 +18,19 @@ import { getLSPDocument } from "./utils";
 
 class Snapshotter {
   private documentVersion = 0;
+  private readonly getNotebookCode: () => {
+    cellIds: CellId[];
+    codes: Record<CellId, string>;
+  };
 
   constructor(
-    private readonly getNotebookCode: () => {
+    getNotebookCode: () => {
       cellIds: CellId[];
       codes: Record<CellId, string>;
     },
-  ) {}
+  ) {
+    this.getNotebookCode = getNotebookCode;
+  }
 
   /**
    * Map from the global document version to the cell id and version.

package/src/core/codemirror/readonly/__tests__/extension.test.ts

@@ -9,7 +9,7 @@ import { connectionAtom } from "../../../network/connection";
 import { dynamicReadonly, isEditorReadonly } from "../extension";
 
 function makeStoreWithConnection(
-  state: WebSocketState.CONNECTING | WebSocketState.OPEN,
+  state: typeof WebSocketState.CONNECTING | typeof WebSocketState.OPEN,
 ) {
   const store = createStore();
   store.set(connectionAtom, { state });

package/src/core/codemirror/rtc/loro/awareness.ts

@@ -202,13 +202,25 @@ export const createSelectionLayer = (): Extension =>
  * Renders a blinking cursor to indicate the cursor of another user.
  */
 export class RemoteCursorMarker implements LayerMarker {
+  private left: number;
+  private top: number;
+  private height: number;
+  private name: string;
+  private colorClassName: string;
+
   constructor(
-    private left: number,
-    private top: number,
-    private height: number,
-    private name: string,
-    private colorClassName: string,
-  ) {}
+    left: number,
+    top: number,
+    height: number,
+    name: string,
+    colorClassName: string,
+  ) {
+    this.left = left;
+    this.top = top;
+    this.height = height;
+    this.name = name;
+    this.colorClassName = colorClassName;
+  }
 
   draw(): HTMLElement {
     const elt = document.createElement("div");
@@ -353,16 +365,30 @@ export interface CursorPosition {
 
 export class AwarenessPlugin implements PluginValue {
   sub: Subscription;
+  public view: EditorView;
+  public doc: LoroDoc;
+  public user: UserState;
+  public awareness: Awareness<AwarenessState>;
+  private getTextFromDoc: (doc: LoroDoc) => LoroText;
+  private scopeId: ScopeId;
+  private getUserId?: () => Uid;
 
   constructor(
-    public view: EditorView,
-    public doc: LoroDoc,
-    public user: UserState,
-    public awareness: Awareness<AwarenessState>,
-    private getTextFromDoc: (doc: LoroDoc) => LoroText,
-    private scopeId: ScopeId,
-    private getUserId?: () => Uid,
+    view: EditorView,
+    doc: LoroDoc,
+    user: UserState,
+    awareness: Awareness<AwarenessState>,
+    getTextFromDoc: (doc: LoroDoc) => LoroText,
+    scopeId: ScopeId,
+    getUserId?: () => Uid,
   ) {
+    this.view = view;
+    this.doc = doc;
+    this.user = user;
+    this.awareness = awareness;
+    this.getTextFromDoc = getTextFromDoc;
+    this.scopeId = scopeId;
+    this.getUserId = getUserId;
     this.sub = this.doc.subscribe((e) => {
       if (e.by === "local") {
         // update remote cursor position
@@ -435,12 +461,21 @@ export class AwarenessPlugin implements PluginValue {
 }
 export class RemoteAwarenessPlugin implements PluginValue {
   _awarenessListener?: AwarenessListener;
+  public view: EditorView;
+  public doc: LoroDoc;
+  public awareness: Awareness<AwarenessState>;
+  private scopeId: ScopeId;
+
   constructor(
-    public view: EditorView,
-    public doc: LoroDoc,
-    public awareness: Awareness<AwarenessState>,
-    private scopeId: ScopeId,
+    view: EditorView,
+    doc: LoroDoc,
+    awareness: Awareness<AwarenessState>,
+    scopeId: ScopeId,
   ) {
+    this.view = view;
+    this.doc = doc;
+    this.awareness = awareness;
+    this.scopeId = scopeId;
     const listener: AwarenessListener = async (arg, origin) => {
       if (origin === "local") {
         return;

package/src/core/codemirror/rtc/loro/sync.ts

@@ -40,13 +40,21 @@ export const loroSyncAnnotation = Annotation.define();
 export class LoroSyncPluginValue implements PluginValue {
   sub?: Subscription;
   private isInitDispatch = false;
+  private view: EditorView;
+  private doc: LoroDoc;
+  private docPath: string[];
+  private getTextFromDoc: (doc: LoroDoc) => LoroText;
 
   constructor(
-    private view: EditorView,
-    private doc: LoroDoc,
-    private docPath: string[],
-    private getTextFromDoc: (doc: LoroDoc) => LoroText,
+    view: EditorView,
+    doc: LoroDoc,
+    docPath: string[],
+    getTextFromDoc: (doc: LoroDoc) => LoroText,
   ) {
+    this.view = view;
+    this.doc = doc;
+    this.docPath = docPath;
+    this.getTextFromDoc = getTextFromDoc;
     this.sub = doc.subscribe(this.onRemoteUpdate);
     Promise.resolve().then(() => {
       this.isInitDispatch = true;

package/src/core/config/config-schema.ts

@@ -144,6 +144,7 @@ export const UserConfigSchema = z
             }
             return width;
           }),
+        locale: z.string().nullable().optional(),
         reference_highlighting: z.boolean().default(false),
       })
       .passthrough()

package/src/core/config/config.ts

@@ -82,6 +82,10 @@ export const editorFontSizeAtom = atom<number>((get) => {
   return get(resolvedMarimoConfigAtom).display.code_editor_font_size;
 });
 
+export const localeAtom = atom<string | null | undefined>((get) => {
+  return get(resolvedMarimoConfigAtom).display.locale;
+});
+
 export function isAiEnabled(config: UserConfig) {
   return (
     Boolean(config.ai?.models?.chat_model) ||

package/src/core/hotkeys/hotkeys.ts

@@ -452,10 +452,13 @@ export class HotkeyProvider implements IHotkeyProvider {
     return new HotkeyProvider(DEFAULT_HOT_KEY, { platform });
   }
 
+  private hotkeys: Record<HotkeyAction, Hotkey>;
+
   constructor(
-    private hotkeys: Record<HotkeyAction, Hotkey>,
+    hotkeys: Record<HotkeyAction, Hotkey>,
     options: HotkeyProviderOptions = {},
   ) {
+    this.hotkeys = hotkeys;
     this.platform = options.platform ?? resolvePlatform();
     this.mod = this.platform === "mac" ? "Cmd" : "Ctrl";
   }
@@ -503,13 +506,14 @@ export class HotkeyProvider implements IHotkeyProvider {
 }
 
 export class OverridingHotkeyProvider extends HotkeyProvider {
+  private readonly overrides: Partial<Record<HotkeyAction, string | undefined>>;
+
   constructor(
-    private readonly overrides: Partial<
-      Record<HotkeyAction, string | undefined>
-    >,
+    overrides: Partial<Record<HotkeyAction, string | undefined>>,
     options: HotkeyProviderOptions = {},
   ) {
     super(DEFAULT_HOT_KEY, options);
+    this.overrides = overrides;
   }
 
   override getHotkey(action: HotkeyAction): ResolvedHotkey {