@marimo-team/frontend 0.21.2-dev5 → 0.21.2-dev52

package/src/components/ui/range-slider.tsx

@@ -18,14 +18,116 @@ const RangeSlider = React.forwardRef<
   React.ElementRef<typeof SliderPrimitive.Root>,
   React.ComponentPropsWithoutRef<typeof SliderPrimitive.Root> & {
     valueMap: (sliderValue: number) => number;
+    steps?: number[];
   }
 >(({ className, valueMap, ...props }, ref) => {
   const [open, openActions] = useBoolean(false);
   const { locale } = useLocale();
 
+  const isDraggingRange = React.useRef(false);
+  const dragStartX = React.useRef(0);
+  const dragStartY = React.useRef(0);
+  const dragStartValue = React.useRef<number[]>([]);
+  const currentDragValue = React.useRef<number[]>([]);
+  const rootRef =
+    React.useRef<React.ElementRef<typeof SliderPrimitive.Root>>(null);
+  const trackRef = React.useRef<HTMLSpanElement>(null);
+  const dragTrackRect = React.useRef<DOMRect | null>(null);
+
+  const mergedRef = React.useCallback(
+    (node: React.ElementRef<typeof SliderPrimitive.Root>) => {
+      rootRef.current = node;
+      if (typeof ref === "function") {
+        ref(node);
+      } else if (ref) {
+        ref.current = node;
+      }
+    },
+    [ref],
+  );
+
+  const handleRangePointerDown = (e: React.PointerEvent<HTMLSpanElement>) => {
+    if (!props.value || props.value.length !== 2) {
+      return;
+    }
+    if (props.disabled) {
+      return;
+    }
+    e.preventDefault();
+    e.stopPropagation();
+
+    isDraggingRange.current = true;
+    dragStartX.current = e.clientX;
+    dragStartY.current = e.clientY;
+    dragStartValue.current = [...props.value];
+    currentDragValue.current = [...props.value];
+    dragTrackRect.current = trackRef.current?.getBoundingClientRect() ?? null;
+
+    (e.currentTarget as HTMLElement).setPointerCapture(e.pointerId);
+  };
+
+  const handleRangePointerMove = (e: React.PointerEvent<HTMLSpanElement>) => {
+    if (!isDraggingRange.current) {
+      return;
+    }
+    e.stopPropagation();
+
+    const trackRect = dragTrackRect.current;
+    if (!trackRect) {
+      return;
+    }
+
+    const isVertical = props.orientation === "vertical";
+    const min = props.min ?? 0;
+    const max = props.max ?? 100;
+    const totalRange = max - min;
+
+    let delta: number;
+    if (isVertical) {
+      const trackLength = trackRect.height;
+      delta = -((e.clientY - dragStartY.current) / trackLength) * totalRange;
+    } else {
+      const trackLength = trackRect.width;
+      delta = ((e.clientX - dragStartX.current) / trackLength) * totalRange;
+    }
+
+    const [origLeft, origRight] = dragStartValue.current;
+    const rangeWidth = origRight - origLeft;
+
+    const steps = props.steps;
+    const step: number =
+      steps && steps.length > 1
+        ? Math.min(...steps.slice(1).map((s, i) => s - steps[i]))
+        : (props.step ?? 1);
+    const snappedDelta = Math.round(delta / step) * step;
+
+    const clampedDelta = Math.max(
+      min - origLeft,
+      Math.min(max - origRight, snappedDelta),
+    );
+
+    const newLeft = origLeft + clampedDelta;
+    const newRight = newLeft + rangeWidth;
+
+    currentDragValue.current = [newLeft, newRight];
+    props.onValueChange?.([newLeft, newRight]);
+  };
+
+  const handleRangePointerUp = (e: React.PointerEvent<HTMLSpanElement>) => {
+    if (!isDraggingRange.current) {
+      return;
+    }
+    (e.currentTarget as HTMLElement).releasePointerCapture(e.pointerId);
+    isDraggingRange.current = false;
+
+    if (currentDragValue.current.length === 2) {
+      props.onValueCommit?.(currentDragValue.current);
+    }
+  };
+
   return (
     <SliderPrimitive.Root
-      ref={ref}
+      ref={mergedRef}
       className={cn(
         "relative flex touch-none select-none hover:cursor-pointer",
         "data-[orientation=horizontal]:w-full data-[orientation=horizontal]:items-center",
@@ -36,6 +138,7 @@ const RangeSlider = React.forwardRef<
       {...props}
     >
       <SliderPrimitive.Track
+        ref={trackRef}
         data-testid="track"
         className={cn(
           "relative grow overflow-hidden rounded-full bg-slate-200 dark:bg-accent/60",
@@ -50,7 +153,11 @@ const RangeSlider = React.forwardRef<
             "data-[orientation=horizontal]:h-full",
             "data-[orientation=vertical]:w-full",
             "data-disabled:opacity-50",
+            "hover:cursor-grab active:cursor-grabbing",
           )}
+          onPointerDown={handleRangePointerDown}
+          onPointerMove={handleRangePointerMove}
+          onPointerUp={handleRangePointerUp}
         />
       </SliderPrimitive.Track>
       <TooltipProvider>

package/src/core/ai/__tests__/staged-cells.test.ts

@@ -3,6 +3,7 @@
 import { renderHook } from "@testing-library/react";
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import { CellId } from "@/core/cells/ids";
 import { updateEditorCodeFromPython } from "../../codemirror/language/utils";
 import {
@@ -61,8 +62,8 @@ describe("staged-cells", () => {
 
   beforeEach(() => {
     store = getDefaultStore();
-    cellId1 = "cell-1" as CellId;
-    cellId2 = "cell-2" as CellId;
+    cellId1 = cellId("cell-1");
+    cellId2 = cellId("cell-2");
 
     // Reset mocks
     vi.clearAllMocks();
@@ -213,7 +214,7 @@ describe("staged-cells", () => {
       const testCode = "print('hello world')";
 
       // Mock CellId.create to return a predictable ID
-      const mockCellId = "mock-cell-id" as CellId;
+      const mockCellId = cellId("mock-cell-id");
       vi.mocked(CellId.create).mockReturnValue(mockCellId);
 
       const returnedCellId = result.current.createStagedCell(testCode);
@@ -229,7 +230,7 @@ describe("staged-cells", () => {
 
     it("should delete a staged cell", () => {
       const { result } = renderHook(() => useStagedCells(store));
-      const testCellId = "test-cell-id" as CellId;
+      const testCellId = cellId("test-cell-id");
 
       result.current.deleteStagedCell(testCellId);
 
@@ -330,7 +331,7 @@ describe("staged-cells", () => {
       const { result } = renderHook(() => useStagedCells(store));
 
       // Create a staged cell
-      const mockCellId = "mock-cell-id" as CellId;
+      const mockCellId = cellId("mock-cell-id");
       vi.mocked(CellId.create).mockReturnValue(mockCellId);
 
       const createdCellId = result.current.createStagedCell("test code");
@@ -341,7 +342,7 @@ describe("staged-cells", () => {
 
       let state = store.get(stagedAICellsAtom);
       expect(state.has(mockCellId)).toBe(true);
-      expect(state.get(mockCellId)).toEqual({ type: "add_cell" });
+      expect(state.get(mockCellId)).toEqual({ type: cellId("add_cell") });
 
       // Delete the staged cell
       result.current.deleteStagedCell(mockCellId);
@@ -423,7 +424,7 @@ describe("onStream", () => {
     result.current.onStream({ type: "text-start", id: "test-id" });
 
     // Mock CellId.create to return a predictable ID
-    const mockCellId = "mock-cell-id" as CellId;
+    const mockCellId = cellId("mock-cell-id");
     vi.mocked(CellId.create).mockReturnValue(mockCellId);
 
     result.current.onStream({
@@ -444,7 +445,7 @@ describe("onStream", () => {
     const { result } = renderHook(() => useStagedCells(store));
     result.current.onStream({ type: "text-start", id: "test-id" });
 
-    const mockCellId = "mock-cell-id" as CellId;
+    const mockCellId = cellId("mock-cell-id");
     vi.mocked(CellId.create).mockReturnValue(mockCellId);
 
     result.current.onStream({

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

@@ -2,6 +2,7 @@
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { Mocks } from "@/__mocks__/common";
+import { cellId } from "@/__tests__/branded";
 
 // Mock the external dependencies
 vi.mock("html-to-image", () => ({
@@ -13,7 +14,6 @@ vi.mock("@/utils/Logger", () => ({
 }));
 
 import type { NotebookState } from "@/core/cells/cells";
-import type { CellId } from "@/core/cells/ids";
 import type { OutputMessage } from "@/core/kernel/messages";
 import type { JotaiStore } from "@/core/state/jotai";
 import { CellOutputContextProvider, getCellContextData } from "../cell-output";
@@ -35,21 +35,21 @@ describe("CellOutputContextProvider", () => {
     // Create a basic mock notebook state
     mockNotebook = {
       cellIds: {
-        inOrderIds: ["cell1" as CellId, "cell2" as CellId, "cell3" as CellId],
+        inOrderIds: ["cell1", "cell2", "cell3"],
       },
       cellData: {
         cell1: {
-          id: "cell1" as CellId,
+          id: "cell1",
           name: "My Cell",
           code: "print('hello world')",
         },
         cell2: {
-          id: "cell2" as CellId,
+          id: "cell2",
           name: "",
           code: "import matplotlib.pyplot as plt\nplt.plot([1,2,3])",
         },
         cell3: {
-          id: "cell3" as CellId,
+          id: "cell3",
           name: "Empty Cell",
           code: "# no output",
         },
@@ -226,9 +226,9 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test" as CellId] },
+          cellIds: { inOrderIds: ["test"] },
           cellData: {
-            test: { id: "test" as CellId, name: "", code: "" },
+            test: { id: "test", name: "", code: "" },
           },
           cellRuntime: {
             test: {
@@ -261,9 +261,9 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test" as CellId] },
+          cellIds: { inOrderIds: ["test"] },
           cellData: {
-            test: { id: "test" as CellId, name: "", code: "" },
+            test: { id: "test", name: "", code: "" },
           },
           cellRuntime: {
             test: {
@@ -295,9 +295,9 @@ describe("Cell output utility functions", () => {
 
       for (const testCase of testCases) {
         mockStore = createMockStore({
-          cellIds: { inOrderIds: ["test" as CellId] },
+          cellIds: { inOrderIds: ["test"] },
           cellData: {
-            test: { id: "test" as CellId, name: "", code: "" },
+            test: { id: "test", name: "", code: "" },
           },
           cellRuntime: {
             test: {
@@ -369,9 +369,9 @@ describe("Cell output utility functions", () => {
       ).mockReturnValue(mockDiv);
 
       mockStore = createMockStore({
-        cellIds: { inOrderIds: ["test" as CellId] },
+        cellIds: { inOrderIds: ["test"] },
         cellData: {
-          test: { id: "test" as CellId, name: "", code: "" },
+          test: { id: "test", name: "", code: "" },
         },
         cellRuntime: {
           test: {
@@ -396,16 +396,16 @@ describe("Cell output utility functions", () => {
     beforeEach(() => {
       mockNotebook = {
         cellIds: {
-          inOrderIds: ["cell1" as CellId, "cell2" as CellId],
+          inOrderIds: ["cell1", "cell2"],
         },
         cellData: {
           cell1: {
-            id: "cell1" as CellId,
+            id: "cell1",
             name: "My Named Cell",
             code: "x = 42",
           },
           cell2: {
-            id: "cell2" as CellId,
+            id: "cell2",
             name: "",
             code: "y = x * 2",
           },
@@ -432,7 +432,7 @@ describe("Cell output utility functions", () => {
     });
 
     it("should extract basic cell data", () => {
-      const result = getCellContextData("cell1" as CellId, mockNotebook);
+      const result = getCellContextData(cellId("cell1"), mockNotebook);
 
       expect(result.cellId).toBe("cell1");
       expect(result.cellName).toBe("My Named Cell");
@@ -440,13 +440,13 @@ describe("Cell output utility functions", () => {
     });
 
     it("should generate cell name for unnamed cells", () => {
-      const result = getCellContextData("cell2" as CellId, mockNotebook);
+      const result = getCellContextData(cellId("cell2"), mockNotebook);
 
       expect(result.cellName).toBe("cell-1"); // 0-indexed, so cell2 is at index 1
     });
 
     it("should include cell output when present", () => {
-      const result = getCellContextData("cell1" as CellId, mockNotebook);
+      const result = getCellContextData(cellId("cell1"), mockNotebook);
 
       expect(result.cellOutput).toBeDefined();
       expect(result.cellOutput?.outputType).toBe("text");
@@ -454,19 +454,19 @@ describe("Cell output utility functions", () => {
     });
 
     it("should not include cell output when not present", () => {
-      const result = getCellContextData("cell2" as CellId, mockNotebook);
+      const result = getCellContextData(cellId("cell2"), mockNotebook);
 
       expect(result.cellOutput).toBeUndefined();
     });
 
     it("should not include console outputs by default", () => {
-      const result = getCellContextData("cell2" as CellId, mockNotebook);
+      const result = getCellContextData(cellId("cell2"), mockNotebook);
 
       expect(result.consoleOutputs).toBeUndefined();
     });
 
     it("should include console outputs when opted in", () => {
-      const result = getCellContextData("cell2" as CellId, mockNotebook, {
+      const result = getCellContextData(cellId("cell2"), mockNotebook, {
         includeConsoleOutput: true,
       });
 
@@ -478,7 +478,7 @@ describe("Cell output utility functions", () => {
     });
 
     it("should filter out empty console outputs", () => {
-      const cell2Runtime = mockNotebook.cellRuntime["cell2" as CellId];
+      const cell2Runtime = mockNotebook.cellRuntime[cellId("cell2")];
       if (cell2Runtime) {
         (cell2Runtime as { consoleOutputs: OutputMessage[] }).consoleOutputs = [
           {
@@ -492,7 +492,7 @@ describe("Cell output utility functions", () => {
         ];
       }
 
-      const result = getCellContextData("cell2" as CellId, mockNotebook, {
+      const result = getCellContextData(cellId("cell2"), mockNotebook, {
         includeConsoleOutput: true,
       });
 
@@ -501,7 +501,7 @@ describe("Cell output utility functions", () => {
     });
 
     it("should handle cells with both cell output and console outputs", () => {
-      const cell1Runtime = mockNotebook.cellRuntime["cell1" as CellId];
+      const cell1Runtime = mockNotebook.cellRuntime[cellId("cell1")];
       if (cell1Runtime) {
         (cell1Runtime as { consoleOutputs: OutputMessage[] }).consoleOutputs = [
           {
@@ -511,7 +511,7 @@ describe("Cell output utility functions", () => {
         ];
       }
 
-      const result = getCellContextData("cell1" as CellId, mockNotebook, {
+      const result = getCellContextData(cellId("cell1"), mockNotebook, {
         includeConsoleOutput: true,
       });
 
@@ -524,7 +524,7 @@ describe("Cell output utility functions", () => {
     });
 
     it("should handle empty console outputs array", () => {
-      const result = getCellContextData("cell1" as CellId, mockNotebook, {
+      const result = getCellContextData(cellId("cell1"), mockNotebook, {
         includeConsoleOutput: true,
       });
 
@@ -532,7 +532,7 @@ describe("Cell output utility functions", () => {
     });
 
     it("should handle media outputs in cell output", () => {
-      const cell1Runtime = mockNotebook.cellRuntime["cell1" as CellId];
+      const cell1Runtime = mockNotebook.cellRuntime[cellId("cell1")];
       if (cell1Runtime) {
         (cell1Runtime as { output: OutputMessage | null }).output = {
           mimetype: "image/png",
@@ -540,7 +540,7 @@ describe("Cell output utility functions", () => {
         } as OutputMessage;
       }
 
-      const result = getCellContextData("cell1" as CellId, mockNotebook);
+      const result = getCellContextData(cellId("cell1"), mockNotebook);
 
       expect(result.cellOutput).toBeDefined();
       expect(result.cellOutput?.outputType).toBe("media");
@@ -548,7 +548,7 @@ describe("Cell output utility functions", () => {
     });
 
     it("should handle media outputs in console outputs", () => {
-      const cell2Runtime = mockNotebook.cellRuntime["cell2" as CellId];
+      const cell2Runtime = mockNotebook.cellRuntime[cellId("cell2")];
       if (cell2Runtime) {
         (cell2Runtime as { consoleOutputs: OutputMessage[] }).consoleOutputs = [
           {
@@ -558,7 +558,7 @@ describe("Cell output utility functions", () => {
         ];
       }
 
-      const result = getCellContextData("cell2" as CellId, mockNotebook, {
+      const result = getCellContextData(cellId("cell2"), mockNotebook, {
         includeConsoleOutput: true,
       });
 

package/src/core/ai/context/providers/__tests__/datasource.test.ts

@@ -2,7 +2,7 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import { beforeEach, describe, expect, it } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId, variableName } from "@/__tests__/branded";
 import type {
   ConnectionsMap,
   DatasetTablesMap,
@@ -88,7 +88,7 @@ const createMockDataTable = (
   source_type: "local",
   num_rows: 100,
   num_columns: 3,
-  variable_name: name, // This makes it a dataframe
+  variable_name: variableName(name), // This makes it a dataframe
   columns: [
     {
       name: "id",
@@ -634,7 +634,7 @@ describe("DatasourceContextProvider", () => {
 
 describe("getDatasourceContext", () => {
   it("should return null if no cell ID is found", () => {
-    const context = getDatasourceContext("1" as CellId);
+    const context = getDatasourceContext(cellId("1"));
     expect(context).toBeNull();
   });
 });

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

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { describe, expect, it } from "vitest";
+import { variableName } from "@/__tests__/branded";
 import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
 import type { DataTable, DataTableColumn } from "@/core/kernel/messages";
 import { type TableContextItem, TableContextProvider } from "../tables";
@@ -86,7 +87,7 @@ describe("TableContextProvider", () => {
 
     it("should handle dataframe tables with variable names", () => {
       const dfTable = createMockTable("df_analysis", {
-        variable_name: "df_analysis",
+        variable_name: variableName("df_analysis"),
         source: "pandas",
         source_type: "local",
         columns: [
@@ -180,7 +181,7 @@ describe("TableContextProvider", () => {
       const table = createMockTable("remote_table", {
         source: "postgresql://localhost:5432/mydb",
         source_type: "connection",
-        engine: "postgresql",
+        engine: variableName("postgresql"),
         columns: [
           createMockColumn("uuid", "string", "uuid"),
           createMockColumn("created_at", "string", "timestamp with time zone"),