@marimo-team/frontend 0.21.2-dev39 → 0.21.2-dev40

package/src/components/editor/ai/__tests__/completion-utils.test.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it, type Mock, vi } from "vitest";
+import { variableName } from "@/__tests__/branded";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
 import { dataSourceConnectionsAtom } from "@/core/datasets/data-source-connections";
 import { DUCKDB_ENGINE } from "@/core/datasets/engines";
@@ -7,7 +8,6 @@ import { datasetsAtom } from "@/core/datasets/state";
 import type { DatasetsState } from "@/core/datasets/types";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
-import type { Variable, VariableName } from "@/core/variables/types";
 import { codeToCells, getAICompletionBody } from "../completion-utils";
 
 // Mock getCodes function
@@ -205,16 +205,16 @@ describe("getAICompletionBody", () => {
 
   it("should return the correct completion body with mentioned variables", () => {
     // Set up test data in the Jotai store
-    const testVariables: Record<VariableName, Variable> = {
-      ["var1" as VariableName]: {
-        name: "var1" as VariableName,
+    const testVariables = {
+      [variableName("var1")]: {
+        name: variableName("var1"),
         value: "string value",
         dataType: "string",
         declaredBy: [],
         usedBy: [],
       },
-      ["var2" as VariableName]: {
-        name: "var2" as VariableName,
+      [variableName("var2")]: {
+        name: variableName("var2"),
         value: "42",
         dataType: "number",
         declaredBy: [],
@@ -253,9 +253,9 @@ describe("getAICompletionBody", () => {
     ];
     store.set(datasetsAtom, { tables: testDatasets } as DatasetsState);
 
-    const testVariables: Record<VariableName, Variable> = {
-      ["var1" as VariableName]: {
-        name: "var1" as VariableName,
+    const testVariables = {
+      [variableName("var1")]: {
+        name: variableName("var1"),
         value: "string value",
         dataType: "string",
         declaredBy: [],
@@ -285,9 +285,9 @@ describe("getAICompletionBody", () => {
 
   it("should handle non-existent variables", () => {
     // Set up test data in the Jotai store
-    const testVariables: Record<VariableName, Variable> = {
-      ["existingVar" as VariableName]: {
-        name: "existingVar" as VariableName,
+    const testVariables = {
+      [variableName("existingVar")]: {
+        name: variableName("existingVar"),
         value: "string value",
         dataType: "string",
         declaredBy: [],
@@ -322,9 +322,9 @@ describe("getAICompletionBody", () => {
     ];
     store.set(datasetsAtom, { tables: testDatasets } as DatasetsState);
 
-    const testVariables: Record<VariableName, Variable> = {
-      ["conflict" as VariableName]: {
-        name: "conflict" as VariableName,
+    const testVariables = {
+      [variableName("conflict")]: {
+        name: variableName("conflict"),
         value: "string value",
         dataType: "string",
         declaredBy: [],

package/src/components/editor/navigation/__tests__/clipboard.test.ts

@@ -4,8 +4,8 @@
 import { act, renderHook } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { asMock, MockModules, Mocks, SetupMocks } from "@/__mocks__/common";
+import { cellId } from "@/__tests__/branded";
 import type { CellActions, NotebookState } from "@/core/cells/cells";
-import type { CellId } from "@/core/cells/ids";
 import { useCellClipboard } from "../clipboard";
 
 // Mock dependencies
@@ -179,7 +179,7 @@ describe("useCellClipboard", () => {
     });
 
     it("should filter out non-existent cells", async () => {
-      const nonExistentCellId = "non-existent" as CellId;
+      const nonExistentCellId = cellId("non-existent");
 
       const { result } = renderHook(() => useCellClipboard());
 

package/src/components/editor/navigation/__tests__/selection.test.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
 import { beforeEach, describe, expect, it } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId } from "@/__tests__/branded";
 import { MultiColumn } from "@/utils/id-tree";
 import type { CellSelectionState } from "../selection";
 import { exportedForTesting } from "../selection";
@@ -8,11 +9,11 @@ import { exportedForTesting } from "../selection";
 const { initialState, reducer, createActions } = exportedForTesting;
 
 const CellIds = {
-  a: "a" as CellId,
-  b: "b" as CellId,
-  c: "c" as CellId,
-  d: "d" as CellId,
-  e: "e" as CellId,
+  a: cellId("a"),
+  b: cellId("b"),
+  c: cellId("c"),
+  d: cellId("d"),
+  e: cellId("e"),
 };
 
 describe("cell selection reducer", () => {

package/src/components/editor/navigation/__tests__/state.test.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import { type CellId, HTMLCellId } from "@/core/cells/ids";
 
 const mockScrollCellIntoView = vi.fn();
@@ -14,12 +15,12 @@ vi.mock("../focus-utils", () => ({
 type TemporarilyShownCodeState = Set<CellId>;
 
 describe("temporarilyShownCodeActions", () => {
-  const cellId = "cell-1" as CellId;
+  const cid = cellId("cell-1");
   let cellElement: HTMLElement;
 
   beforeEach(() => {
     cellElement = document.createElement("div");
-    cellElement.id = HTMLCellId.create(cellId);
+    cellElement.id = HTMLCellId.create(cid);
     document.body.append(cellElement);
     cellElement.focus();
 
@@ -37,17 +38,17 @@ describe("temporarilyShownCodeActions", () => {
   });
 
   it("should scroll cell into view when removing cell causes layout shift", () => {
-    const state = new Set<CellId>([cellId]);
-    removeCell(state, cellId);
+    const state = new Set<CellId>([cid]);
+    removeCell(state, cid);
 
-    expect(mockScrollCellIntoView).toHaveBeenCalledWith(cellId);
+    expect(mockScrollCellIntoView).toHaveBeenCalledWith(cid);
   });
 
   it("should not scroll when focused cell is not found", () => {
     vi.spyOn(HTMLCellId, "findElementThroughShadowDOMs").mockReturnValue(null);
 
-    const state = new Set<CellId>([cellId]);
-    removeCell(state, cellId);
+    const state = new Set<CellId>([cid]);
+    removeCell(state, cid);
 
     expect(mockScrollCellIntoView).not.toHaveBeenCalled();
   });

package/src/components/editor/output/MarimoErrorOutput.tsx

@@ -316,7 +316,7 @@ export const MarimoErrorOutput = ({
               <ul className="list-disc">
                 {error.cells.map((cid, cidIdx) => (
                   <li className={liStyle} key={`cell-${cidIdx}`}>
-                    <CellLinkError cellId={cid as CellId} />
+                    <CellLinkError cellId={cid} />
                   </li>
                 ))}
               </ul>
@@ -491,7 +491,7 @@ export const MarimoErrorOutput = ({
                 ) : (
                   <div>
                     {processTextForUrls(error.msg, `exception-${idx}`)}
-                    <CellLinkError cellId={error.raising_cell as CellId} />
+                    <CellLinkError cellId={error.raising_cell} />
                   </div>
                 )}
               </li>
@@ -518,7 +518,7 @@ export const MarimoErrorOutput = ({
               ) : (
                 <div>
                   {error.msg}
-                  <CellLinkError cellId={error.blamed_cell as CellId} />
+                  <CellLinkError cellId={error.blamed_cell} />
                 </div>
               )}
             </li>
@@ -554,13 +554,13 @@ export const MarimoErrorOutput = ({
               {error.msg}
               {error.blamed_cell == null ? (
                 <span>
-                  (<CellLinkError cellId={error.raising_cell as CellId} />)
+                  (<CellLinkError cellId={error.raising_cell} />)
                 </span>
               ) : (
                 <span>
-                  (<CellLinkError cellId={error.raising_cell as CellId} />
+                  (<CellLinkError cellId={error.raising_cell} />
                   &nbsp;blames&nbsp;
-                  <CellLinkError cellId={error.blamed_cell as CellId} />)
+                  <CellLinkError cellId={error.blamed_cell} />)
                 </span>
               )}
             </div>
@@ -578,7 +578,7 @@ export const MarimoErrorOutput = ({
           {ancestorStoppedErrors.map((error, idx) => (
             <div key={`ancestor-stopped-${idx}`}>
               {error.msg}
-              <CellLinkError cellId={error.raising_cell as CellId} />
+              <CellLinkError cellId={error.raising_cell} />
             </div>
           ))}
           {cellId && (

package/src/components/editor/output/__tests__/traceback.test.tsx

@@ -3,8 +3,8 @@
 import { render } from "@testing-library/react";
 import { beforeAll, beforeEach, describe, expect, test, vi } from "vitest";
 import { Tracebacks } from "@/__mocks__/tracebacks";
+import { cellId } from "@/__tests__/branded";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import type { CellId } from "@/core/cells/ids";
 import { initialModeAtom } from "@/core/mode";
 import { store } from "@/core/state/jotai";
 import { renderHTML } from "@/plugins/core/RenderHTML";
@@ -14,7 +14,7 @@ import {
   replaceTracebackPrefix,
 } from "../MarimoTracebackOutput";
 
-const cellId = "1" as CellId;
+const cid = cellId("1");
 
 describe("traceback component", () => {
   beforeEach(() => {
@@ -25,7 +25,7 @@ describe("traceback component", () => {
   test("extracts cell-link", () => {
     const traceback = (
       <TooltipProvider>
-        <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cellId} />
+        <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cid} />
       </TooltipProvider>
     );
     const { unmount, getAllByRole } = render(traceback);
@@ -45,7 +45,7 @@ describe("traceback component", () => {
   test("renames File to Cell for relevant lines", () => {
     const traceback = (
       <TooltipProvider>
-        <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cellId} />
+        <MarimoTracebackOutput traceback={Tracebacks.raw} cellId={cid} />
       </TooltipProvider>
     );
     const { unmount, container } = render(traceback);

package/src/components/editor/output/console/__tests__/ConsoleOutput.test.tsx

@@ -3,8 +3,8 @@
 import { act, fireEvent, render, screen } from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { SetupMocks } from "@/__mocks__/common";
+import { cellId } from "@/__tests__/branded";
 import { TooltipProvider } from "@/components/ui/tooltip";
-import type { CellId } from "@/core/cells/ids";
 import type { WithResponse } from "@/core/cells/types";
 import type { OutputMessage } from "@/core/kernel/messages";
 import { CONSOLE_CLEAR_DEBOUNCE_MS, ConsoleOutput } from "../ConsoleOutput";
@@ -24,7 +24,7 @@ describe("ConsoleOutput integration", () => {
   });
 
   const defaultProps = {
-    cellId: "cell-1" as CellId,
+    cellId: cellId("cell-1"),
     cellName: "test_cell",
     consoleOutputs: [] as WithResponse<OutputMessage>[],
     stale: false,
@@ -55,7 +55,7 @@ describe("ConsoleOutput integration", () => {
 
 describe("ConsoleOutput pdb history", () => {
   const defaultProps = {
-    cellId: "cell-1" as CellId,
+    cellId: cellId("cell-1"),
     cellName: "test_cell",
     consoleOutputs: [] as WithResponse<OutputMessage>[],
     stale: false,
@@ -215,7 +215,7 @@ describe("ConsoleOutput debounced clearing", () => {
   });
 
   const defaultProps = {
-    cellId: "cell-1" as CellId,
+    cellId: cellId("cell-1"),
     cellName: "test_cell",
     consoleOutputs: [] as WithResponse<OutputMessage>[],
     stale: false,

package/src/components/editor/renderers/vertical-layout/useFocusFirstEditor.ts

@@ -91,7 +91,7 @@ function focusCellByName(cellName: string) {
 
       // Look for an editor to focus
       const { cellHandles } = getNotebook();
-      const cellId = cellElement.dataset.cellId as CellId | undefined;
+      const cellId = extractCellIdFromDomElement(cellElement);
 
       if (!cellId) {
         Logger.error(`Missing cellId for cell with name ${cellName}`);
@@ -111,3 +111,10 @@ function focusCellByName(cellName: string) {
     focusFirstEditor();
   }
 }
+
+function extractCellIdFromDomElement(
+  cellElement: HTMLElement,
+): CellId | undefined {
+  const cellIdStr = cellElement.dataset.cellId ?? undefined;
+  return cellIdStr as CellId | undefined;
+}

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";
@@ -474,7 +473,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

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/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,
       });