@marimo-team/frontend 0.22.1-dev2 → 0.22.1-dev21

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,6 +165,7 @@ 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",
@@ -183,6 +192,7 @@ export class LazyWebsocketTransport extends Transport {
   override close(): void {
     this.delegate?.close();
     this.delegate = undefined;
+    this.needsReInitialization = true;
   }
 
   override async sendData(
@@ -202,6 +212,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 +240,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);
   }
 }

package/src/core/islands/__tests__/bridge.test.ts

@@ -1,5 +1,15 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
+import type { components } from "@marimo-team/marimo-api";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  cellId,
+  requestId,
+  uiElementId,
+  widgetModelId,
+} from "@/__tests__/branded";
+
+type Base64String = components["schemas"]["Base64String"];
 
 // Mock browser APIs before any imports
 vi.stubGlobal(
@@ -89,7 +99,7 @@ describe("IslandsPyodideBridge", () => {
   describe("sendComponentValues", () => {
     it("should include type field and token in control request", async () => {
       const request = {
-        objectIds: ["Hbol-0"],
+        objectIds: [uiElementId("Hbol-0")],
         values: [58],
       };
 
@@ -108,7 +118,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("should preserve all request properties", async () => {
       const request = {
-        objectIds: ["slider-1", "slider-2"],
+        objectIds: [uiElementId("slider-1"), uiElementId("slider-2")],
         values: [10, 20],
       };
 
@@ -128,7 +138,7 @@ describe("IslandsPyodideBridge", () => {
   describe("sendFunctionRequest", () => {
     it("should include type field in control request", async () => {
       const request = {
-        functionCallId: "call-123",
+        functionCallId: requestId("call-123"),
         namespace: "test_namespace",
         functionName: "my_function",
         args: { x: 1, y: 2 },
@@ -152,7 +162,7 @@ describe("IslandsPyodideBridge", () => {
   describe("sendRun", () => {
     it("should include type field in control request", async () => {
       const request = {
-        cellIds: ["cell-1", "cell-2"],
+        cellIds: [cellId("cell-1"), cellId("cell-2")],
         codes: ["print('hello')", "print('world')"],
       };
 
@@ -170,7 +180,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("should call loadPackages before putControlRequest", async () => {
       const request = {
-        cellIds: ["cell-1"],
+        cellIds: [cellId("cell-1")],
         codes: ["import pandas"],
       };
 
@@ -190,13 +200,13 @@ describe("IslandsPyodideBridge", () => {
   describe("sendModelValue", () => {
     it("should include type field in control request", async () => {
       const request = {
-        modelId: "widget-1",
+        modelId: widgetModelId("widget-1"),
         message: {
           method: "update" as const,
           state: { value: 42 },
           bufferPaths: [],
         },
-        buffers: [],
+        buffers: [] as Base64String[],
       };
 
       await bridge.sendModelValue(request);
@@ -222,16 +232,16 @@ describe("IslandsPyodideBridge", () => {
       // Test all methods to ensure they include the type field
       await bridge.sendComponentValues({ objectIds: [], values: [] });
       await bridge.sendFunctionRequest({
-        functionCallId: "",
+        functionCallId: requestId(""),
         namespace: "",
         functionName: "",
         args: {},
       });
       await bridge.sendRun({ cellIds: [], codes: [] });
       await bridge.sendModelValue({
-        modelId: "",
+        modelId: widgetModelId(""),
         message: { method: "update", state: {}, bufferPaths: [] },
-        buffers: [],
+        buffers: [] as Base64String[],
       });
 
       // All calls should have the type field

package/src/core/network/__tests__/requests-lazy.test.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId, requestId, uiElementId } from "@/__tests__/branded";
 import type { RuntimeManager } from "../../runtime/runtime";
 import { createLazyRequests } from "../requests-lazy";
 import type { EditRequests, RunRequests } from "../types";
@@ -59,7 +60,7 @@ describe("createLazyRequests", () => {
       mockGetRuntimeManager,
     );
 
-    await lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] });
+    await lazyRequests.sendRun({ cellIds: [cellId("cell1")], codes: ["code"] });
 
     expect(mockInit).toHaveBeenCalledTimes(1);
   });
@@ -70,10 +71,13 @@ describe("createLazyRequests", () => {
       mockGetRuntimeManager,
     );
 
-    await lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] });
-    await lazyRequests.sendInstantiate({ objectIds: ["obj1"], values: [] });
+    await lazyRequests.sendRun({ cellIds: [cellId("cell1")], codes: ["code"] });
+    await lazyRequests.sendInstantiate({
+      objectIds: [uiElementId("obj1")],
+      values: [],
+    });
     await lazyRequests.sendFunctionRequest({
-      functionCallId: "func1",
+      functionCallId: requestId("func1"),
       functionName: "testFunc",
       args: {},
       namespace: "test",
@@ -89,7 +93,7 @@ describe("createLazyRequests", () => {
       mockGetRuntimeManager,
     );
 
-    await lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] });
+    await lazyRequests.sendRun({ cellIds: [cellId("cell1")], codes: ["code"] });
 
     expect(waitForConnectionOpen).toHaveBeenCalled();
   });
@@ -100,7 +104,7 @@ describe("createLazyRequests", () => {
       mockGetRuntimeManager,
     );
 
-    const args = { cellIds: ["cell1"], codes: ["code"] };
+    const args = { cellIds: [cellId("cell1")], codes: ["code"] };
     await lazyRequests.sendRun(args);
 
     expect(mockDelegate.sendRun).toHaveBeenCalledWith(args);
@@ -113,7 +117,7 @@ describe("createLazyRequests", () => {
     );
 
     const result = await lazyRequests.sendFunctionRequest({
-      functionCallId: "func1",
+      functionCallId: requestId("func1"),
       functionName: "testFunc",
       args: {},
       namespace: "test",
@@ -143,7 +147,7 @@ describe("createLazyRequests", () => {
     );
 
     await expect(
-      lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] }),
+      lazyRequests.sendRun({ cellIds: [cellId("cell1")], codes: ["code"] }),
     ).rejects.toThrow("Init failed");
   });
 
@@ -157,7 +161,7 @@ describe("createLazyRequests", () => {
     );
 
     await expect(
-      lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] }),
+      lazyRequests.sendRun({ cellIds: [cellId("cell1")], codes: ["code"] }),
     ).rejects.toThrow("Request failed");
   });
 
@@ -169,7 +173,10 @@ describe("createLazyRequests", () => {
       );
 
       // First request with first runtime manager
-      await lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] });
+      await lazyRequests.sendRun({
+        cellIds: [cellId("cell1")],
+        codes: ["code"],
+      });
       expect(mockInit).toHaveBeenCalledTimes(1);
 
       // Create a new runtime manager
@@ -187,7 +194,10 @@ describe("createLazyRequests", () => {
       );
 
       // Second request with second runtime manager
-      await lazyRequests2.sendRun({ cellIds: ["cell2"], codes: ["code2"] });
+      await lazyRequests2.sendRun({
+        cellIds: [cellId("cell2")],
+        codes: ["code2"],
+      });
 
       // Both inits should have been called
       expect(mockInit).toHaveBeenCalledTimes(1);
@@ -201,9 +211,15 @@ describe("createLazyRequests", () => {
       );
 
       // Multiple requests
-      await lazyRequests.sendRun({ cellIds: ["cell1"], codes: ["code"] });
-      await lazyRequests.sendDeleteCell({ cellId: "cell2" });
-      await lazyRequests.sendInstantiate({ objectIds: ["obj1"], values: [] });
+      await lazyRequests.sendRun({
+        cellIds: [cellId("cell1")],
+        codes: ["code"],
+      });
+      await lazyRequests.sendDeleteCell({ cellId: cellId("cell2") });
+      await lazyRequests.sendInstantiate({
+        objectIds: [uiElementId("obj1")],
+        values: [],
+      });
 
       // Init should only be called once
       expect(mockInit).toHaveBeenCalledTimes(1);

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -186,7 +186,7 @@ export function useMarimoKernelConnection(opts: {
         return;
 
       case "completion-result":
-        AUTOCOMPLETER.resolve(msg.data.completion_id as RequestId, msg.data);
+        AUTOCOMPLETER.resolve(msg.data.completion_id, msg.data);
         return;
       case "function-call-result":
         FUNCTIONS_REGISTRY.resolve(msg.data.function_call_id, msg.data);
@@ -207,20 +207,14 @@ export function useMarimoKernelConnection(opts: {
       case "variables":
         setVariables(
           msg.data.variables.map((v) => ({
-            name: v.name as VariableName,
+            name: v.name,
             declaredBy: v.declared_by,
             usedBy: v.used_by,
           })),
         );
-        filterDatasetsFromVariables(
-          msg.data.variables.map((v) => v.name as VariableName),
-        );
-        filterDataSourcesFromVariables(
-          msg.data.variables.map((v) => v.name as VariableName),
-        );
-        filterStorageFromVariables(
-          msg.data.variables.map((v) => v.name as VariableName),
-        );
+        filterDatasetsFromVariables(msg.data.variables.map((v) => v.name));
+        filterDataSourcesFromVariables(msg.data.variables.map((v) => v.name));
+        filterStorageFromVariables(msg.data.variables.map((v) => v.name));
         return;
       case "variable-values":
         setMetadata(

package/src/core/websocket/useWebSocket.tsx

@@ -30,6 +30,8 @@ function createConnectionTransport(
   // Create a connection transport using the ReconnectingWebSocket from partysocket
   // This handles reconnecting when the connection is lost.
   const urlProvider = options.url; // We don't call the URL provider now since it may change (i.e. if the runtime redirects)
+  // Cast needed: ReconnectingWebSocket types readyState as `number`
+  // but IConnectionTransport expects `0 | 1 | 2 | 3`
   return new ReconnectingWebSocket(urlProvider, undefined, {
     // We don't want Infinity retries
     maxRetries: 10,
@@ -38,7 +40,7 @@ function createConnectionTransport(
     // long timeout -- the server can become slow when many notebooks
     // are open.
     connectionTimeout: 10_000,
-  });
+  }) as unknown as IConnectionTransport;
 }
 
 /**

package/src/css/app/Cell.css

@@ -98,9 +98,30 @@
 
     /* Special case for particular components */
 
-    .output-area:has(> .output > marimo-ui-element > marimo-table) {
+    .output-area:has(
+        > .output:only-child > marimo-ui-element:only-child > marimo-table
+      ) {
+      padding: 0 0 5px;
       max-height: none;
       overflow: hidden;
+
+      /* Flush table: remove border and configure edge padding via CSS variable */
+      --marimo-table-edge-padding: 0.75rem;
+
+      marimo-table::part(table-tabs) {
+        margin-top: 0.25rem;
+        border: none;
+        border-radius: 0;
+      }
+
+      marimo-table::part(table-wrapper) {
+        border: none;
+        border-radius: 0;
+      }
+
+      marimo-table::part(table-footer) {
+        padding-inline: 0.25rem;
+      }
     }
 
     & > :first-child {

package/src/css/table.css

@@ -1,5 +1,22 @@
 @reference "../css/globals.css";
 
+/* Edge padding for flush tables (--marimo-table-edge-padding inherits through shadow DOM) */
+[part="table-wrapper"] th:first-child {
+  padding-left: var(--marimo-table-edge-padding, 0.5rem);
+}
+
+[part="table-wrapper"] th:last-child {
+  padding-right: var(--marimo-table-edge-padding, 0.5rem);
+}
+
+[part="table-wrapper"] td:first-child {
+  padding-left: var(--marimo-table-edge-padding, 0.375rem);
+}
+
+[part="table-wrapper"] td:last-child {
+  padding-right: var(--marimo-table-edge-padding, 0.375rem);
+}
+
 .markdown table,
 table.dataframe {
   display: block;

package/src/plugins/impl/DataTablePlugin.tsx

@@ -763,6 +763,7 @@ export const LoadingDataTableComponent = memo(
         {props.showChartBuilder ? (
           <TablePanel
             displayHeader={displayHeader}
+            onCloseChartBuilder={() => setDisplayHeader(false)}
             data={data?.rows || []}
             columns={props.totalColumns}
             totalRows={props.totalRows}

package/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx

@@ -4,7 +4,6 @@
 import type { AnyWidget } from "@anywidget/types";
 import { useEffect, useRef } from "react";
 import { z } from "zod";
-import { RANDOM_ID_ATTR } from "@/core/dom/ui-element-constants";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import type { HTMLElementNotDerivedFromRef } from "@/hooks/useEventListener";
 import { createPlugin } from "@/plugins/core/builder";
@@ -145,18 +144,10 @@ const AnyWidgetSlot = (props: IPluginProps<ModelIdRef, Data>) => {
     return <ErrorBanner error={error} />;
   }
 
-  // Find the closest parent element with an attribute of `random-id`
-  const randomId = props.host
-    .closest(`[${RANDOM_ID_ATTR}]`)
-    ?.getAttribute(RANDOM_ID_ATTR);
-  const key = randomId ?? jsUrl;
-
   return (
     <LoadedSlot
-      // Use the a key to force a re-render when the randomId (or jsUrl) changes
-      // Plugins may be stateful and we cannot make assumptions that we won't be
-      // so it is safer to just re-render.
-      key={key}
+      // Force remount when the widget module or model changes (cell re-run).
+      key={`${jsHash}:${modelId}`}
       widget={jsModule.default}
       modelId={modelId}
       host={host}