@marimo-team/frontend 0.22.1-dev105 → 0.22.1-dev15

package/src/components/editor/ai/add-cell-with-ai.tsx

@@ -122,7 +122,7 @@ export const AddCellWithAI: React.FC<{
     transport: new StreamingChunkTransport(
       {
         api: runtimeManager.getAiURL("completion").toString(),
-        headers: runtimeManager.headers(),
+        headers: () => runtimeManager.headers(),
         prepareSendMessagesRequest: async (options) => {
           const completionBody = await buildCompletionRequestBody(
             options.messages,
@@ -136,6 +136,7 @@ export const AddCellWithAI: React.FC<{
           };
 
           return {
+            api: runtimeManager.getAiURL("completion").toString(),
             body: body,
           };
         },

package/src/components/editor/chrome/panels/context-aware-panel/context-aware-panel.tsx

@@ -101,7 +101,7 @@ export const ContextAwarePanel: React.FC = () => {
   const renderBody = () => {
     return (
       <div className="pb-7 mb-4 h-full overflow-auto">
-        <div className="p-3 border-b flex justify-between items-center">
+        <div className="px-3 py-2 border-b flex justify-between items-center">
           {renderModeToggle()}
           <Slot name={SlotNames.CONTEXT_AWARE_PANEL_HEADER} />
           <Button

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/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/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/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

@@ -466,6 +466,7 @@ interface DataTableProps<T> extends Data<T>, DataTableFunctions {
   hoverTemplate?: string | null;
   cellHoverTexts?: Record<string, Record<string, string | null>> | null;
   toggleDisplayHeader?: () => void;
+  isChartBuilderOpen?: boolean;
   host: HTMLElement;
   cellId?: CellId | null;
 }
@@ -749,6 +750,7 @@ export const LoadingDataTableComponent = memo(
           > | null
         }
         toggleDisplayHeader={toggleDisplayHeader}
+        isChartBuilderOpen={displayHeader}
         getRow={getRow}
         cellId={cellId}
         maxHeight={props.maxHeight}
@@ -761,6 +763,7 @@ export const LoadingDataTableComponent = memo(
         {props.showChartBuilder ? (
           <TablePanel
             displayHeader={displayHeader}
+            onCloseChartBuilder={() => setDisplayHeader(false)}
             data={data?.rows || []}
             columns={props.totalColumns}
             totalRows={props.totalRows}
@@ -821,6 +824,7 @@ const DataTableComponent = ({
   hoverTemplate,
   cellHoverTexts,
   toggleDisplayHeader,
+  isChartBuilderOpen,
   calculate_top_k_rows,
   preview_column,
   getRow,
@@ -1088,6 +1092,7 @@ const DataTableComponent = ({
             getRowIds={get_row_ids}
             toggleDisplayHeader={toggleDisplayHeader}
             showChartBuilder={showChartBuilder}
+            isChartBuilderOpen={isChartBuilderOpen}
             showPageSizeSelector={showPageSizeSelector}
             // Hidden in VSCode (for now) because we don't have a panel to show
             // the table explorer.

package/src/plugins/impl/vega/__tests__/utils.test.ts

@@ -33,4 +33,72 @@ describe("getContainerWidth", () => {
   it("should return undefined when width is explicitly undefined", () => {
     expect(getContainerWidth({ width: undefined })).toBeUndefined();
   });
+
+  it("should find width in nested facet spec", () => {
+    expect(
+      getContainerWidth({
+        $schema: "https://vega.github.io/schema/vega-lite/v6.json",
+        facet: { column: { field: "Origin", type: "nominal" } },
+        spec: {
+          mark: "point",
+          encoding: {},
+          width: "container",
+        },
+      }),
+    ).toBe("container");
+  });
+
+  it("should find width in nested repeat spec", () => {
+    expect(
+      getContainerWidth({
+        $schema: "https://vega.github.io/schema/vega-lite/v6.json",
+        repeat: { row: ["A", "B"] },
+        spec: {
+          mark: "point",
+          encoding: {},
+          width: "container",
+        },
+      }),
+    ).toBe("container");
+  });
+
+  it("should return undefined for nested spec without width", () => {
+    expect(
+      getContainerWidth({
+        facet: { column: { field: "Origin" } },
+        spec: { mark: "point", encoding: {} },
+      }),
+    ).toBeUndefined();
+  });
+
+  it("should return undefined for hconcat (width on sub-specs)", () => {
+    expect(
+      getContainerWidth({
+        hconcat: [{ width: "container" }, { width: "container" }],
+      }),
+    ).toBeUndefined();
+  });
+
+  it("should return undefined for vconcat (width on sub-specs)", () => {
+    expect(
+      getContainerWidth({
+        vconcat: [{ width: "container" }, { width: "container" }],
+      }),
+    ).toBeUndefined();
+  });
+
+  it("should return undefined for compiled Vega spec (width as signal)", () => {
+    expect(
+      getContainerWidth({
+        $schema: "https://vega.github.io/schema/vega/v6.json",
+        autosize: { contains: "padding", type: "fit-x" },
+        signals: [
+          {
+            name: "width",
+            init: "isFinite(containerSize()[0]) ? containerSize()[0] : 300",
+          },
+        ],
+      }),
+    ).toBeUndefined();
+  });
 });

package/src/plugins/impl/vega/utils.ts

@@ -6,13 +6,22 @@ import type { DataType, FieldTypes, VegaDataType } from "./vega-loader";
 export type ContainerWidth = number | "container";
 
 /**
- * Get the container width from a VegaLite spec.
- * @param spec - The VegaLite spec.
- * @returns The container width.
+ * Get the container width from a Vega-Lite spec.
+ *
+ * For unit specs, `width` is at the top level. For facet/repeat specs,
+ * `width` is nested inside `spec`. This does not handle hconcat/vconcat and Vega spec
+ * where the width may be a signal. These cases are covered by
+ * the CSS fallback `.vega-embed:has(> .chart-wrapper.fit-x)`.
  */
 export function getContainerWidth(spec: unknown): ContainerWidth | undefined {
-  if (typeof spec === "object" && spec !== null && "width" in spec) {
-    return spec.width as ContainerWidth | undefined;
+  if (typeof spec === "object" && spec !== null) {
+    if ("width" in spec) {
+      return spec.width as ContainerWidth | undefined;
+    }
+    // Faceted/repeated spec
+    if ("spec" in spec) {
+      return getContainerWidth(spec.spec);
+    }
   }
   return undefined;
 }

package/src/plugins/impl/vega/vega.css

@@ -4,7 +4,8 @@
   max-width: 100%;
 }
 
-.vega-embed[data-container-width="container"] {
+.vega-embed[data-container-width="container"],
+.vega-embed:has(> .chart-wrapper.fit-x) {
   width: 100%;
 }
 

package/src/utils/time.ts

@@ -1,8 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-export type Milliseconds = number & { __type__: "milliseconds" };
+import type { TypedNumber } from "./typed";
 
-export type Seconds = number & { __type__: "seconds" };
+export type Milliseconds = TypedNumber<"milliseconds">;
+
+export type Seconds = TypedNumber<"seconds">;
 
 export class Time {
   private readonly ms: Milliseconds;