@marimo-team/frontend 0.22.1-dev3 → 0.22.1-dev32

package/src/components/ui/textarea.tsx

@@ -6,8 +6,7 @@ import { useDebounceControlledState } from "@/hooks/useDebounce";
 import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 
-export interface TextareaProps
-  extends React.TextareaHTMLAttributes<HTMLTextAreaElement> {
+export interface TextareaProps extends React.TextareaHTMLAttributes<HTMLTextAreaElement> {
   bottomAdornment?: React.ReactNode;
 }
 

package/src/components/ui/toast.tsx

@@ -1,8 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import * as ToastPrimitives from "@radix-ui/react-toast";
 import { cva, type VariantProps } from "class-variance-authority";
 import { X } from "lucide-react";
+import { Toast as ToastPrimitives } from "radix-ui";
 import * as React from "react";
 
 import { cn } from "@/utils/cn";

package/src/components/ui/toggle.tsx

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import * as TogglePrimitive from "@radix-ui/react-toggle";
 import { cva, type VariantProps } from "class-variance-authority";
+import { Toggle as TogglePrimitive } from "radix-ui";
 import * as React from "react";
 
 import { cn } from "@/utils/cn";

package/src/components/ui/tooltip.tsx

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import * as TooltipPrimitive from "@radix-ui/react-tooltip";
+import { Tooltip as TooltipPrimitive } from "radix-ui";
 import * as React from "react";
 import { StyleNamespace } from "@/theme/namespace";
 import { cn } from "@/utils/cn";

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

@@ -34,8 +34,7 @@ export interface CellOutputData {
 // For the context provider
 // Currently, we enforce that cellOutput is present.
 interface CellOutputContextData
-  extends CellOutputData,
-    Record<string, unknown> {
+  extends CellOutputData, Record<string, unknown> {
   cellOutput: BaseOutput;
 }
 

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

@@ -72,9 +72,10 @@ type EditNotebookInput = z.infer<typeof editNotebookSchema>;
 type EditOperation = EditNotebookInput["edit"];
 export type EditType = EditOperation["type"];
 
-export class EditNotebookTool
-  implements AiTool<EditNotebookInput, ToolOutputBase>
-{
+export class EditNotebookTool implements AiTool<
+  EditNotebookInput,
+  ToolOutputBase
+> {
   readonly name = "edit_notebook_tool";
   readonly description = description;
   readonly schema = editNotebookSchema;

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

@@ -50,9 +50,10 @@ const filePartSchema = z.object({
   mediaType: z.string(),
 }) satisfies z.ZodType<FileUIPart>;
 
-export class RunStaleCellsTool
-  implements AiTool<EmptyToolInput, RunStaleCellsOutput>
-{
+export class RunStaleCellsTool implements AiTool<
+  EmptyToolInput,
+  RunStaleCellsOutput
+> {
   readonly name = "run_stale_cells_tool";
   readonly description = description;
   readonly schema = z.object({});

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

@@ -48,28 +48,29 @@ describe("maybeAddMissingImport", () => {
     "import marimo as mo\nimport marimo as mo",
     "import   marimo    as   mo",
   ];
-  it.each(
-    VALID_IMPORTS,
-  )("should not add an import if the import statement already exists in the notebook", (code) => {
-    const appStore = createStore();
-    appStore.set(variablesAtom, {} as Variables);
-    appStore.set(
-      notebookAtom,
-      MockNotebook.notebookState({
-        cellData: {
-          [Cell1]: { code: code },
-        },
-      }),
-    );
-    const onAddImport = vi.fn();
-    maybeAddMissingImport({
-      moduleName: "marimo",
-      variableName: "mo",
-      onAddImport,
-      appStore,
-    });
-    expect(onAddImport).not.toHaveBeenCalled();
-  });
+  it.each(VALID_IMPORTS)(
+    "should not add an import if the import statement already exists in the notebook",
+    (code) => {
+      const appStore = createStore();
+      appStore.set(variablesAtom, {} as Variables);
+      appStore.set(
+        notebookAtom,
+        MockNotebook.notebookState({
+          cellData: {
+            [Cell1]: { code: code },
+          },
+        }),
+      );
+      const onAddImport = vi.fn();
+      maybeAddMissingImport({
+        moduleName: "marimo",
+        variableName: "mo",
+        onAddImport,
+        appStore,
+      });
+      expect(onAddImport).not.toHaveBeenCalled();
+    },
+  );
 
   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();

package/src/core/cells/__tests__/apply-transaction.test.ts

@@ -18,6 +18,7 @@ import { python } from "@codemirror/lang-python";
 import { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import type { CellHandle } from "@/components/editor/notebook-cell";
 import { adaptiveLanguageConfiguration } from "@/core/codemirror/language/extension";
 import { OverridingHotkeyProvider } from "@/core/hotkeys/hotkeys";
@@ -121,7 +122,7 @@ describe("applyTransactionChanges edge cases", () => {
     apply([
       {
         type: "create-cell",
-        cellId: "new-cell",
+        cellId: cellId("new-cell"),
         code: "configured",
         name: "",
         config: { hide_code: true, disabled: true, column: 1 },
@@ -141,12 +142,12 @@ describe("applyTransactionChanges edge cases", () => {
     apply([
       {
         type: "create-cell",
-        cellId: "new-cell",
+        cellId: cellId("new-cell"),
         code: "new",
         name: "",
         config: {},
       },
-      { type: "move-cell", cellId: "new-cell", before: a },
+      { type: "move-cell", cellId: cellId("new-cell"), before: a },
     ]);
     expect(pretty(state)).toMatchInlineSnapshot(`
       "
@@ -162,12 +163,12 @@ describe("applyTransactionChanges edge cases", () => {
     apply([
       {
         type: "create-cell",
-        cellId: "new-cell",
+        cellId: cellId("new-cell"),
         code: "initial",
         name: "",
         config: {},
       },
-      { type: "set-code", cellId: "new-cell", code: "updated" },
+      { type: "set-code", cellId: cellId("new-cell"), code: "updated" },
     ]);
     expect(pretty(state)).toMatchInlineSnapshot(`
       "
@@ -182,12 +183,12 @@ describe("applyTransactionChanges edge cases", () => {
     apply([
       {
         type: "create-cell",
-        cellId: "ephemeral",
+        cellId: cellId("ephemeral"),
         code: "tmp",
         name: "",
         config: {},
       },
-      { type: "delete-cell", cellId: "ephemeral" },
+      { type: "delete-cell", cellId: cellId("ephemeral") },
     ]);
     expect(pretty(state)).toMatchInlineSnapshot(`
       "
@@ -228,7 +229,7 @@ describe("applyTransactionChanges edge cases", () => {
   it("move-cell with missing after anchor falls back to end", () => {
     setup("a", "b");
     const [a] = state.cellIds.inOrderIds;
-    apply([{ type: "move-cell", cellId: a, after: "nonexistent" as CellId }]);
+    apply([{ type: "move-cell", cellId: a, after: cellId("nonexistent") }]);
     expect(pretty(state)).toMatchInlineSnapshot(`
       "
       1: 'b'
@@ -240,7 +241,7 @@ describe("applyTransactionChanges edge cases", () => {
   it("move-cell with missing before anchor falls back to start", () => {
     setup("a", "b");
     const [, b] = state.cellIds.inOrderIds;
-    apply([{ type: "move-cell", cellId: b, before: "nonexistent" as CellId }]);
+    apply([{ type: "move-cell", cellId: b, before: cellId("nonexistent") }]);
     expect(pretty(state)).toMatchInlineSnapshot(`
       "
       1: 'b'
@@ -254,8 +255,8 @@ describe("applyTransactionChanges edge cases", () => {
     apply([
       {
         type: "move-cell",
-        cellId: "nonexistent" as CellId,
-        after: "0" as CellId,
+        cellId: cellId("nonexistent"),
+        after: cellId("0"),
       },
     ]);
     expect(pretty(state)).toMatchInlineSnapshot(`

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

@@ -43,19 +43,20 @@ describe("outputIsLoading", () => {
 });
 
 describe("outputIsStale", () => {
-  it.each(
-    STATUSES,
-  )("should return true if the cell is edited and status is %s", (status) => {
-    const cell = {
-      status: status,
-      staleInputs: true,
-      output: null,
-      runStartTimestamp: null,
-      interrupted: false,
-    };
-    const edited = true;
-    expect(outputIsStale(cell, edited)).toBe(true);
-  });
+  it.each(STATUSES)(
+    "should return true if the cell is edited and status is %s",
+    (status) => {
+      const cell = {
+        status: status,
+        staleInputs: true,
+        output: null,
+        runStartTimestamp: null,
+        interrupted: false,
+      };
+      const edited = true;
+      expect(outputIsStale(cell, edited)).toBe(true);
+    },
+  );
 
   it("should return true if the cell is loading", () => {
     const cell = {

package/src/core/cells/document-changes.ts

@@ -425,7 +425,7 @@ export function fromDocumentChanges(
             cellId,
             before,
             code: change.code,
-            newCellId: change.cellId as CellId,
+            newCellId: change.cellId,
             autoFocus: false,
             hideCode: change.config?.hide_code ?? false,
           },
@@ -433,14 +433,14 @@ export function fromDocumentChanges(
         if (change.name) {
           actions.push({
             type: "updateCellName",
-            payload: { cellId: change.cellId as CellId, name: change.name },
+            payload: { cellId: change.cellId, name: change.name },
           });
         }
         if (change.config?.disabled != null || change.config?.column != null) {
           actions.push({
             type: "updateCellConfig",
             payload: {
-              cellId: change.cellId as CellId,
+              cellId: change.cellId,
               config: {
                 ...(change.config.disabled != null && {
                   disabled: change.config.disabled,
@@ -460,7 +460,7 @@ export function fromDocumentChanges(
       case "delete-cell":
         actions.push({
           type: "deleteCell",
-          payload: { cellId: change.cellId as CellId },
+          payload: { cellId: change.cellId },
         });
         break;
 
@@ -471,7 +471,7 @@ export function fromDocumentChanges(
       // missing. No-ops if the cell itself doesn't exist.
       case "move-cell": {
         const ids = [...getCurrentCellIds()];
-        const cellId = change.cellId as CellId;
+        const cellId = change.cellId;
         const idx = ids.indexOf(cellId);
         if (idx < 0) {
           break;
@@ -507,7 +507,7 @@ export function fromDocumentChanges(
       case "reorder-cells":
         actions.push({
           type: "setCellIds",
-          payload: { cellIds: change.cellIds as CellId[] },
+          payload: { cellIds: change.cellIds },
         });
         break;
 
@@ -518,7 +518,7 @@ export function fromDocumentChanges(
         actions.push({
           type: "setCellCodes",
           payload: {
-            ids: [change.cellId as CellId],
+            ids: [change.cellId],
             codes: [change.code],
             codeIsStale: true,
           },
@@ -530,7 +530,7 @@ export function fromDocumentChanges(
       case "set-name":
         actions.push({
           type: "updateCellName",
-          payload: { cellId: change.cellId as CellId, name: change.name },
+          payload: { cellId: change.cellId, name: change.name },
         });
         break;
 
@@ -542,7 +542,7 @@ export function fromDocumentChanges(
         actions.push({
           type: "updateCellConfig",
           payload: {
-            cellId: change.cellId as CellId,
+            cellId: change.cellId,
             config: {
               ...(change.hideCode != null && { hide_code: change.hideCode }),
               ...(change.disabled != null && { disabled: change.disabled }),

package/src/core/cells/logs.ts

@@ -98,7 +98,7 @@ const CellLogLogger = {
           ? "red"
           : "orange";
     const status = payload.level.toUpperCase();
-    /** biome-ignore lint/suspicious/noConsole: for debugging */
+    // oxlint-disable-next-line no-console -- intentional debug logging
     console.log(
       `%c[${status}]`,
       `color:${color}; padding:2px 0; border-radius:2px; font-weight:bold`,

package/src/core/codemirror/cells/__tests__/extensions.test.ts

@@ -18,24 +18,22 @@ function createTransaction(spec: TransactionSpec) {
 }
 
 describe("shouldAutorunMarkdownUpdate", () => {
-  it.each([
-    "input.type",
-    "delete.backward",
-    "undo",
-    "redo",
-  ])("accepts local %s transactions", (userEvent) => {
-    const transaction = createTransaction({
-      changes: { from: 0, insert: "#" },
-      annotations: [Transaction.userEvent.of(userEvent)],
-    });
+  it.each(["input.type", "delete.backward", "undo", "redo"])(
+    "accepts local %s transactions",
+    (userEvent) => {
+      const transaction = createTransaction({
+        changes: { from: 0, insert: "#" },
+        annotations: [Transaction.userEvent.of(userEvent)],
+      });
 
-    expect(
-      shouldAutorunMarkdownUpdate({
-        docChanged: transaction.docChanged,
-        transactions: [transaction],
-      }),
-    ).toBe(true);
-  });
+      expect(
+        shouldAutorunMarkdownUpdate({
+          docChanged: transaction.docChanged,
+          transactions: [transaction],
+        }),
+      ).toBe(true);
+    },
+  );
 
   it("ignores formatting changes", () => {
     const transaction = createTransaction({

package/src/core/codemirror/copilot/__tests__/transport.test.ts

@@ -365,7 +365,7 @@ describe("LazyWebsocketTransport", () => {
       expect(delegate.connect).toHaveBeenCalled();
     });
 
-    it("should clamp timeout to maxTimeoutMs", async () => {
+    it("should use maxTimeoutMs as default when no timeout is provided", async () => {
       const transport = new LazyWebsocketTransport({
         getWsUrl: mockGetWsUrl,
         waitForReady: mockWaitForReady,
@@ -376,12 +376,29 @@ describe("LazyWebsocketTransport", () => {
       await transport.connect();
 
       const data: any = { method: "test", params: [] };
-      await transport.sendData(data, 10_000);
+      await transport.sendData(data, undefined);
 
       const delegate = (transport as any).delegate;
       expect(delegate.sendData).toHaveBeenCalledWith(data, 5000);
     });
 
+    it("should respect caller-provided timeout without clamping", async () => {
+      const transport = new LazyWebsocketTransport({
+        getWsUrl: mockGetWsUrl,
+        waitForReady: mockWaitForReady,
+        showError: mockShowError,
+        maxTimeoutMs: 5000,
+      });
+
+      await transport.connect();
+
+      const data: any = { method: "test", params: [] };
+      await transport.sendData(data, 30_000);
+
+      const delegate = (transport as any).delegate;
+      expect(delegate.sendData).toHaveBeenCalledWith(data, 30_000);
+    });
+
     it("should throw error if reconnection fails", async () => {
       const connectionError = new Error("Connection failed");
       (WebSocketTransport as any).mockImplementation(function (this: any) {
@@ -407,6 +424,115 @@ describe("LazyWebsocketTransport", () => {
     });
   });
 
+  describe("onReconnect", () => {
+    it("should call onReconnect after close + sendData reconnection", async () => {
+      const transport = new LazyWebsocketTransport({
+        getWsUrl: mockGetWsUrl,
+        waitForReady: mockWaitForReady,
+        showError: mockShowError,
+      });
+
+      const onReconnect = vi.fn().mockResolvedValue(undefined);
+      transport.onReconnect = onReconnect;
+
+      // Initial connect
+      await transport.connect();
+      expect(onReconnect).not.toHaveBeenCalled();
+
+      // Close the transport (simulates failure handling)
+      transport.close();
+
+      // sendData should reconnect and call onReconnect
+      const data: any = { method: "test", params: [] };
+      await transport.sendData(data, 5000);
+
+      expect(onReconnect).toHaveBeenCalledTimes(1);
+    });
+
+    it("should NOT call onReconnect on initial sendData connection", async () => {
+      const transport = new LazyWebsocketTransport({
+        getWsUrl: mockGetWsUrl,
+        waitForReady: mockWaitForReady,
+        showError: mockShowError,
+      });
+
+      const onReconnect = vi.fn().mockResolvedValue(undefined);
+      transport.onReconnect = onReconnect;
+
+      // sendData without prior connect — this is the initial connection
+      const data: any = { method: "test", params: [] };
+      await transport.sendData(data, 5000);
+
+      expect(onReconnect).not.toHaveBeenCalled();
+    });
+
+    it("should call onReconnect after tryConnect failure + retry via sendData", async () => {
+      let connectAttempt = 0;
+      (WebSocketTransport as any).mockImplementation(function (this: any) {
+        this.connect = vi.fn().mockImplementation(() => {
+          connectAttempt++;
+          // First 2 attempts fail (retries=2 exhausted), then succeed on next sendData
+          if (connectAttempt <= 2) {
+            return Promise.reject(new Error("Connection failed"));
+          }
+          return Promise.resolve(undefined);
+        });
+        this.close = vi.fn();
+        this.sendData = vi.fn().mockResolvedValue({ result: "success" });
+        this.subscribe = vi.fn();
+        this.unsubscribe = vi.fn();
+      });
+
+      const transport = new LazyWebsocketTransport({
+        getWsUrl: mockGetWsUrl,
+        waitForReady: mockWaitForReady,
+        showError: mockShowError,
+        retries: 2,
+        retryDelayMs: 10,
+      });
+
+      const onReconnect = vi.fn().mockResolvedValue(undefined);
+      transport.onReconnect = onReconnect;
+
+      // Initial connect fails (all retries exhausted)
+      await expect(transport.connect()).rejects.toThrow("Connection failed");
+      expect(onReconnect).not.toHaveBeenCalled();
+
+      // Now sendData reconnects successfully
+      const data: any = { method: "test", params: [] };
+      await transport.sendData(data, 5000);
+
+      expect(onReconnect).toHaveBeenCalledTimes(1);
+    });
+
+    it("should propagate onReconnect rejection and allow retry", async () => {
+      const transport = new LazyWebsocketTransport({
+        getWsUrl: mockGetWsUrl,
+        waitForReady: mockWaitForReady,
+        showError: mockShowError,
+      });
+
+      const reconnectError = new Error("Re-initialization failed");
+      const onReconnect = vi.fn().mockRejectedValueOnce(reconnectError);
+      transport.onReconnect = onReconnect;
+
+      await transport.connect();
+      transport.close();
+
+      const data: any = { method: "test", params: [] };
+      await expect(transport.sendData(data, 5000)).rejects.toThrow(
+        "Re-initialization failed",
+      );
+      expect(onReconnect).toHaveBeenCalledTimes(1);
+
+      // needsReInitialization should still be true, so a subsequent retry
+      // will attempt onReconnect again
+      onReconnect.mockResolvedValueOnce(undefined);
+      await transport.sendData(data, 5000);
+      expect(onReconnect).toHaveBeenCalledTimes(2);
+    });
+  });
+
   describe("close", () => {
     it("should close delegate and clear it", async () => {
       const transport = new LazyWebsocketTransport({

package/src/core/codemirror/copilot/client.ts

@@ -37,13 +37,20 @@ export const createWSTransport = once(() => {
 export const getCopilotClient = once(() => {
   const userConfig = store.get(resolvedMarimoConfigAtom);
   const copilotSettings = userConfig.ai?.github?.copilot_settings ?? {};
+  const transport = createWSTransport();
 
-  return new CopilotLanguageServerClient({
+  const client = new CopilotLanguageServerClient({
     rootUri: FILE_URI,
     workspaceFolders: null,
-    transport: createWSTransport(),
+    transport,
     copilotSettings,
   });
+
+  // Re-run the LSP initialize handshake when the transport reconnects
+  // after a close or connection failure.
+  transport.onReconnect = () => client.reInitialize();
+
+  return client;
 });
 
 export function copilotServer() {

package/src/core/codemirror/copilot/language-server.ts

@@ -86,6 +86,17 @@ export class CopilotLanguageServerClient extends LanguageServerClient {
     });
   }
 
+  /**
+   * Re-run the LSP initialize handshake and send configuration.
+   * Called by the transport's onReconnect callback after reconnecting.
+   */
+  async reInitialize(): Promise<void> {
+    logger.log("#reInitialize: Re-initializing LSP connection");
+    this.initializePromise = this.initialize();
+    await this.initializePromise;
+    await this.sendConfiguration();
+  }
+
   private async sendConfiguration() {
     const settings = this.copilotSettings;
     // Skip if no settings are provided

package/src/core/codemirror/copilot/transport.ts

@@ -36,7 +36,8 @@ export interface LazyWebsocketTransportOptions {
   retryDelayMs?: number;
 
   /**
-   * Maximum timeout for sendData operations in milliseconds.
+   * Default timeout for sendData operations in milliseconds.
+   * Used when the caller does not provide an explicit timeout.
    * @default 5000
    */
   maxTimeoutMs?: number;
@@ -60,6 +61,13 @@ export class LazyWebsocketTransport extends Transport {
   private delegate: WebSocketTransport | undefined;
   private pendingSubscriptions: Subscription[] = [];
   private readonly options: Required<LazyWebsocketTransportOptions>;
+  private needsReInitialization = false;
+
+  /**
+   * Callback invoked after the transport reconnects following a close or connection failure.
+   * Used by the LSP client to re-run the initialize handshake on the new connection.
+   */
+  onReconnect?: () => Promise<void>;
 
   constructor(options: LazyWebsocketTransportOptions) {
     super();
@@ -157,11 +165,11 @@ export class LazyWebsocketTransport extends Transport {
         );
         if (attempt === this.options.retries) {
           this.delegate = undefined;
+          this.needsReInitialization = true;
           // Show error toast on final retry
           this.options.showError(
             "GitHub Copilot Connection Error",
-            "Failed to connect to GitHub Copilot. Please check your settings and try again.\n\n" +
-              prettyError(error),
+            `Failed to connect to GitHub Copilot. Please check your settings and try again.\n\n${prettyError(error)}`,
           );
           throw error;
         }
@@ -183,6 +191,7 @@ export class LazyWebsocketTransport extends Transport {
   override close(): void {
     this.delegate?.close();
     this.delegate = undefined;
+    this.needsReInitialization = true;
   }
 
   override async sendData(
@@ -202,6 +211,25 @@ export class LazyWebsocketTransport extends Transport {
           "Unable to connect to GitHub Copilot. Please check your settings and try again.",
         );
       }
+
+      // Re-run LSP initialization handshake after reconnecting
+      if (this.needsReInitialization && this.onReconnect) {
+        Logger.log(
+          "Copilot#sendData: Re-initializing LSP after reconnection...",
+        );
+        try {
+          await this.onReconnect();
+          this.needsReInitialization = false;
+        } catch (error) {
+          // Close the uninitialized connection so the next attempt starts fresh
+          this.close();
+          Logger.error(
+            "Copilot#sendData: LSP re-initialization after reconnection failed",
+            error,
+          );
+          throw error;
+        }
+      }
     }
 
     // After reconnection, delegate should be initialized
@@ -211,11 +239,8 @@ export class LazyWebsocketTransport extends Transport {
       );
     }
 
-    // Clamp timeout to maxTimeoutMs
-    timeout = Math.min(
-      timeout ?? this.options.maxTimeoutMs,
-      this.options.maxTimeoutMs,
-    );
+    // Use maxTimeoutMs as default when no timeout is provided
+    timeout = timeout ?? this.options.maxTimeoutMs;
     return this.delegate.sendData(data, timeout);
   }
 }