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

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"),

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

@@ -2,9 +2,9 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import { describe, expect, it, vi } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId, variableName } from "@/__tests__/branded";
 import type { DatasetTablesMap } from "@/core/datasets/data-source-connections";
-import type { Variable, VariableName, Variables } from "@/core/variables/types";
+import type { Variable, Variables } from "@/core/variables/types";
 import { type VariableContextItem, VariableContextProvider } from "../variable";
 
 // Mock the variable completions module
@@ -17,7 +17,7 @@ vi.mock("@/core/codemirror/completion/variable-completions", () => ({
       prefix: string,
     ) => {
       return Object.entries(variables).map(([name, variable]) => ({
-        name: `${prefix}${name}`,
+        name: variableName(`${prefix}${name}`),
         displayname: name,
         detail: variable.dataType || "unknown",
         boost,
@@ -34,9 +34,9 @@ const createMockVariable = (
   name: string,
   options: Partial<Variable> = {},
 ): Variable => ({
-  name: name as VariableName,
-  declaredBy: ["cell1" as CellId],
-  usedBy: ["cell2" as CellId],
+  name: variableName(name),
+  declaredBy: [cellId("cell1")],
+  usedBy: [cellId("cell2")],
   value: `value_of_${name}`,
   dataType: "str",
   ...options,
@@ -55,10 +55,13 @@ describe("VariableContextProvider", () => {
 
     it("should return variable items for single variable", () => {
       const variables: Variables = {
-        ["user_name" as VariableName]: createMockVariable("user_name", {
-          value: '"John Doe"',
-          dataType: "str",
-        }),
+        [variableName("user_name")]: createMockVariable(
+          variableName("user_name"),
+          {
+            value: '"John Doe"',
+            dataType: "str",
+          },
+        ),
       };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
@@ -69,29 +72,32 @@ describe("VariableContextProvider", () => {
 
     it("should return variable items for multiple variables with different types", () => {
       const variables: Variables = {
-        ["user_id" as VariableName]: createMockVariable("user_id", {
+        [variableName("user_id")]: createMockVariable(variableName("user_id"), {
           value: "123",
           dataType: "int",
-          declaredBy: ["cell1" as CellId],
-          usedBy: ["cell2" as CellId, "cell3" as CellId],
-        }),
-        ["is_active" as VariableName]: createMockVariable("is_active", {
-          value: "True",
-          dataType: "bool",
-          declaredBy: ["cell2" as CellId],
-          usedBy: [],
+          declaredBy: [cellId("cell1")],
+          usedBy: [cellId("cell2"), cellId("cell3")],
         }),
-        ["scores" as VariableName]: createMockVariable("scores", {
+        [variableName("is_active")]: createMockVariable(
+          variableName("is_active"),
+          {
+            value: "True",
+            dataType: "bool",
+            declaredBy: [cellId("cell2")],
+            usedBy: [],
+          },
+        ),
+        [variableName("scores")]: createMockVariable(variableName("scores"), {
           value: "[1, 2, 3, 4, 5]",
           dataType: "list",
-          declaredBy: ["cell3" as CellId],
-          usedBy: ["cell4" as CellId],
+          declaredBy: [cellId("cell3")],
+          usedBy: [cellId("cell4")],
         }),
-        ["config" as VariableName]: createMockVariable("config", {
+        [variableName("config")]: createMockVariable(variableName("config"), {
           value: '{"debug": true, "timeout": 30}',
           dataType: "dict",
-          declaredBy: ["cell1" as CellId],
-          usedBy: ["cell2" as CellId, "cell4" as CellId],
+          declaredBy: [cellId("cell1")],
+          usedBy: [cellId("cell2"), cellId("cell4")],
         }),
       };
       const tablesMap: DatasetTablesMap = new Map();
@@ -103,18 +109,27 @@ describe("VariableContextProvider", () => {
 
     it("should handle variables with null/undefined values", () => {
       const variables: Variables = {
-        ["null_var" as VariableName]: createMockVariable("null_var", {
-          value: null,
-          dataType: "NoneType",
-        }),
-        ["undefined_var" as VariableName]: createMockVariable("undefined_var", {
-          value: undefined,
-          dataType: null,
-        }),
-        ["empty_string" as VariableName]: createMockVariable("empty_string", {
-          value: '""',
-          dataType: "str",
-        }),
+        [variableName("null_var")]: createMockVariable(
+          variableName("null_var"),
+          {
+            value: null,
+            dataType: "NoneType",
+          },
+        ),
+        [variableName("undefined_var")]: createMockVariable(
+          variableName("undefined_var"),
+          {
+            value: undefined,
+            dataType: null,
+          },
+        ),
+        [variableName("empty_string")]: createMockVariable(
+          variableName("empty_string"),
+          {
+            value: '""',
+            dataType: "str",
+          },
+        ),
       };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
@@ -125,23 +140,23 @@ describe("VariableContextProvider", () => {
 
     it("should handle complex data types", () => {
       const variables: Variables = {
-        ["df" as VariableName]: createMockVariable("df", {
+        [variableName("df")]: createMockVariable(variableName("df"), {
           value: "<DataFrame shape: (100, 5)>",
           dataType: "pandas.DataFrame",
-          declaredBy: ["cell1" as CellId],
-          usedBy: ["cell2" as CellId, "cell3" as CellId],
+          declaredBy: [cellId("cell1")],
+          usedBy: [cellId("cell2"), cellId("cell3")],
         }),
-        ["model" as VariableName]: createMockVariable("model", {
+        [variableName("model")]: createMockVariable(variableName("model"), {
           value: "<sklearn.linear_model.LinearRegression>",
           dataType: "sklearn.linear_model._base.LinearRegression",
-          declaredBy: ["cell4" as CellId],
-          usedBy: ["cell5" as CellId],
+          declaredBy: [cellId("cell4")],
+          usedBy: [cellId("cell5")],
         }),
-        ["array" as VariableName]: createMockVariable("array", {
+        [variableName("array")]: createMockVariable(variableName("array"), {
           value: "array([1, 2, 3, 4, 5])",
           dataType: "numpy.ndarray",
-          declaredBy: ["cell2" as CellId],
-          usedBy: ["cell3" as CellId],
+          declaredBy: [cellId("cell2")],
+          usedBy: [cellId("cell3")],
         }),
       };
       const tablesMap: DatasetTablesMap = new Map();
@@ -153,21 +168,27 @@ describe("VariableContextProvider", () => {
 
     it("should handle variables with special characters in names", () => {
       const variables: Variables = {
-        ["_private_var" as VariableName]: createMockVariable("_private_var", {
-          value: "42",
-          dataType: "int",
-        }),
-        ["var_with_numbers123" as VariableName]: createMockVariable(
-          "var_with_numbers123",
+        [variableName("_private_var")]: createMockVariable(
+          variableName("_private_var"),
+          {
+            value: "42",
+            dataType: "int",
+          },
+        ),
+        [variableName("var_with_numbers123")]: createMockVariable(
+          variableName("var_with_numbers123"),
           {
             value: '"test"',
             dataType: "str",
           },
         ),
-        ["CONSTANT_VAR" as VariableName]: createMockVariable("CONSTANT_VAR", {
-          value: "3.14159",
-          dataType: "float",
-        }),
+        [variableName("CONSTANT_VAR")]: createMockVariable(
+          variableName("CONSTANT_VAR"),
+          {
+            value: "3.14159",
+            dataType: "float",
+          },
+        ),
       };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
@@ -191,7 +212,7 @@ describe("VariableContextProvider", () => {
         data: { variable },
       };
 
-      const variables: Variables = { ["username" as VariableName]: variable };
+      const variables: Variables = { [variableName("username")]: variable };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
 
@@ -213,7 +234,7 @@ describe("VariableContextProvider", () => {
       };
 
       const variables: Variables = {
-        ["mystery_var" as VariableName]: variable,
+        [variableName("mystery_var")]: variable,
       };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
@@ -237,7 +258,7 @@ describe("VariableContextProvider", () => {
       };
 
       const variables: Variables = {
-        ["complex_data" as VariableName]: variable,
+        [variableName("complex_data")]: variable,
       };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
@@ -260,7 +281,7 @@ describe("VariableContextProvider", () => {
         data: { variable },
       };
 
-      const variables: Variables = { ["sales_df" as VariableName]: variable };
+      const variables: Variables = { [variableName("sales_df")]: variable };
       const tablesMap: DatasetTablesMap = new Map();
       const provider = new VariableContextProvider(variables, tablesMap);
 
@@ -284,12 +305,12 @@ describe("VariableContextProvider", () => {
   describe("integration with tables", () => {
     it("should work with both variables and tables", () => {
       const variables: Variables = {
-        ["df" as VariableName]: createMockVariable("df", {
+        [variableName("df")]: createMockVariable(variableName("df"), {
           dataType: "pandas.DataFrame",
           value: "<DataFrame shape: (50, 3)>",
         }),
-        ["connection_string" as VariableName]: createMockVariable(
-          "connection_string",
+        [variableName("connection_string")]: createMockVariable(
+          variableName("connection_string"),
           {
             dataType: "str",
             value: '"postgresql://localhost:5432/mydb"',

package/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts

@@ -5,6 +5,7 @@ import { EditorView } from "@codemirror/view";
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { MockNotebook } from "@/__mocks__/notebook";
+import { cellId } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { updateEditorCodeFromPython } from "@/core/codemirror/language/utils";
@@ -32,7 +33,7 @@ function createMockEditorView(code: string): EditorView {
       doc: code,
       extensions: [
         adaptiveLanguageConfiguration({
-          cellId: "cell1" as CellId,
+          cellId: cellId("cell1"),
           completionConfig: {
             copilot: false,
             activate_on_typing: true,
@@ -44,7 +45,7 @@ function createMockEditorView(code: string): EditorView {
           lspConfig: {},
         }),
         cellConfigExtension({
-          cellId: "cell1" as CellId,
+          cellId: cellId("cell1"),
           completionConfig: {
             copilot: false,
             activate_on_typing: true,
@@ -93,9 +94,9 @@ describe("EditNotebookTool", () => {
     };
     tool = new EditNotebookTool();
 
-    cellId1 = "cell-1" as CellId;
-    cellId2 = "cell-2" as CellId;
-    cellId3 = "cell-3" as CellId;
+    cellId1 = cellId("cell-1");
+    cellId2 = cellId("cell-2");
+    cellId3 = cellId("cell-3");
 
     // Reset mocks
     vi.clearAllMocks();
@@ -270,7 +271,7 @@ describe("EditNotebookTool", () => {
           {
             edit: {
               type: "update_cell",
-              position: { cellId: "nonexistent" as CellId },
+              position: { cellId: cellId("nonexistent") },
               code: "x = 2",
             },
           },
@@ -283,7 +284,7 @@ describe("EditNotebookTool", () => {
           {
             edit: {
               type: "update_cell",
-              position: { cellId: "nonexistent" as CellId },
+              position: { cellId: cellId("nonexistent") },
               code: "x = 2",
             },
           },
@@ -496,7 +497,7 @@ describe("EditNotebookTool", () => {
               type: "add_cell",
               position: {
                 type: "relative",
-                cellId: "nonexistent" as CellId,
+                cellId: cellId("nonexistent"),
                 before: true,
               },
               code: "y = 2",
@@ -660,7 +661,7 @@ describe("EditNotebookTool", () => {
           {
             edit: {
               type: "delete_cell",
-              position: { cellId: "nonexistent" as CellId },
+              position: { cellId: cellId("nonexistent") },
             },
           },
           toolContext as never,

package/src/core/ai/tools/__tests__/run-cells-tool.test.ts

@@ -3,6 +3,7 @@
 import { getDefaultStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { MockNotebook } from "@/__mocks__/notebook";
+import { cellId } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { RunStaleCellsTool } from "../run-cells-tool";
@@ -46,9 +47,9 @@ describe("RunStaleCellsTool", () => {
 
     tool = new RunStaleCellsTool({ postExecutionDelay: 0 });
 
-    cellId1 = "cell-1" as CellId;
-    cellId2 = "cell-2" as CellId;
-    cellId3 = "cell-3" as CellId;
+    cellId1 = cellId("cell-1");
+    cellId2 = cellId("cell-2");
+    cellId3 = cellId("cell-3");
 
     // Reset mocks
     vi.clearAllMocks();
@@ -490,9 +491,8 @@ describe("RunStaleCellsTool", () => {
     });
 
     it("should omit output for cells that exceed total output budget", async () => {
-      const cellIds = Array.from(
-        { length: 25 },
-        (_, i) => `budget-cell-${i}` as CellId,
+      const cellIds = Array.from({ length: 25 }, (_, i) =>
+        cellId(`budget-cell-${i}`),
       );
       const cellData: Record<string, { code: string; edited: boolean }> = {};
       for (const id of cellIds) {