feat(notebook): quiet connection/identity slot for cloud and desktop (#3421 PR 3)

apps/mcp-app/package.json

@@ -7,7 +7,7 @@
   },
   "dependencies": {
     "@modelcontextprotocol/ext-apps": "^1.3.0",
-    "@modelcontextprotocol/sdk": "^1.28.0",
+    "@modelcontextprotocol/sdk": "^1.29.0",
     "react": "^19.1.0",
     "react-dom": "^19.1.0"
   },

apps/notebook-cloud/scripts/hosted-widget-cross-window-smoke.mjs

@@ -277,9 +277,8 @@ async function openNotebookShell(page, href, timeout) {
 }
 
 async function waitForWidgetSlider(page, timeout) {
-  const cell = page.locator('[data-cell-type="code"]').first();
-  await cell.waitFor({ state: "visible", timeout });
-  const slider = cell
+  await page.locator('[data-cell-type="code"]').first().waitFor({ state: "visible", timeout });
+  const slider = page
     .frameLocator('[data-slot="isolated-frame"]')
     .locator('[data-widget-type="IntSlider"]')
     .getByRole("slider")

apps/notebook-cloud/src/cloudflare-types.ts

@@ -52,6 +52,15 @@ export interface DurableObjectState {
   waitUntil(promise: Promise<unknown>): void;
   acceptWebSocket?(socket: CloudflareWebSocket, tags?: string[]): void;
   getWebSockets?(tag?: string): CloudflareWebSocket[];
+  // Optional like the other hibernation APIs: the runtime answers matching
+  // text messages without waking the DO; fakes in tests need not implement
+  // it (the room feature-detects before calling).
+  setWebSocketAutoResponse?(pair: WebSocketRequestResponsePair): void;
+}
+
+export interface WebSocketRequestResponsePair {
+  readonly request: string;
+  readonly response: string;
 }
 
 export interface DurableObjectStorage {

apps/notebook-cloud/src/index.ts

@@ -128,6 +128,11 @@ const VIEWER_RUNTIME_WASM_ASSET_MANIFEST_PATH = "/assets/runtime-wasm-assets.jso
 const VIEWER_RUNTIMED_WASM_MODULE_NAME = "runtimed_wasm.js";
 const VIEWER_RUNTIMED_WASM_NAME = "runtimed_wasm_bg.wasm";
 const SNAPSHOT_BLOB_HEAD_CONCURRENCY = 16;
+// One R2 HEAD is issued per referenced blob during snapshot-pair validation.
+// Cap the total so a single publish cannot fan out unbounded billable R2
+// operations. Sized ~10x the largest blob_ref_count observed in
+// snapshot_pair.validation.completed logs; raise it if legitimate notebooks hit it.
+const MAX_SNAPSHOT_BLOB_REFS = 2000;
 const ULID_ALPHABET = "0123456789ABCDEFGHJKMNPQRSTVWXYZ";
 const CREATE_NOTEBOOK_ID_ATTEMPTS = 8;
 
@@ -2807,6 +2812,29 @@ async function validateSnapshotPair(options: {
     };
   }
 
+  const blobRefs = snapshotBlobRefsOverCap(render);
+  if (blobRefs.over) {
+    cloudLog("warn", "snapshot_pair.validation.blob_refs_over_cap", {
+      notebook_id: options.notebookId,
+      notebook_heads_hash: options.notebookHeadsHash,
+      runtime_heads_hash: options.runtimeHeadsHash,
+      duration_ms: durationMs(startedAt),
+      blob_ref_count: blobRefs.count,
+      blob_ref_cap: blobRefs.cap,
+      counter: "snapshot_pair_validation_blob_ref_cap_rejections",
+      counter_delta: 1,
+    });
+    return {
+      ok: false,
+      status: 422,
+      body: {
+        error: "snapshot references too many blobs",
+        blob_ref_count: blobRefs.count,
+        blob_ref_cap: blobRefs.cap,
+      },
+    };
+  }
+
   const missingBlobs = await findMissingSnapshotBlobs(bucket, options.notebookId, render);
   if (missingBlobs.length > 0) {
     cloudLog("warn", "snapshot_pair.validation.missing_blobs", {
@@ -2848,6 +2876,9 @@ async function findMissingSnapshotBlobs(
   render: unknown,
 ): Promise<MissingSnapshotBlob[]> {
   const refs = collectSnapshotBlobRefs(render);
+  if (refs.length > MAX_SNAPSHOT_BLOB_REFS) {
+    throw new Error(`snapshot blob ref count ${refs.length} exceeds cap ${MAX_SNAPSHOT_BLOB_REFS}`);
+  }
   const missing: Array<MissingSnapshotBlob | null> = [];
 
   for (let index = 0; index < refs.length; index += SNAPSHOT_BLOB_HEAD_CONCURRENCY) {
@@ -2884,6 +2915,14 @@ function collectSnapshotBlobRefs(render: unknown): BlobRef[] {
   return Object.values(refs);
 }
 
+export function snapshotBlobRefsOverCap(
+  render: unknown,
+  cap = MAX_SNAPSHOT_BLOB_REFS,
+): { count: number; cap: number; over: boolean } {
+  const count = collectSnapshotBlobRefs(render).length;
+  return { count, cap, over: count > cap };
+}
+
 function isRecord(value: unknown): value is Record<string, unknown> {
   return typeof value === "object" && value !== null && !Array.isArray(value);
 }

apps/notebook-cloud/src/notebook-room.ts

@@ -1,4 +1,9 @@
-import type { CloudflareWebSocket, DurableObjectState, Env } from "./cloudflare-types.ts";
+import type {
+  CloudflareWebSocket,
+  DurableObjectState,
+  Env,
+  WebSocketRequestResponsePair,
+} from "./cloudflare-types.ts";
 import type { WorkstationAttachmentState } from "runtimed";
 import { identityDisplayLabel } from "./display-label.ts";
 import {
@@ -15,6 +20,8 @@ import {
   frameSizeLimits,
   frameTypeName,
   isClientWritableFrame,
+  LIVENESS_PING,
+  LIVENESS_PONG,
   splitTypedFrame,
   type SessionControlMessage,
   type TypedFrame,
@@ -117,6 +124,23 @@ export class NotebookRoom {
     private readonly state: DurableObjectState,
     private readonly env: Env,
   ) {
+    // CF-native liveness: the runtime answers the client's text ping for
+    // hibernatable sockets WITHOUT waking the DO. The pair persists across
+    // hibernation, so re-setting it per wake is idempotent; matching pings
+    // never reach webSocketMessage, so frame budgets and the hibernation
+    // restore path are untouched. Feature-detected like the other
+    // hibernation APIs (handleMessage answers manually as the fallback).
+    const pairCtor = (
+      globalThis as {
+        WebSocketRequestResponsePair?: new (
+          request: string,
+          response: string,
+        ) => WebSocketRequestResponsePair;
+      }
+    ).WebSocketRequestResponsePair;
+    if (pairCtor && this.state.setWebSocketAutoResponse) {
+      this.state.setWebSocketAutoResponse(new pairCtor(LIVENESS_PING, LIVENESS_PONG));
+    }
     this.restoredPeersReady = this.restoreHibernatedPeers();
     this.state.waitUntil(this.restoredPeersReady);
   }
@@ -383,6 +407,17 @@ export class NotebookRoom {
     message: string | ArrayBuffer | ArrayBufferView,
   ): Promise<void> {
     const incomingByteLength = webSocketMessageByteLength(message);
+    if (message === LIVENESS_PING) {
+      // Fallback for non-hibernation sockets and runtimes without
+      // setWebSocketAutoResponse — must answer BEFORE the binary-only
+      // rejection so pings never count toward consecutiveRejectedFrames.
+      try {
+        peer.socket.send(LIVENESS_PONG);
+      } catch {
+        // socket is closing; the close handler owns cleanup
+      }
+      return;
+    }
     if (typeof message === "string") {
       this.recordFrameBudget(notebookId, peer, "incoming", "text", incomingByteLength);
       this.rejectFrame(

apps/notebook-cloud/src/protocol.ts

@@ -11,6 +11,17 @@ export type { FrameSizeLimits, FrameTypeValue };
 
 export const NOTEBOOK_PROTOCOL = "v4";
 
+/**
+ * Liveness probe text messages. The client pings on an interval; the room
+ * DO answers via the runtime's WebSocket auto-response (no DO wake), with
+ * a manual fallback in the room's message handler for runtimes without
+ * auto-response support. Text (not typed binary frames) on purpose: the
+ * auto-response API matches string messages, and the typed-frame channel
+ * stays binary-only.
+ */
+export const LIVENESS_PING = "nteract-liveness-ping";
+export const LIVENESS_PONG = "nteract-liveness-pong";
+
 export interface TypedFrame {
   type: FrameTypeValue;
   payload: Uint8Array;

apps/notebook-cloud/test/cloud-projection-flicker.test.ts

@@ -0,0 +1,170 @@
+import { afterEach, describe, it } from "node:test";
+import assert from "node:assert/strict";
+import { readFileSync } from "node:fs";
+import { getCellIdsSnapshot } from "@/components/notebook/state/cell-store";
+import {
+  getCellExecutionId,
+  getExecutionById,
+  resetNotebookExecutions,
+} from "@/components/notebook/state/execution-store";
+import { getOutputById, resetNotebookOutputs } from "@/components/notebook/state/output-store";
+import {
+  projectCloudCellsIntoNotebookViewStores,
+  resetCloudProjectionUnlessPreserved,
+  resetCloudViewStoreProjection,
+} from "../viewer/notebook-view-store-bridge.ts";
+import type { ResolvedCell } from "../viewer/render-resolution.ts";
+
+// Field-observed flicker: with IndexedDB seeding, the persisted snapshot
+// paints cells (and their outputs) well before the live-room effect
+// settles. The effect re-runs for reasons that are NOT a notebook switch,
+// and clearing ANY of the projected stores on those re-runs blanks the
+// notebook into a full→empty→full flash. CodeCell renders outputs and
+// execution counts exclusively through the execution/output stores, so the
+// gate must keep or clear ALL projected stores together — these tests run
+// the real gate against the real stores.
+describe("cloud projection flicker gate", () => {
+  const PAINTED = "id:nb-1";
+
+  afterEach(() => {
+    resetCloudViewStoreProjection();
+    resetNotebookExecutions();
+    resetNotebookOutputs();
+  });
+
+  function paintNotebook(): void {
+    projectCloudCellsIntoNotebookViewStores([
+      {
+        id: "cell-code",
+        cellType: "code",
+        source: "print('painted')",
+        language: "python",
+        executionId: "exec-1",
+        executionCount: 3,
+        outputs: [
+          {
+            output_type: "stream",
+            name: "stdout",
+            text: "painted output\n",
+            output_id: "out-1",
+          },
+        ],
+        metadata: {},
+      },
+      {
+        id: "cell-md",
+        cellType: "markdown",
+        source: "# Painted",
+        language: null,
+        executionId: null,
+        executionCount: null,
+        outputs: [],
+        metadata: {},
+      },
+    ] satisfies ResolvedCell[]);
+  }
+
+  function assertPainted(): void {
+    assert.deepEqual(getCellIdsSnapshot(), ["cell-code", "cell-md"]);
+    assert.equal(getCellExecutionId("cell-code"), "exec-1");
+    assert.equal(getExecutionById("exec-1")?.execution_count, 3);
+    assert.deepEqual(getExecutionById("exec-1")?.output_ids, ["out-1"]);
+    const output = getOutputById("out-1");
+    assert.equal(output?.output_type, "stream");
+  }
+
+  it("preserves cells, outputs, and execution pointers across a same-notebook re-run", () => {
+    paintNotebook();
+    assertPainted();
+
+    const preserved = resetCloudProjectionUnlessPreserved({
+      paintedNotebookIdentity: PAINTED,
+      nextNotebookIdentity: PAINTED,
+    });
+
+    assert.equal(preserved, true);
+    // The whole painted surface survives — not just the cell list. Wiping
+    // the execution/output stores while keeping cells still flickers every
+    // output and execution count (the dominant visual mass).
+    assertPainted();
+  });
+
+  it("clears every projected store on a real notebook switch", () => {
+    paintNotebook();
+
+    const preserved = resetCloudProjectionUnlessPreserved({
+      paintedNotebookIdentity: PAINTED,
+      nextNotebookIdentity: "id:nb-2",
+    });
+
+    assert.equal(preserved, false);
+    assert.deepEqual(getCellIdsSnapshot(), [] as string[]);
+    assert.equal(getCellExecutionId("cell-code"), null);
+    assert.equal(getExecutionById("exec-1"), undefined);
+    assert.equal(getOutputById("out-1"), undefined);
+  });
+
+  it("fails closed when no painted identity was recorded", () => {
+    paintNotebook();
+
+    const preserved = resetCloudProjectionUnlessPreserved({
+      paintedNotebookIdentity: null,
+      nextNotebookIdentity: PAINTED,
+    });
+
+    assert.equal(preserved, false);
+    assert.deepEqual(getCellIdsSnapshot(), [] as string[]);
+    assert.equal(getOutputById("out-1"), undefined);
+  });
+
+  it("does not preserve an empty projection", () => {
+    const preserved = resetCloudProjectionUnlessPreserved({
+      paintedNotebookIdentity: PAINTED,
+      nextNotebookIdentity: PAINTED,
+    });
+
+    assert.equal(preserved, false);
+    assert.deepEqual(getCellIdsSnapshot(), [] as string[]);
+  });
+
+  // Source pins for the session wiring that cannot run under node (the hook
+  // imports the component-bearing notebook surface): the cleanup and the
+  // next run's body must both route through the shared gate.
+  describe("cloud-viewer-session wiring", () => {
+    const sessionSource = readFileSync(
+      new URL("../viewer/cloud-viewer-session.ts", import.meta.url),
+      "utf8",
+    );
+
+    it("clears real notebook switches in the next run's body, before connecting", () => {
+      // The cleanup closes over its own run's config, so switch-clearing
+      // can only happen here — and a cleared switch also drops the painted
+      // identity so later gates fail closed.
+      assert.match(
+        sessionSource,
+        /const preservedAcrossRuns = resetCloudProjectionUnlessPreserved\(\{\s*paintedNotebookIdentity: paintedNotebookIdentityRef\.current,\s*nextNotebookIdentity: `id:\$\{config\.notebookId\}`,\s*\}\);\s*if \(!preservedAcrossRuns\) \{\s*paintedNotebookIdentityRef\.current = null;\s*\}/,
+      );
+    });
+
+    it("gates the cleanup on the shared store gate with only the pool reset unconditional", () => {
+      assert.match(
+        sessionSource,
+        /resetCloudProjectionUnlessPreserved\(\{\s*paintedNotebookIdentity: paintedNotebookIdentityRef\.current,\s*nextNotebookIdentity: `id:\$\{config\.notebookId\}`,\s*\}\);\s*resetPoolState\(\);/,
+      );
+    });
+
+    it("tracks the painted notebook identity only when cells actually painted", () => {
+      assert.match(
+        sessionSource,
+        /if \(resolvedCells\.length > 0\) \{\s*paintedNotebookIdentityRef\.current = `id:\$\{config\.notebookId\}`;\s*\}/,
+      );
+    });
+
+    it("keeps an unconditional full clear on true unmount", () => {
+      assert.match(
+        sessionSource,
+        /useEffect\(\s*\(\) => \(\) => \{\s*resetCloudViewStoreProjection\(\);\s*resetRuntimeState\(\);\s*resetRuntimeStoresProjection\(\);\s*\},\s*\[\],\s*\);/,
+      );
+    });
+  });
+});