@marimo-team/frontend 0.21.2-dev1 → 0.21.2-dev101

package/src/core/datasets/__tests__/data-source.test.ts

@@ -1,6 +1,7 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it } from "vitest";
-import type { DataTable } from "@/core/kernel/messages";
+import { variableName } from "@/__tests__/branded";
+import type { DatabaseSchema, DataTable } from "@/core/kernel/messages";
 import type { VariableName } from "@/core/variables/types";
 import {
   type DataSourceConnection,
@@ -182,7 +183,7 @@ describe("filtering data sources", () => {
   });
 
   it("keeps matching variables and internal engines", () => {
-    const filtered = filterDataSources(["conn1" as unknown as VariableName]);
+    const filtered = filterDataSources([variableName("conn1")]);
     expect(filtered.connectionsMap.size).toBe(defaultConnSize + 1);
     expect(filtered.connectionsMap.has("conn1" as ConnectionName)).toBe(true);
     for (const engine of INTERNAL_SQL_ENGINES) {
@@ -191,16 +192,14 @@ describe("filtering data sources", () => {
   });
 
   it("filters out non-matching variables", () => {
-    const filtered = filterDataSources([
-      "non_existent" as unknown as VariableName,
-    ]);
+    const filtered = filterDataSources([variableName("non_existent")]);
     expect(filtered.connectionsMap.size).toBe(defaultConnSize);
   });
 
   it("handles mix of matching and non-matching variables", () => {
     const filtered = filterDataSources([
-      "conn1" as unknown as VariableName,
-      "non_existent" as unknown as VariableName,
+      variableName("conn1"),
+      variableName("non_existent"),
     ]);
     expect(filtered.connectionsMap.size).toBe(defaultConnSize + 1);
     expect(filtered.connectionsMap.has("conn1" as ConnectionName)).toBe(true);
@@ -210,6 +209,93 @@ describe("filtering data sources", () => {
   });
 });
 
+describe("add schema list", () => {
+  const connections: DataSourceConnection[] = [
+    {
+      name: "conn1" as ConnectionName,
+      source: "sqlite",
+      display_name: "SQLite DB",
+      dialect: "sqlite",
+      databases: [
+        {
+          name: "db1",
+          schemas: [],
+          dialect: "sqlite",
+        },
+      ],
+    },
+  ];
+
+  // Helper function to add schema list
+  const addSchemaList = (
+    schemas: DatabaseSchema[],
+    engine: string,
+    database: string,
+  ) => {
+    return reducer(baseState, {
+      type: "addSchemaList",
+      payload: {
+        schemas,
+        sqlSchemaContext: { engine, database },
+      },
+    });
+  };
+
+  let baseState: DataSourceState;
+
+  beforeEach(() => {
+    baseState = addConnection(connections, baseState);
+    expect(baseState.connectionsMap.size).toBe(defaultConnSize + 1);
+  });
+
+  it("adds schema list to a specific database", () => {
+    const schemaList: DatabaseSchema[] = [
+      { name: "public", tables: [] },
+      { name: "analytics", tables: [] },
+    ];
+    const newState = addSchemaList(schemaList, "conn1", "db1");
+
+    const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
+    const db1 = conn1?.databases.find((db) => db.name === "db1");
+    expect(db1?.schemas).toEqual(schemaList);
+  });
+
+  it("updates schema list for a database", () => {
+    const schemaList: DatabaseSchema[] = [
+      { name: "public", tables: [] },
+      { name: "analytics", tables: [] },
+    ];
+    const newState = addSchemaList(schemaList, "conn1", "db1");
+
+    const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
+    const db1 = conn1?.databases.find((db) => db.name === "db1");
+    expect(db1?.schemas).toEqual(schemaList);
+
+    // update with new schema list
+    const newSchemaList: DatabaseSchema[] = [
+      { name: "public", tables: [] },
+      { name: "sales", tables: [] },
+    ];
+    const updatedState = addSchemaList(newSchemaList, "conn1", "db1");
+
+    const newConn = updatedState.connectionsMap.get("conn1" as ConnectionName);
+    const newDb1 = newConn?.databases.find((db) => db.name === "db1");
+    expect(newDb1?.schemas).toEqual(newSchemaList);
+  });
+
+  it("does not add schema list if database does not exist", () => {
+    const schemaList: DatabaseSchema[] = [
+      { name: "public", tables: [] },
+      { name: "analytics", tables: [] },
+    ];
+    const newState = addSchemaList(schemaList, "conn1", "non_existent_db");
+
+    const conn1 = newState.connectionsMap.get("conn1" as ConnectionName);
+    const db1 = conn1?.databases.find((db) => db.name === "db1");
+    expect(db1?.schemas.length).toBe(0);
+  });
+});
+
 describe("add table list", () => {
   const connections: DataSourceConnection[] = [
     {

package/src/core/datasets/data-source-connections.ts

@@ -5,6 +5,7 @@ import { isSchemaless } from "@/components/datasources/utils";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import { Logger } from "@/utils/Logger";
 import type {
+  DatabaseSchema,
   DataSourceConnection as DataSourceConnectionType,
   DataTable,
 } from "../kernel/messages";
@@ -43,6 +44,11 @@ export interface DataSourceState {
   connectionsMap: ConnectionsMap;
 }
 
+export interface SQLSchemaContext {
+  engine: string;
+  database: string;
+}
+
 export interface SQLTableContext {
   engine: string;
   database: string;
@@ -134,6 +140,44 @@ const {
     };
   },
 
+  // Add schema list to a specific database in a connection
+  addSchemaList: (
+    state: DataSourceState,
+    opts: {
+      schemas: DatabaseSchema[];
+      sqlSchemaContext: SQLSchemaContext;
+    },
+  ): DataSourceState => {
+    const { schemas, sqlSchemaContext } = opts;
+    const { connectionsMap, latestEngineSelected } = state;
+    const connectionName = sqlSchemaContext.engine as ConnectionName;
+    const conn = connectionsMap.get(connectionName);
+
+    if (!conn) {
+      return state;
+    }
+
+    const newMap = new Map(connectionsMap);
+    const newConn: DataSourceConnection = {
+      ...conn,
+      databases: conn.databases.map((db) => {
+        if (db.name !== sqlSchemaContext.database) {
+          return db;
+        }
+        return {
+          ...db,
+          schemas: schemas,
+        };
+      }),
+    };
+    newMap.set(connectionName, newConn);
+
+    return {
+      latestEngineSelected: latestEngineSelected,
+      connectionsMap: newMap,
+    };
+  },
+
   // Add table list to a specific schema in a connection
   addTableList: (
     state: DataSourceState,

package/src/core/datasets/request-registry.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import type {
+  SQLSchemaListPreview,
   SQLTableListPreview,
   SQLTablePreview,
   ValidateSQLResult,
@@ -8,6 +9,7 @@ import { CachingRequestRegistry } from "../network/CachingRequestRegistry";
 import { DeferredRequestRegistry } from "../network/DeferredRequestRegistry";
 import { getRequestClient } from "../network/requests";
 import type {
+  ListSQLSchemasRequest,
   ListSQLTablesRequest,
   PreviewSQLTableRequest,
   ValidateSQLRequest,
@@ -28,6 +30,17 @@ export const PreviewSQLTable = new DeferredRequestRegistry<
   });
 });
 
+export const PreviewSQLSchemaList = new DeferredRequestRegistry<
+  Omit<ListSQLSchemasRequest, "requestId">,
+  SQLSchemaListPreview
+>("sql-schema-list-preview", async (requestId, req) => {
+  const client = getRequestClient();
+  await client.previewSQLSchemaList({
+    requestId: requestId,
+    ...req,
+  });
+});
+
 export const PreviewSQLTableList = new DeferredRequestRegistry<
   Omit<ListSQLTablesRequest, "requestId">,
   SQLTableListPreview

package/src/core/datasets/state.ts

@@ -83,7 +83,7 @@ const {
       if (!table.variable_name) {
         return true;
       }
-      return names.has(table.variable_name as VariableName);
+      return names.has(table.variable_name);
     });
     return { ...state, tables };
   },

package/src/core/edit-app.tsx

@@ -20,13 +20,11 @@ import { CellArray } from "../components/editor/renderers/cell-array";
 import { CellsRenderer } from "../components/editor/renderers/cells-renderer";
 import { useHotkey } from "../hooks/useHotkey";
 import {
-  cellIdsAtom,
   hasCellsAtom,
   notebookIsRunningAtom,
   numColumnsAtom,
   useCellActions,
 } from "./cells/cells";
-import { CellEffects } from "./cells/effects";
 import type { AppConfig, UserConfig } from "./config/config-schema";
 import { RuntimeState } from "./kernel/RuntimeState";
 import { getSessionId } from "./kernel/session";
@@ -35,7 +33,6 @@ import { viewStateAtom } from "./mode";
 import { useRequestClient } from "./network/requests";
 import { useFilename } from "./saving/filename";
 import { lastSavedNotebookAtom } from "./saving/state";
-import { useJotaiEffect } from "./state/jotai";
 import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
 
 interface AppProps {
@@ -58,8 +55,6 @@ export const EditApp: React.FC<AppProps> = ({
   appConfig,
   hideControls = false,
 }) => {
-  useJotaiEffect(cellIdsAtom, CellEffects.onCellIdsChange);
-
   const { setCells, mergeAllColumns, collapseAllCells, expandAllCells } =
     useCellActions();
   const viewState = useAtomValue(viewStateAtom);

package/src/core/errors/__tests__/errors.test.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import type { MarimoError } from "@/core/kernel/messages";
 import { getAutoFixes, getImportCode } from "../errors";
 
@@ -25,7 +26,7 @@ describe("getAutoFixes", () => {
     const error: MarimoError = {
       type: "multiple-defs",
       name: "foo",
-      cells: ["foo"],
+      cells: [cellId("foo")],
     };
 
     const fixes = getAutoFixes(error, opts);

package/src/core/export/__tests__/hooks.test.ts

@@ -6,6 +6,7 @@ import type { ReactNode } from "react";
 import * as React from "react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { MockModules, Mocks } from "@/__mocks__/common";
+import { cellId } from "@/__tests__/branded";
 import type { CellId } from "@/core/cells/ids";
 import { CellOutputId } from "@/core/cells/ids";
 import type { CellRuntimeState } from "@/core/cells/types";
@@ -63,7 +64,7 @@ describe("useEnrichCellOutputs", () => {
   ): Record<CellId, CellRuntimeState> => {
     return Object.fromEntries(
       Object.entries(cells).map(([cellId, cell]) => [
-        cellId as CellId,
+        cellId,
         {
           output: cell.output || null,
           status: cell.status || "idle",
@@ -106,7 +107,7 @@ describe("useEnrichCellOutputs", () => {
   });
 
   it("should capture screenshots for cells with text/html output", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockElement = document.createElement("div");
     const mockDataUrl = "data:image/png;base64,mockImageData";
 
@@ -116,7 +117,7 @@ describe("useEnrichCellOutputs", () => {
 
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -133,7 +134,7 @@ describe("useEnrichCellOutputs", () => {
     const output = await takeScreenshots({ progress });
 
     expect(document.getElementById).toHaveBeenCalledWith(
-      CellOutputId.create(cellId),
+      CellOutputId.create(cid),
     );
     expect(toPng).toHaveBeenCalledWith(
       mockElement,
@@ -143,12 +144,12 @@ describe("useEnrichCellOutputs", () => {
       }),
     );
     expect(output).toEqual({
-      [cellId]: ["image/png", mockDataUrl],
+      [cid]: ["image/png", mockDataUrl],
     });
   });
 
   it("should skip cells where output has not changed", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockElement = document.createElement("div");
     const mockDataUrl = "data:image/png;base64,mockImageData";
     const htmlData = "<div>Chart</div>";
@@ -158,7 +159,7 @@ describe("useEnrichCellOutputs", () => {
 
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -176,7 +177,7 @@ describe("useEnrichCellOutputs", () => {
     // First call - should capture
     let takeScreenshots = result.current;
     let output = await takeScreenshots({ progress });
-    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
+    expect(output).toEqual({ [cid]: ["image/png", mockDataUrl] });
     expect(toPng).toHaveBeenCalledTimes(1);
 
     // Rerender to get updated atom state
@@ -190,7 +191,7 @@ describe("useEnrichCellOutputs", () => {
   });
 
   it("should handle screenshot errors gracefully", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockElement = document.createElement("div");
     const error = new Error("Screenshot failed");
 
@@ -199,7 +200,7 @@ describe("useEnrichCellOutputs", () => {
 
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -217,13 +218,13 @@ describe("useEnrichCellOutputs", () => {
 
     expect(output).toEqual({}); // Failed screenshot should be filtered out
     expect(Logger.error).toHaveBeenCalledWith(
-      `Error screenshotting cell ${cellId}:`,
+      `Error screenshotting cell ${cid}:`,
       error,
     );
   });
 
   it("should retry failed screenshots on next call", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockElement = document.createElement("div");
     const error = new Error("Screenshot failed");
     const mockDataUrl = "data:image/png;base64,retrySuccess";
@@ -236,7 +237,7 @@ describe("useEnrichCellOutputs", () => {
 
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -262,18 +263,18 @@ describe("useEnrichCellOutputs", () => {
     // Second call - should retry since the first one failed
     takeScreenshots = result.current;
     output = await takeScreenshots({ progress });
-    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl] });
+    expect(output).toEqual({ [cid]: ["image/png", mockDataUrl] });
     expect(toPng).toHaveBeenCalledTimes(2);
   });
 
   it("should handle missing DOM elements", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
 
     vi.spyOn(document, "getElementById").mockReturnValue(null);
 
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -291,14 +292,14 @@ describe("useEnrichCellOutputs", () => {
 
     expect(output).toEqual({});
     expect(Logger.error).toHaveBeenCalledWith(
-      `Output element not found for cell ${cellId}`,
+      `Output element not found for cell ${cid}`,
     );
     expect(toPng).not.toHaveBeenCalled();
   });
 
   it("should process multiple cells in parallel", async () => {
-    const cell1 = "cell-1" as CellId;
-    const cell2 = "cell-2" as CellId;
+    const cell1 = "cell-1";
+    const cell2 = "cell-2";
     const mockElement1 = document.createElement("div");
     const mockElement2 = document.createElement("div");
     const mockDataUrl1 = "data:image/png;base64,image1";
@@ -347,8 +348,8 @@ describe("useEnrichCellOutputs", () => {
 
   it("should filter out null results from failed screenshots", async () => {
     // Setup: one successful, one failed screenshot
-    const cell1 = "cell-1" as CellId;
-    const cell2 = "cell-2" as CellId;
+    const cell1 = "cell-1";
+    const cell2 = "cell-2";
     const mockElement1 = document.createElement("div");
     const mockDataUrl = "data:image/png;base64,image1";
 
@@ -394,7 +395,7 @@ describe("useEnrichCellOutputs", () => {
   });
 
   it("should only capture screenshots for cells with changed output", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockElement = document.createElement("div");
     const mockDataUrl1 = "data:image/png;base64,image1";
     const mockDataUrl2 = "data:image/png;base64,image2";
@@ -407,7 +408,7 @@ describe("useEnrichCellOutputs", () => {
     // First call - cell should be captured
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -425,7 +426,7 @@ describe("useEnrichCellOutputs", () => {
     // First screenshot
     let takeScreenshots = result.current;
     let output = await takeScreenshots({ progress });
-    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl1] });
+    expect(output).toEqual({ [cid]: ["image/png", mockDataUrl1] });
 
     // Second call - same output, should not be captured
     rerender();
@@ -436,7 +437,7 @@ describe("useEnrichCellOutputs", () => {
     // Third call - output changed, should be captured
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -450,7 +451,7 @@ describe("useEnrichCellOutputs", () => {
     rerender();
     takeScreenshots = result.current;
     output = await takeScreenshots({ progress });
-    expect(output).toEqual({ [cellId]: ["image/png", mockDataUrl2] });
+    expect(output).toEqual({ [cid]: ["image/png", mockDataUrl2] });
     expect(toPng).toHaveBeenCalledTimes(2);
   });
 
@@ -523,7 +524,7 @@ describe("useEnrichCellOutputs", () => {
 
   it("should return correctly formatted result with CellId and tuple", async () => {
     // Expected format: Record<CellId, ["image/png", string]>
-    const cellId = "test-cell" as CellId;
+    const cid = cellId("test-cell");
     const mockElement = document.createElement("div");
     const mockDataUrl = "data:image/png;base64,testData";
 
@@ -532,7 +533,7 @@ describe("useEnrichCellOutputs", () => {
 
     setCellsRuntime(
       createMockCellRuntimes({
-        [cellId]: {
+        [cid]: {
           output: {
             channel: "output",
             mimetype: "text/html",
@@ -549,8 +550,8 @@ describe("useEnrichCellOutputs", () => {
     const output = await takeScreenshots({ progress });
 
     // Verify the exact return type structure
-    expect(output).toHaveProperty(cellId);
-    const cellOutput = output[cellId];
+    expect(output).toHaveProperty(cid);
+    const cellOutput = output[cid];
     expect(cellOutput).toBeDefined();
     expect(Array.isArray(cellOutput)).toBe(true);
     if (cellOutput) {
@@ -566,9 +567,9 @@ describe("updateCellOutputsWithScreenshots", () => {
   });
 
   it("should call updateCellOutputs when there are screenshots", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockScreenshots = {
-      [cellId]: ["image/png", "data:image/png;base64,test"] as [
+      [cid]: ["image/png", "data:image/png;base64,test"] as [
         "image/png",
         string,
       ],
@@ -603,8 +604,8 @@ describe("updateCellOutputsWithScreenshots", () => {
   });
 
   it("should handle multiple cell screenshots", async () => {
-    const cell1 = "cell-1" as CellId;
-    const cell2 = "cell-2" as CellId;
+    const cell1 = "cell-1";
+    const cell2 = "cell-2";
     const mockScreenshots = {
       [cell1]: ["image/png", "data:image/png;base64,image1"] as [
         "image/png",
@@ -654,9 +655,9 @@ describe("updateCellOutputsWithScreenshots", () => {
   });
 
   it("should catch errors from updateCellOutputs and show toast", async () => {
-    const cellId = "cell-1" as CellId;
+    const cid = cellId("cell-1");
     const mockScreenshots = {
-      [cellId]: ["image/png", "data:image/png;base64,test"] as [
+      [cid]: ["image/png", "data:image/png;base64,test"] as [
         "image/png",
         string,
       ],

package/src/core/islands/bridge.ts

@@ -162,12 +162,13 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   saveAppConfig = throwNotImplemented;
   saveCellConfig = throwNotImplemented;
   sendRestart = throwNotImplemented;
-  syncCellIds = throwNotImplemented;
+  sendDocumentTransaction = throwNotImplemented;
   readCode = throwNotImplemented;
   readSnippets = throwNotImplemented;
   previewDatasetColumn = throwNotImplemented;
   previewSQLTable = throwNotImplemented;
   previewSQLTableList = throwNotImplemented;
+  previewSQLSchemaList = throwNotImplemented;
   previewDataSourceConnection = throwNotImplemented;
   validateSQL = throwNotImplemented;
   openFile = throwNotImplemented;

package/src/core/islands/main.ts

@@ -28,7 +28,6 @@ import {
   notebookAtom,
   notebookReducer,
 } from "../cells/cells";
-import type { UIElementId } from "../cells/ids";
 import { defineCustomElement } from "../dom/defineCustomElement";
 import { MarimoValueInputEvent } from "../dom/events";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
@@ -41,7 +40,6 @@ import {
 import { queryParamHandlers } from "../kernel/queryParamHandlers";
 import { RuntimeState } from "../kernel/RuntimeState";
 import { initialModeAtom } from "../mode";
-import type { RequestId } from "../network/DeferredRequestRegistry";
 import { requestClientAtom } from "../network/requests";
 import { store } from "../state/jotai";
 import { IslandsPyodideBridge } from "./bridge";
@@ -115,14 +113,13 @@ export async function initialize() {
       case "installing-package-alert":
       case "completion-result":
       case "reload":
-      case "update-cell-codes":
-      case "update-cell-ids":
       case "focus-cell":
       case "variables":
       case "variable-values":
       case "data-column-preview":
       case "sql-table-preview":
       case "sql-table-list-preview":
+      case "sql-schema-list-preview":
       case "datasets":
       case "data-source-connections":
       case "validate-sql-result":
@@ -153,7 +150,7 @@ export async function initialize() {
         return;
       case "send-ui-element-message":
         UI_ELEMENT_REGISTRY.broadcastMessage(
-          msg.data.ui_element as UIElementId,
+          msg.data.ui_element,
           msg.data.message,
           safeExtractSetUIElementMessageBuffers(msg.data),
         );
@@ -163,10 +160,7 @@ export async function initialize() {
         handleRemoveUIElements(msg.data);
         return;
       case "function-call-result":
-        FUNCTIONS_REGISTRY.resolve(
-          msg.data.function_call_id as RequestId,
-          msg.data,
-        );
+        FUNCTIONS_REGISTRY.resolve(msg.data.function_call_id, msg.data);
         return;
       case "cell-op":
         handleCellNotificationeration(msg.data, actions.handleCellMessage);
@@ -201,6 +195,8 @@ export async function initialize() {
         return;
       case "kernel-startup-error":
         return;
+      case "notebook-document-transaction":
+        return;
       case "model-lifecycle":
         handleWidgetMessage(MODEL_MANAGER, msg.data);
         return;

package/src/core/kernel/__tests__/handlers.test.ts

@@ -1,5 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
 import { buildCellData, buildLayoutState } from "../handlers";
 import type { NotificationMessageData } from "../messages";
 import { queryParamHandlers } from "../queryParamHandlers";
@@ -65,7 +66,7 @@ describe("queryParamHandlers", () => {
 describe("buildCellData", () => {
   it("should build cell data from kernel-ready data", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: ["cell1", "cell2"],
+      cell_ids: [cellId("cell1"), cellId("cell2")],
       codes: ["x = 1", "y = 2"],
       names: ["__", "__"],
       configs: [
@@ -127,7 +128,7 @@ describe("buildCellData", () => {
 
   it("should mark cells as edited when code differs from last executed code", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: ["cell1", "cell2"],
+      cell_ids: [cellId("cell1"), cellId("cell2")],
       codes: ["x = 1", "y = 3"],
       names: ["__", "__"],
       configs: [
@@ -201,7 +202,7 @@ describe("buildCellData", () => {
 describe("buildLayoutState", () => {
   it("should build default layout state when no layout is provided", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: ["cell1"],
+      cell_ids: [cellId("cell1")],
       codes: ["x = 1"],
       names: ["__"],
       configs: [{ disabled: false, hide_code: false }],
@@ -243,7 +244,7 @@ describe("buildLayoutState", () => {
 
   it("should build layout state with vertical layout", () => {
     const kernelReadyData: NotificationMessageData<"kernel-ready"> = {
-      cell_ids: ["cell1", "cell2"],
+      cell_ids: [cellId("cell1"), cellId("cell2")],
       codes: ["x = 1", "y = 2"],
       names: ["__", "__"],
       configs: [