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

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) {

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

@@ -90,7 +90,7 @@ export class EditNotebookTool
     switch (edit.type) {
       case "update_cell": {
         const { position, code } = edit;
-        const cellId = position.cellId as CellId;
+        const cellId = position.cellId as CellId | undefined;
 
         if (!cellId) {
           throw new ToolExecutionError(
@@ -152,7 +152,7 @@ export class EditNotebookTool
 
         switch (position.type) {
           case "relative": {
-            const cellId = position.cellId as CellId;
+            const cellId = position.cellId as CellId | undefined;
             if (!cellId) {
               throw new ToolExecutionError(
                 "Cell ID is required for add_cell with relative position",
@@ -214,7 +214,7 @@ export class EditNotebookTool
       }
       case "delete_cell": {
         const { position } = edit;
-        const cellId = position.cellId as CellId;
+        const cellId = position.cellId as CellId | undefined;
 
         if (!cellId) {
           throw new ToolExecutionError(

package/src/core/cells/__tests__/add-missing-import.test.ts

@@ -4,6 +4,7 @@ import { createStore } from "jotai";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { MockNotebook } from "@/__mocks__/notebook";
 import { MockRequestClient } from "@/__mocks__/requests";
+import { cellId } from "@/__tests__/branded";
 import { store } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
 import type { Variables } from "@/core/variables/types";
@@ -13,7 +14,6 @@ import {
   maybeAddMissingImport,
 } from "../add-missing-import";
 import { notebookAtom } from "../cells";
-import type { CellId } from "../ids";
 
 // Mock the getRequestClient function
 const mockRequestClient = MockRequestClient.create();
@@ -21,8 +21,8 @@ vi.mock("@/core/network/requests", () => ({
   getRequestClient: () => mockRequestClient,
 }));
 
-const Cell1 = "1" as CellId;
-const Cell2 = "2" as CellId;
+const Cell1 = cellId("1");
+const Cell2 = cellId("2");
 
 describe("maybeAddMissingImport", () => {
   beforeEach(() => {