@marimo-team/frontend 0.21.2-dev6 → 0.21.2-dev61

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

@@ -163,6 +163,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   saveCellConfig = throwNotImplemented;
   sendRestart = throwNotImplemented;
   syncCellIds = throwNotImplemented;
+  sendDocumentTransaction = throwNotImplemented;
   readCode = throwNotImplemented;
   readSnippets = throwNotImplemented;
   previewDatasetColumn = 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";
@@ -153,7 +151,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 +161,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 +196,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: [

package/src/core/kernel/handlers.ts

@@ -3,7 +3,7 @@ import { deserializeLayout } from "@/components/editor/renderers/plugins";
 import type { LayoutType } from "@/components/editor/renderers/types";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
-import type { CellId, UIElementId } from "../cells/ids";
+import type { UIElementId } from "../cells/ids";
 import { type CellData, createCell } from "../cells/types";
 import { type AppConfig, AppConfigSchema } from "../config/config-schema";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
@@ -47,7 +47,7 @@ export function buildCellData(data: KernelReadyData): CellData[] {
     const edited = lastCodeRun ? lastCodeRun !== code : false;
 
     return createCell({
-      id: cellId as CellId,
+      id: cellId,
       code,
       edited,
       name: names[i],
@@ -160,7 +160,10 @@ export function handleKernelReady(
   // If resumed, restore UI element values from kernel and we're done
   if (resumed) {
     for (const [objectId, value] of Objects.entries(ui_values || {})) {
-      UI_ELEMENT_REGISTRY.set(objectId as UIElementId, value);
+      // @ts-expect-error - TODO: We need to fix the types for KernelReadyNotification["ui_values"]
+      // It is currently typed as Record<string, any> and should be Record<UIElementId, unknown>
+      const uiObjectId: UIElementId = objectId;
+      UI_ELEMENT_REGISTRY.set(uiObjectId, value);
     }
     return;
   }
@@ -196,7 +199,7 @@ export function handleRemoveUIElements(
   // This removes the element from the registry to (1) clean-up
   // memory and (2) make sure that the old value doesn't get re-used
   // if the same cell-id is later reused for another element.
-  const cellId = data.cell_id as CellId;
+  const cellId = data.cell_id;
   UI_ELEMENT_REGISTRY.removeElementsByCell(cellId);
   VirtualFileTracker.INSTANCE.removeForCellId(cellId);
 }

package/src/core/network/DeferredRequestRegistry.ts

@@ -1,10 +1,10 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* eslint-disable @typescript-eslint/no-redeclare */
+import type { components } from "@marimo-team/marimo-api";
 import { Deferred } from "@/utils/Deferred";
 import { generateUUID } from "@/utils/uuid";
-import type { TypedString } from "../../utils/typed";
 
-export type RequestId = TypedString<"RequestId">;
+export type RequestId = components["schemas"]["RequestId"];
 export const RequestId = {
   create(): RequestId {
     return generateUUID() as RequestId;

package/src/core/network/__tests__/CachingRequestRegistry.test.ts

@@ -1,17 +1,16 @@
 /* Copyright 2026 Marimo. All rights reserved. */
+
 import { beforeEach, describe, expect, it, vi } from "vitest";
+import { requestId } from "@/__tests__/branded";
 import { CachingRequestRegistry } from "../CachingRequestRegistry";
-import {
-  DeferredRequestRegistry,
-  type RequestId,
-} from "../DeferredRequestRegistry";
+import { DeferredRequestRegistry } from "../DeferredRequestRegistry";
 
 vi.mock("@/utils/uuid", () => ({
   generateUUID: vi.fn().mockReturnValue("uuid"),
 }));
 
 describe("CachingRequestRegistry", () => {
-  const REQUEST_ID = "uuid" as RequestId;
+  const REQUEST_ID = requestId("uuid");
   let makeRequestMock = vi.fn();
   let delegate: DeferredRequestRegistry<unknown, unknown>;
   let caching: CachingRequestRegistry<unknown, unknown>;
@@ -31,7 +30,7 @@ describe("CachingRequestRegistry", () => {
 
     // Resolve first request
     delegate.resolve(REQUEST_ID, "response");
-    await expect(p1).resolves.toBe("response");
+    await expect(p1).resolves.toBe(requestId("response"));
 
     // Second call with equivalent request gets served from cache
     const p2 = caching.request({ a: 1 });
@@ -50,8 +49,8 @@ describe("CachingRequestRegistry", () => {
     expect(p1).toStrictEqual(p2);
 
     // Resolve and ensure both resolve to same result
-    delegate.resolve(REQUEST_ID, "ok");
-    await expect(p1).resolves.toBe("ok");
+    delegate.resolve(REQUEST_ID, requestId("ok"));
+    await expect(p1).resolves.toBe(requestId("ok"));
     await expect(p2).resolves.toBe("ok");
   });
 
@@ -67,7 +66,7 @@ describe("CachingRequestRegistry", () => {
     expect(makeRequestMock).toHaveBeenCalledTimes(2);
 
     // Resolve the second request
-    delegate.resolve(REQUEST_ID, "ok");
-    await expect(p2).resolves.toBe("ok");
+    delegate.resolve(REQUEST_ID, requestId("ok"));
+    await expect(p2).resolves.toBe(requestId("ok"));
   });
 });

package/src/core/network/__tests__/DeferredRequestRegistry.test.ts

@@ -1,16 +1,14 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { beforeEach, describe, expect, it, vi } from "vitest";
-import {
-  DeferredRequestRegistry,
-  type RequestId,
-} from "../DeferredRequestRegistry";
+import { requestId } from "@/__tests__/branded";
+import { DeferredRequestRegistry } from "../DeferredRequestRegistry";
 
 vi.mock("@/utils/uuid", () => ({
   generateUUID: vi.fn().mockReturnValue("uuid"),
 }));
 
 describe("DeferredRequestRegistry", () => {
-  const REQUEST_ID = "uuid" as RequestId;
+  const REQUEST_ID = requestId("uuid");
   let makeRequestMock = vi.fn();
   let registry: DeferredRequestRegistry<unknown, unknown>;
 
@@ -26,7 +24,7 @@ describe("DeferredRequestRegistry", () => {
 
     // resolve
     registry.resolve(REQUEST_ID, "response");
-    await expect(promise).resolves.toBe("response");
+    await expect(promise).resolves.toBe(requestId("response"));
   });
 
   it("should handle request failure", async () => {

package/src/core/network/requests-lazy.ts

@@ -119,6 +119,7 @@ const ACTIONS: Record<keyof AllRequests, Action> = {
 
   // Sync operations that wait for connection
   syncCellIds: "waitForConnectionOpen",
+  sendDocumentTransaction: "waitForConnectionOpen",
   sendCodeCompletionRequest: "waitForConnectionOpen",
 };
 

package/src/core/network/requests-network.ts

@@ -58,6 +58,15 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
+    sendDocumentTransaction: async (request) => {
+      await waitForConnectionOpen();
+      return getClient()
+        .POST("/api/document/transaction", {
+          body: request,
+          params: getParams(),
+        })
+        .then(handleResponseReturnNull);
+    },
     sendRename: (request) => {
       return getClient()
         .POST("/api/kernel/rename", {

package/src/core/network/requests-static.ts

@@ -37,6 +37,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
     },
     sendRestart: throwNotInEditMode,
     syncCellIds: throwNotInEditMode,
+    sendDocumentTransaction: throwNotInEditMode,
     sendRun: throwNotInEditMode,
     sendRunScratchpad: throwNotInEditMode,
     sendRename: throwNotInEditMode,

package/src/core/network/requests-toasting.tsx

@@ -22,6 +22,7 @@ export function createErrorToastingRequests(
     sendFunctionRequest: "Failed to send function request",
     sendRestart: "Failed to restart",
     syncCellIds: "Failed to sync cell IDs",
+    sendDocumentTransaction: "Failed to sync document transaction",
     sendRun: "Failed to run",
     sendRunScratchpad: "Failed to run scratchpad",
     sendRename: "Failed to rename",

package/src/core/network/types.ts

@@ -80,6 +80,8 @@ export interface SetCellConfigRequest {
 export type UpdateUIElementRequest = schemas["UpdateUIElementRequest"];
 export type ModelRequest = schemas["ModelRequest"];
 export type UpdateCellIdsRequest = schemas["UpdateCellIdsRequest"];
+export type NotebookDocumentTransactionRequest =
+  schemas["NotebookDocumentTransactionRequest"];
 export type UpdateUserConfigRequest = schemas["UpdateUserConfigRequest"];
 export type ShutdownSessionRequest = schemas["ShutdownSessionRequest"];
 export type Snippet = schemas["Snippet"];
@@ -138,6 +140,9 @@ export interface EditRequests {
   saveCellConfig: (request: SetCellConfigRequest) => Promise<null>;
   sendRestart: () => Promise<null>;
   syncCellIds: (request: UpdateCellIdsRequest) => Promise<null>;
+  sendDocumentTransaction: (
+    request: NotebookDocumentTransactionRequest,
+  ) => Promise<null>;
   sendInstallMissingPackages: (
     request: InstallPackagesRequest,
   ) => Promise<null>;

package/src/core/static/__tests__/virtual-file-tracker.test.ts

@@ -1,6 +1,6 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it } from "vitest";
-import type { CellId } from "@/core/cells/ids";
+import { cellId } from "@/__tests__/branded";
 import type { OutputMessage } from "@/core/kernel/messages";
 import { findVirtualFiles, VirtualFileTracker } from "../virtual-file-tracker";
 
@@ -103,32 +103,32 @@ describe("findVirtualFiles", () => {
 describe("VirtualFileTracker", () => {
   it("should track virtual files, append, clear", () => {
     const tracker = VirtualFileTracker.INSTANCE;
-    const cellId = "test-cell-id" as CellId;
+    const cid = cellId("test-cell-id");
     tracker.track({
-      cell_id: cellId,
+      cell_id: cid,
       output: {
         mimetype: "text/html",
         data: "Some text /@file/test-file.js more text",
       } as OutputMessage,
     });
-    expect(tracker.virtualFiles.get(cellId)).toEqual(
+    expect(tracker.virtualFiles.get(cid)).toEqual(
       new Set(["/@file/test-file.js"]),
     );
 
     // can append
     tracker.track({
-      cell_id: cellId,
+      cell_id: cid,
       output: {
         mimetype: "text/html",
         data: "Some text /@file/another-file.txt more text",
       } as OutputMessage,
     });
-    expect(tracker.virtualFiles.get(cellId)).toEqual(
+    expect(tracker.virtualFiles.get(cid)).toEqual(
       new Set(["/@file/test-file.js", "/@file/another-file.txt"]),
     );
 
     // can clear
-    tracker.removeForCellId(cellId);
-    expect(tracker.virtualFiles.get(cellId)).toEqual(undefined);
+    tracker.removeForCellId(cid);
+    expect(tracker.virtualFiles.get(cid)).toEqual(undefined);
   });
 });

package/src/core/static/virtual-file-tracker.ts

@@ -29,7 +29,7 @@ export class VirtualFileTracker {
 
   track(message: Pick<CellMessage, "cell_id" | "output">): void {
     const output = message.output;
-    const cellId = message.cell_id as CellId;
+    const cellId = message.cell_id;
     if (!output) {
       return;
     }