@marimo-team/islands 0.21.2-dev1 → 0.21.2-dev100

package/src/components/storage/storage-file-viewer.tsx

@@ -4,27 +4,32 @@ import { FileIcon, LoaderCircle, RefreshCwIcon } from "lucide-react";
 import type React from "react";
 import { useCallback } from "react";
 import { useLocale } from "react-aria";
+import { useAddCodeToNewCell } from "@/components/editor/cell/useAddCell";
 import { FilePreviewHeader } from "@/components/editor/file-tree/file-header";
 import { renderFileIcon } from "@/components/editor/file-tree/file-icons";
 import {
   FileContentRenderer,
   isMediaMime,
 } from "@/components/editor/file-tree/renderers";
+import { Tooltip } from "@/components/ui/tooltip";
 import { toast } from "@/components/ui/use-toast";
 import { DownloadStorage } from "@/core/storage/request-registry";
-import type { StorageEntry } from "@/core/storage/types";
+import type { StorageEntry, StorageNamespace } from "@/core/storage/types";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { downloadByURL } from "@/utils/download";
 import { formatBytes } from "@/utils/formatting";
 import { Logger } from "@/utils/Logger";
 import { CopyClipboardIcon } from "../icons/copy-icon";
 import { Button } from "../ui/button";
+import { STORAGE_SNIPPETS } from "./storage-snippets";
 
 const MAX_MEDIA_PREVIEW_SIZE = 100 * 1024 * 1024; // 100 MB
 
 interface Props {
   entry: StorageEntry;
   namespace: string;
+  protocol: string;
+  backendType: StorageNamespace["backendType"];
   onBack: () => void;
 }
 
@@ -41,9 +46,12 @@ type PreviewData =
 export const StorageFileViewer: React.FC<Props> = ({
   entry,
   namespace,
+  protocol,
+  backendType,
   onBack,
 }) => {
   const { locale } = useLocale();
+  const addCodeToNewCell = useAddCodeToNewCell();
   const name = displayName(entry.path);
   const mime = entry.mimeType || "text/plain";
   const isMedia = isMediaMime(mime);
@@ -111,12 +119,38 @@ export const StorageFileViewer: React.FC<Props> = ({
     }
   }, [namespace, entry.path, name]);
 
+  const snippetActions = STORAGE_SNIPPETS.map((snippet) => {
+    const code = snippet.getCode({
+      variableName: namespace,
+      protocol,
+      entry,
+      backendType,
+    });
+    if (code === null) {
+      return null;
+    }
+    const Icon = snippet.icon;
+    return (
+      <Tooltip key={snippet.id} content={snippet.label}>
+        <Button
+          variant="text"
+          size="xs"
+          onClick={() => addCodeToNewCell(code)}
+          aria-label={snippet.label}
+        >
+          <Icon className="h-3.5 w-3.5" />
+        </Button>
+      </Tooltip>
+    );
+  });
+
   const header = (
     <FilePreviewHeader
       filename={name}
       filenameIcon={renderFileIcon(name)}
       onBack={onBack}
       onDownload={handleDownload}
+      actions={snippetActions}
     />
   );
 

package/src/components/storage/storage-inspector.tsx

@@ -50,7 +50,6 @@ import type {
   StoragePathKey,
 } from "@/core/storage/types";
 import { storagePathKey } from "@/core/storage/types";
-import type { VariableName } from "@/core/variables/types";
 import { cn } from "@/utils/cn";
 import { copyToClipboard } from "@/utils/copy";
 import { downloadByURL } from "@/utils/download";
@@ -65,6 +64,8 @@ import { STORAGE_SNIPPETS } from "./storage-snippets";
 interface OpenFileInfo {
   entry: StorageEntry;
   namespace: string;
+  protocol: string;
+  backendType: StorageNamespace["backendType"];
 }
 
 // Pixels per depth level. Applied as paddingLeft on each full-width item
@@ -316,7 +317,7 @@ const StorageEntryRow: React.FC<{
           if (isDir) {
             setIsExpanded(!effectiveExpanded);
           } else {
-            onOpenFile({ entry, namespace });
+            onOpenFile({ entry, namespace, protocol, backendType });
           }
         }}
       >
@@ -362,7 +363,9 @@ const StorageEntryRow: React.FC<{
             >
               {!isDir && (
                 <DropdownMenuItem
-                  onSelect={() => onOpenFile({ entry, namespace })}
+                  onSelect={() =>
+                    onOpenFile({ entry, namespace, protocol, backendType })
+                  }
                 >
                   <ViewIcon className={MENU_ITEM_ICON_CLASS} />
                   View
@@ -474,7 +477,7 @@ const StorageNamespaceSection: React.FC<{
         <span>{namespace.displayName}</span>
         {namespace.name && (
           <span className="text-xs text-muted-foreground font-normal">
-            (<EngineVariable variableName={namespace.name as VariableName} />)
+            (<EngineVariable variableName={namespace.name} />)
           </span>
         )}
         <RefreshIconButton
@@ -586,6 +589,8 @@ export const StorageInspector: React.FC = () => {
         <StorageFileViewer
           entry={openFile.entry}
           namespace={openFile.namespace}
+          protocol={openFile.protocol}
+          backendType={openFile.backendType}
           onBack={() => setOpenFile(null)}
         />
       )}

package/src/components/storage/storage-snippets.ts

@@ -1,7 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import type { LucideIcon } from "lucide-react";
-import { BookOpenIcon, LinkIcon } from "lucide-react";
+import { BookPlusIcon, FileSymlink } from "lucide-react";
 import type { StorageEntry, StorageNamespace } from "@/core/storage/types";
 
 type BackendType = StorageNamespace["backendType"];
@@ -31,7 +31,7 @@ export const STORAGE_SNIPPETS: StorageSnippet[] = [
   {
     id: "read-file",
     label: "Insert read snippet",
-    icon: BookOpenIcon,
+    icon: BookPlusIcon,
     getCode: (ctx) => {
       if (ctx.entry.kind === "directory") {
         return null;
@@ -46,7 +46,7 @@ export const STORAGE_SNIPPETS: StorageSnippet[] = [
   {
     id: "download-file",
     label: "Insert download snippet",
-    icon: LinkIcon,
+    icon: FileSymlink,
     getCode: (ctx) => {
       if (ctx.entry.kind === "directory") {
         return null;

package/src/components/tracing/tracing.tsx

@@ -243,7 +243,9 @@ const TraceBlockBody: React.FC<{
         signalName: VEGA_HOVER_SIGNAL,
         handler: (_name: string, value: unknown) => {
           const signalValue = value as VegaHoverCellSignal;
-          const hoveredCell = signalValue.cell?.[0] as CellId | undefined;
+          const hoveredCell = (signalValue.cell?.[0] ?? undefined) as
+            | CellId
+            | undefined;
           setHoveredCellId(hoveredCell ?? null);
         },
       },

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/components/variables/variables-table.tsx

@@ -11,7 +11,6 @@ import {
   type SortingState,
   useReactTable,
 } from "@tanstack/react-table";
-import { sortBy } from "lodash-es";
 import { SquareEqualIcon, WorkflowIcon } from "lucide-react";
 import React, { memo, useMemo } from "react";
 import { useLocale } from "react-aria";
@@ -21,6 +20,7 @@ import type { CellId } from "@/core/cells/ids";
 import { isInternalCellName } from "@/core/cells/names";
 import { goToVariableDefinition } from "@/core/codemirror/go-to-definition/commands";
 import type { Variable, Variables } from "@/core/variables/types";
+import { sortBy } from "@/utils/arrays";
 import { cn } from "@/utils/cn";
 import { DataTableColumnHeader } from "../data-table/column-header";
 import { CellLinkList } from "../editor/links/cell-link-list";

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