From 8b673230142b26ae1117d5bfdd8da0c00463792c Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Sun, 3 May 2026 14:16:40 -0700 Subject: [PATCH 1/2] Fix notebook SDK open race after createLocal Signed-off-by: Jeremy lewi --- app/src/lib/notebookData.test.ts | 98 ++++++++++++++++++++++++----- app/src/lib/runtime/appJsGlobals.ts | 25 ++++++++ 2 files changed, 109 insertions(+), 14 deletions(-) diff --git a/app/src/lib/notebookData.test.ts b/app/src/lib/notebookData.test.ts index 2ad663a..3398774 100644 --- a/app/src/lib/notebookData.test.ts +++ b/app/src/lib/notebookData.test.ts @@ -1210,15 +1210,72 @@ describe("NotebookData.runCodeCell", () => { }); it("supports drive.saveAsCurrentNotebook in appkernel cells", async () => { + const notebooksByUri = new Map>(); + const savedNotebooksByUri = new Map(); + let activeUri = "local://file/original"; + + const resolveNotebook = (target?: unknown) => { + if (!target) { + return notebooksByUri.get(activeUri) ?? null; + } + if (typeof target === "string") { + return notebooksByUri.get(target) ?? null; + } + if ( + typeof target === "object" && + target && + "uri" in target && + typeof (target as { uri?: string }).uri === "string" + ) { + return notebooksByUri.get((target as { uri: string }).uri) ?? null; + } + if ( + typeof target === "object" && + target && + "handle" in target && + (target as { handle?: { uri?: string } }).handle?.uri + ) { + return ( + notebooksByUri.get((target as { handle: { uri: string } }).handle.uri) ?? + null + ); + } + return null; + }; + + const listNotebooks = () => Array.from(notebooksByUri.values()); + const createRemote = vi.fn().mockResolvedValue({ uri: "https://drive.google.com/file/d/saveas123/view", }); const saveContent = vi.fn().mockResolvedValue(undefined); appState.setDriveNotebookStore({ create: createRemote, saveContent } as any); const addFile = vi.fn().mockResolvedValue("local://file/saveas-copy"); - const saveLocal = vi.fn().mockResolvedValue(undefined); + const saveLocal = vi.fn().mockImplementation( + async (uri: string, notebook: parser_pb.Notebook) => { + savedNotebooksByUri.set(uri, create(parser_pb.NotebookSchema, notebook)); + } + ); appState.setLocalNotebooks({ addFile, save: saveLocal } as any); - const openNotebook = vi.fn().mockResolvedValue(undefined); + const openNotebook = vi.fn().mockImplementation(async (uri: string) => { + activeUri = uri; + setTimeout(() => { + const saved = savedNotebooksByUri.get(uri); + if (!saved) { + return; + } + const createdModel = new NotebookData({ + notebook: create(parser_pb.NotebookSchema, saved), + uri, + name: "copy.json", + notebookStore: null, + loaded: true, + resolveNotebookForAppKernel: resolveNotebook, + listNotebooksForAppKernel: listNotebooks, + }); + notebooksByUri.set(uri, createdModel); + }, 0); + }); appState.setOpenNotebookHandler(openNotebook); const cell = create(parser_pb.CellSchema, { @@ -1242,7 +1299,10 @@ describe("NotebookData.runCodeCell", () => { name: "saveas-source.json", notebookStore: null, loaded: true, + resolveNotebookForAppKernel: resolveNotebook, + listNotebooksForAppKernel: listNotebooks, }); + notebooksByUri.set(model.getUri(), model); model.runCodeCell(cell); await waitForCondition(() => { @@ -1269,6 +1329,7 @@ describe("NotebookData.runCodeCell", () => { it("supports notebook creation and append helpers in appkernel cells", async () => { const notebooksByUri = new Map>(); + const savedNotebooksByUri = new Map(); let activeUri = "nb://seed"; const resolveNotebook = (target?: unknown) => { @@ -1306,21 +1367,30 @@ describe("NotebookData.runCodeCell", () => { uri: "local://file/helloworld", name: "helloworld", }); - const saveLocal = vi.fn().mockImplementation(async (uri: string, notebook: parser_pb.Notebook) => { - const createdModel = new NotebookData({ - notebook, - uri, - name: "helloworld", - notebookStore: null, - loaded: true, - resolveNotebookForAppKernel: resolveNotebook, - listNotebooksForAppKernel: listNotebooks, - }); - notebooksByUri.set(uri, createdModel); - }); + const saveLocal = vi.fn().mockImplementation( + async (uri: string, notebook: parser_pb.Notebook) => { + savedNotebooksByUri.set(uri, create(parser_pb.NotebookSchema, notebook)); + } + ); appState.setLocalNotebooks({ create: createLocal, save: saveLocal } as any); const openNotebook = vi.fn().mockImplementation(async (uri: string) => { activeUri = uri; + setTimeout(() => { + const saved = savedNotebooksByUri.get(uri); + if (!saved) { + return; + } + const createdModel = new NotebookData({ + notebook: create(parser_pb.NotebookSchema, saved), + uri, + name: "helloworld", + notebookStore: null, + loaded: true, + resolveNotebookForAppKernel: resolveNotebook, + listNotebooksForAppKernel: listNotebooks, + }); + notebooksByUri.set(uri, createdModel); + }, 0); }); appState.setOpenNotebookHandler(openNotebook); diff --git a/app/src/lib/runtime/appJsGlobals.ts b/app/src/lib/runtime/appJsGlobals.ts index 2763d74..f74f5ee 100644 --- a/app/src/lib/runtime/appJsGlobals.ts +++ b/app/src/lib/runtime/appJsGlobals.ts @@ -120,6 +120,13 @@ function createEmptyNotebook(): parser_pb.Notebook { return create(parser_pb.NotebookSchema, { cells: [], metadata: {} }) } +const NOTEBOOK_OPEN_RESOLVE_TIMEOUT_MS = 1500 +const NOTEBOOK_OPEN_RESOLVE_POLL_MS = 10 + +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)) +} + function defaultEnsureFilesystemStore(): FilesystemNotebookStore | null { if (appState.filesystemStore) { return appState.filesystemStore @@ -174,12 +181,30 @@ export function createAppJsGlobals({ appState.removeWorkspaceItem(uri) } const resolveStore = () => resolveNotebookStore?.() ?? appState.localNotebooks + const waitForOpenedNotebook = async (uri: string) => { + const startedAt = Date.now() + while (Date.now() - startedAt < NOTEBOOK_OPEN_RESOLVE_TIMEOUT_MS) { + const candidates = [ + resolveNotebook?.(uri), + resolveNotebook?.({ uri }), + resolveNotebook?.(), + runme.getCurrentNotebook(), + ] + if (candidates.some((notebook) => notebook?.getUri() === uri)) { + return + } + await sleep(NOTEBOOK_OPEN_RESOLVE_POLL_MS) + } + throw new Error(`Timed out waiting for notebook ${uri} to load after opening.`) + } const openNotebookForRuntime = async (uri: string) => { if (openNotebook) { await openNotebook(uri) + await waitForOpenedNotebook(uri) return } await appState.openNotebook(uri) + await waitForOpenedNotebook(uri) } const resolveLocalMirrorStore = () => { if (!appState.localNotebooks) { From de883729fed070712b3f4fbcc456d5413d3fdf00 Mon Sep 17 00:00:00 2001 From: Jeremy lewi Date: Sun, 3 May 2026 14:38:26 -0700 Subject: [PATCH 2/2] Avoid fixed timeout when opening runtime notebooks Signed-off-by: Jeremy lewi --- app/src/lib/runtime/appJsGlobals.ts | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/app/src/lib/runtime/appJsGlobals.ts b/app/src/lib/runtime/appJsGlobals.ts index f74f5ee..c867cc1 100644 --- a/app/src/lib/runtime/appJsGlobals.ts +++ b/app/src/lib/runtime/appJsGlobals.ts @@ -120,7 +120,6 @@ function createEmptyNotebook(): parser_pb.Notebook { return create(parser_pb.NotebookSchema, { cells: [], metadata: {} }) } -const NOTEBOOK_OPEN_RESOLVE_TIMEOUT_MS = 1500 const NOTEBOOK_OPEN_RESOLVE_POLL_MS = 10 function sleep(ms: number): Promise { @@ -182,8 +181,11 @@ export function createAppJsGlobals({ } const resolveStore = () => resolveNotebookStore?.() ?? appState.localNotebooks const waitForOpenedNotebook = async (uri: string) => { - const startedAt = Date.now() - while (Date.now() - startedAt < NOTEBOOK_OPEN_RESOLVE_TIMEOUT_MS) { + // Notebook navigation only updates current-doc first. The concrete + // NotebookData model is created later by React effects, so callers that + // need to operate on the opened notebook must wait for eventual + // materialization instead of failing on a fixed timeout window. + while (true) { const candidates = [ resolveNotebook?.(uri), resolveNotebook?.({ uri }), @@ -195,7 +197,6 @@ export function createAppJsGlobals({ } await sleep(NOTEBOOK_OPEN_RESOLVE_POLL_MS) } - throw new Error(`Timed out waiting for notebook ${uri} to load after opening.`) } const openNotebookForRuntime = async (uri: string) => { if (openNotebook) {