diff --git a/src/main/ipc/notebook.test.ts b/src/main/ipc/notebook.test.ts index 9201793592a..961c652f316 100644 --- a/src/main/ipc/notebook.test.ts +++ b/src/main/ipc/notebook.test.ts @@ -1,5 +1,5 @@ import { EventEmitter } from 'node:events' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import type { KernelFrame } from '../../shared/notebook-kernel-types' const handlers = new Map unknown>() @@ -22,22 +22,39 @@ import type { Store } from '../persistence' function fakeKernel() { let onFrame: (frame: KernelFrame) => void = () => {} + const exited = Promise.withResolvers() const kernel = { execute: vi.fn(), interrupt: vi.fn(), shutdown: vi.fn() } startNotebookKernelMock.mockImplementationOnce((options) => { onFrame = options.onFrame - return { kernel, ready: Promise.resolve({ status: 'ready' }), exited: new Promise(() => {}) } + return { kernel, ready: Promise.resolve({ status: 'ready' }), exited: exited.promise } }) - return { kernel, emit: (frame: KernelFrame) => onFrame(frame) } + return { kernel, emit: (frame: KernelFrame) => onFrame(frame), exit: () => exited.resolve() } } +let nextOwnerId = 0 +const owners: EventEmitter[] = [] + function fakeOwner() { - return Object.assign(new EventEmitter(), { send: vi.fn(), isDestroyed: () => false }) + const owner = Object.assign(new EventEmitter(), { + id: ++nextOwnerId, + send: vi.fn(), + isDestroyed: (): boolean => false + }) + owners.push(owner) + return owner } +afterEach(() => { + for (const owner of owners) { + owner.emit('destroyed') + } + owners.length = 0 +}) + describe('notebook IPC', () => { beforeEach(() => { handlers.clear() - vi.clearAllMocks() + vi.resetAllMocks() resolveAuthorizedPathMock.mockImplementation(async (path: string) => `/real${path}`) // oxlint-disable-next-line typescript/consistent-type-assertions -- SAFETY: the handlers under test only pass the store to the mocked authorizer. registerNotebookHandlers({} as Store) @@ -104,4 +121,122 @@ describe('notebook IPC', () => { expect(second.kernel.execute).toHaveBeenCalledWith('x') expect(first.kernel.execute).not.toHaveBeenCalled() }) + + it.each(['shutdown', 'destroyed', 'did-navigate', 'render-process-gone'])( + 'does not start a kernel after %s while path authorization is pending', + async (boundary) => { + const authorization = Promise.withResolvers() + resolveAuthorizedPathMock.mockReturnValueOnce(authorization.promise) + fakeKernel() + const owner = fakeOwner() + const args = { filePath: '/repo/nb.ipynb', python: '/py' } + const pending = handlers.get('notebook:startKernel')!({ sender: owner }, args) + if (boundary === 'shutdown') { + handlers.get('notebook:shutdownKernel')!({ sender: owner }, args) + } else { + owner.emit(boundary) + } + authorization.resolve('/real/repo/nb.ipynb') + + await expect(pending).resolves.toMatchObject({ status: 'failed' }) + expect(startNotebookKernelMock).not.toHaveBeenCalled() + expect(owner.listenerCount('did-navigate')).toBe(boundary === 'shutdown' ? 1 : 0) + } + ) + + it('refuses a renderer that was already destroyed before invocation', async () => { + fakeKernel() + const owner = fakeOwner() + owner.isDestroyed = () => true + + await expect( + handlers.get('notebook:startKernel')!( + { sender: owner }, + { filePath: '/repo/nb.ipynb', python: '/py' } + ) + ).resolves.toMatchObject({ status: 'failed' }) + expect(resolveAuthorizedPathMock).not.toHaveBeenCalled() + expect(startNotebookKernelMock).not.toHaveBeenCalled() + }) + + it('keeps a reopened pending start cancellable after the old start finishes', async () => { + const oldAuthorization = Promise.withResolvers() + const freshAuthorization = Promise.withResolvers() + resolveAuthorizedPathMock + .mockReturnValueOnce(oldAuthorization.promise) + .mockReturnValueOnce(freshAuthorization.promise) + fakeKernel() + fakeKernel() + const owner = fakeOwner() + const args = { filePath: '/repo/nb.ipynb', python: '/py' } + const start = handlers.get('notebook:startKernel')! + const shutdown = handlers.get('notebook:shutdownKernel')! + const old = start({ sender: owner }, args) + shutdown({ sender: owner }, args) + const fresh = start({ sender: owner }, args) + oldAuthorization.resolve('/real/repo/nb.ipynb') + await expect(old).resolves.toMatchObject({ status: 'failed' }) + shutdown({ sender: owner }, args) + freshAuthorization.resolve('/real/repo/nb.ipynb') + + await expect(fresh).resolves.toMatchObject({ status: 'failed' }) + expect(startNotebookKernelMock).not.toHaveBeenCalled() + expect(owner.listenerCount('did-navigate')).toBe(1) + }) + + it.each([0, 1])( + 'preserves concurrent authorization order %s and ignores a replaced kernel exit', + async (firstIndex) => { + const authorizations = [Promise.withResolvers(), Promise.withResolvers()] + resolveAuthorizedPathMock + .mockReturnValueOnce(authorizations[0].promise) + .mockReturnValueOnce(authorizations[1].promise) + const firstConstructed = fakeKernel() + const lastConstructed = fakeKernel() + const owner = fakeOwner() + const start = handlers.get('notebook:startKernel')! + const args = { filePath: '/repo/nb.ipynb', python: '/py' } + const pending = [start({ sender: owner }, args), start({ sender: owner }, args)] + authorizations[firstIndex].resolve('/real/repo/nb.ipynb') + await expect(pending[firstIndex]).resolves.toEqual({ status: 'ready' }) + authorizations[1 - firstIndex].resolve('/real/repo/nb.ipynb') + await expect(pending[1 - firstIndex]).resolves.toEqual({ status: 'ready' }) + expect(firstConstructed.kernel.shutdown).toHaveBeenCalledOnce() + expect(lastConstructed.kernel.shutdown).not.toHaveBeenCalled() + firstConstructed.exit() + await Promise.resolve() + + handlers.get('notebook:execute')!( + { sender: owner }, + { filePath: args.filePath, code: 'current' } + ) + expect(lastConstructed.kernel.execute).toHaveBeenCalledWith('current') + expect(firstConstructed.kernel.execute).not.toHaveBeenCalled() + } + ) + + it('cancels one raw alias without canceling the canonical path’s pending start', async () => { + const aliasAuthorization = Promise.withResolvers() + const canonicalAuthorization = Promise.withResolvers() + resolveAuthorizedPathMock + .mockReturnValueOnce(aliasAuthorization.promise) + .mockReturnValueOnce(canonicalAuthorization.promise) + const current = fakeKernel() + const owner = fakeOwner() + const start = handlers.get('notebook:startKernel')! + const alias = { filePath: '/tmp/repo/nb.ipynb', python: '/py' } + const canonical = { filePath: '/private/tmp/repo/nb.ipynb', python: '/py' } + const old = start({ sender: owner }, alias) + const fresh = start({ sender: owner }, canonical) + handlers.get('notebook:shutdownKernel')!({ sender: owner }, alias) + aliasAuthorization.resolve(canonical.filePath) + await expect(old).resolves.toMatchObject({ status: 'failed' }) + canonicalAuthorization.resolve(canonical.filePath) + + await expect(fresh).resolves.toEqual({ status: 'ready' }) + handlers.get('notebook:execute')!({ sender: owner }, { ...canonical, code: 'canonical' }) + expect(current.kernel.execute).toHaveBeenCalledWith('canonical') + expect(current.kernel.shutdown).not.toHaveBeenCalled() + expect(startNotebookKernelMock).toHaveBeenCalledOnce() + }) }) diff --git a/src/main/ipc/notebook.ts b/src/main/ipc/notebook.ts index 70f8224011b..f599734a03c 100644 --- a/src/main/ipc/notebook.ts +++ b/src/main/ipc/notebook.ts @@ -1,7 +1,9 @@ +import { randomUUID } from 'node:crypto' import { dirname } from 'node:path' import { ipcMain, type WebContents } from 'electron' import type { Store } from '../persistence' import { resolveAuthorizedPath } from './filesystem-auth' +import { createSenderScopedRequestCancellations } from './sender-scoped-request-cancellation' import { startNotebookKernel, type NotebookKernel } from '../notebook/notebook-kernel' import { createNotebookVenv, @@ -20,6 +22,17 @@ import type { /** Each renderer document's kernels, by notebook file. */ const kernelsByOwner = new Map>() +const startCancellations = createSenderScopedRequestCancellations() +const startsByOwner = new WeakMap>>() + +function cancelPendingStarts(owner: WebContents, filePath: string): void { + const starts = startsByOwner.get(owner) + const pending = starts?.get(filePath) + starts?.delete(filePath) + for (const controller of pending ?? []) { + controller.abort() + } +} // Why: a reloaded, crashed or closed renderer has lost its sessions, so its kernels go with it. function kernelsOf(owner: WebContents): Map { @@ -67,30 +80,55 @@ export function registerNotebookHandlers(store: Store): void { ipcMain.handle( 'notebook:startKernel', async (event, args: { filePath: string; python: string }): Promise => { - // Why: run from the notebook's folder so relative imports and data paths resolve as on disk. - const cwd = dirname(await resolveAuthorizedPath(args.filePath, store)) const owner = event.sender - const kernels = kernelsOf(owner) - kernels.get(args.filePath)?.shutdown() - const { kernel, ready, exited } = startNotebookKernel({ - python: args.python, - cwd, - onFrame: (frame) => { - if (!owner.isDestroyed()) { - owner.send('notebook:kernelFrame', { - filePath: args.filePath, - frame - } satisfies KernelFrameEvent) + if (owner.isDestroyed()) { + return { status: 'failed', detail: 'The notebook closed before its kernel started.' } + } + // Each start stays independent until the notebook or issuing document closes. + const requestToken = randomUUID() + const controller = startCancellations.begin(event, requestToken) + if (!controller) { + return { status: 'failed', detail: 'The notebook closed before its kernel started.' } + } + const starts = startsByOwner.get(owner) ?? new Map>() + startsByOwner.set(owner, starts) + const pending = starts.get(args.filePath) ?? new Set() + starts.set(args.filePath, pending) + pending.add(controller) + try { + // Why: run from the notebook's folder so relative imports and data paths resolve as on disk. + const cwd = dirname(await resolveAuthorizedPath(args.filePath, store)) + if (controller.signal.aborted || owner.isDestroyed()) { + return { status: 'failed', detail: 'The notebook closed before its kernel started.' } + } + const kernels = kernelsOf(owner) + kernels.get(args.filePath)?.shutdown() + const { kernel, ready, exited } = startNotebookKernel({ + python: args.python, + cwd, + onFrame: (frame) => { + if (!owner.isDestroyed()) { + owner.send('notebook:kernelFrame', { + filePath: args.filePath, + frame + } satisfies KernelFrameEvent) + } } + }) + kernels.set(args.filePath, kernel) + void exited.then(() => { + if (kernels.get(args.filePath) === kernel) { + kernels.delete(args.filePath) + } + }) + return await ready + } finally { + pending.delete(controller) + if (pending.size === 0 && starts.get(args.filePath) === pending) { + starts.delete(args.filePath) } - }) - kernels.set(args.filePath, kernel) - void exited.then(() => { - if (kernels.get(args.filePath) === kernel) { - kernels.delete(args.filePath) - } - }) - return ready + startCancellations.finish(event, requestToken, controller) + } } ) @@ -123,6 +161,7 @@ export function registerNotebookHandlers(store: Store): void { }) ipcMain.handle('notebook:shutdownKernel', (event, args: { filePath: string }): void => { + cancelPendingStarts(event.sender, args.filePath) const kernels = kernelsOf(event.sender) kernels.get(args.filePath)?.shutdown() kernels.delete(args.filePath) diff --git a/src/renderer/src/components/editor/ipynb-kernel-session-start-lifetime.test.ts b/src/renderer/src/components/editor/ipynb-kernel-session-start-lifetime.test.ts new file mode 100644 index 00000000000..9d0292895e0 --- /dev/null +++ b/src/renderer/src/components/editor/ipynb-kernel-session-start-lifetime.test.ts @@ -0,0 +1,188 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { toast } from 'sonner' +import type { + KernelFrameEvent, + KernelStartResult, + PythonEnvironments +} from '../../../../shared/notebook-kernel-types' + +type OpenFilesState = { openFiles: { filePath: string }[] } +type AppStoreListener = (state: OpenFilesState, previous: OpenFilesState) => void + +const { appStoreListeners } = vi.hoisted(() => { + const appStoreListeners: AppStoreListener[] = [] + return { appStoreListeners } +}) + +vi.mock('@/i18n/i18n', () => ({ translate: (_key: string, fallback: string) => fallback })) +vi.mock('sonner', () => ({ toast: { error: vi.fn() } })) +vi.mock('@/store', () => ({ + useAppStore: { + subscribe: (listener: AppStoreListener) => { + appStoreListeners.push(listener) + return () => { + const index = appStoreListeners.indexOf(listener) + if (index !== -1) { + appStoreListeners.splice(index, 1) + } + } + } + } +})) + +const FILE = '/notebook.ipynb' +const ENVIRONMENT = { path: '/workspace/.venv/bin/python', name: '.venv' } +let emitFrame: (event: KernelFrameEvent) => void = () => {} +const notebookApi = { + listPythonEnvironments: vi.fn<() => Promise>(), + startKernel: vi.fn<() => Promise>(), + execute: vi.fn(), + shutdownKernel: vi.fn(), + onKernelFrame: (listener: (event: KernelFrameEvent) => void) => { + emitFrame = listener + return () => {} + } +} +Object.defineProperty(globalThis, 'window', { + configurable: true, + value: { api: { notebook: notebookApi } } +}) + +const session = await import('./ipynb-kernel-session') +const { getSession, setEnvironment, store } = await import('./ipynb-kernel-store') + +function closeNotebook(): void { + for (const listener of appStoreListeners) { + listener({ openFiles: [] }, { openFiles: [{ filePath: FILE }] }) + } +} + +function reopenNotebook(): void { + session.trustNotebook(FILE) + setEnvironment(FILE, ENVIRONMENT) +} + +function finishFreshCell(): void { + emitFrame({ + filePath: FILE, + frame: { type: 'stream', content: { name: 'stdout', text: 'fresh output\n' } } + }) + emitFrame({ + filePath: FILE, + frame: { type: 'done', status: 'ok', execution_count: 1 } + }) +} + +beforeEach(() => { + closeNotebook() + store.setState({ environments: {} }) + vi.clearAllMocks() + notebookApi.startKernel.mockReset().mockResolvedValue({ status: 'ready' }) + notebookApi.listPythonEnvironments.mockReset().mockResolvedValue({ + workspace: [ENVIRONMENT], + path: [] + }) + reopenNotebook() +}) + +afterEach(() => closeNotebook()) + +describe('notebook start request session ownership', () => { + it.each([ + ['ready', false], + ['ready', true], + ['failed', false], + ['failed', true], + ['missing', false], + ['missing', true], + ['rejected', false], + ['rejected', true] + ] as const)( + 'ignores old %s completion when replacement ready first is %s', + async (oldOutcome, freshReadyFirst) => { + const oldReply = Promise.withResolvers() + const freshReply = Promise.withResolvers() + notebookApi.startKernel + .mockReturnValueOnce(oldReply.promise) + .mockReturnValueOnce(freshReply.promise) + const old = session.runCells(FILE, [{ key: 'old', code: 'old cell' }], null) + closeNotebook() + reopenNotebook() + const fresh = session.runCells(FILE, [{ key: 'fresh', code: 'fresh cell' }], null) + if (freshReadyFirst) { + freshReply.resolve({ status: 'ready' }) + await fresh + finishFreshCell() + } + const expected = structuredClone(getSession(FILE)) + if (oldOutcome === 'rejected') { + oldReply.reject(new Error('Old start rejected')) + } else { + oldReply.resolve( + oldOutcome === 'ready' + ? { status: 'ready' } + : oldOutcome === 'missing' + ? { status: 'missing-ipykernel', externallyManaged: true } + : { status: 'failed', detail: 'Old start failed' } + ) + } + await old + + expect(getSession(FILE)).toEqual(expected) + expect(store.getState().environments[FILE]).toEqual(ENVIRONMENT) + expect(notebookApi.execute).toHaveBeenCalledTimes(freshReadyFirst ? 1 : 0) + expect(toast.error).not.toHaveBeenCalled() + if (!freshReadyFirst) { + freshReply.resolve({ status: 'ready' }) + await fresh + finishFreshCell() + } + expect(notebookApi.execute).toHaveBeenCalledOnce() + expect(notebookApi.execute).toHaveBeenCalledWith({ filePath: FILE, code: 'fresh cell' }) + expect(getSession(FILE)).toMatchObject({ + status: 'ready', + setup: null, + queue: [], + runs: { fresh: { outputs: [{ output_type: 'stream', text: 'fresh output\n' }] } } + }) + } + ) + + it.each([false, true])( + 'ignores old discovery when replacement ready first is %s', + async (freshReadyFirst) => { + const discovered = Promise.withResolvers() + const freshReply = Promise.withResolvers() + store.setState({ environments: {} }) + notebookApi.listPythonEnvironments.mockReturnValueOnce(discovered.promise) + const old = session.runCells(FILE, [{ key: 'old', code: 'old cell' }], '/workspace') + closeNotebook() + reopenNotebook() + notebookApi.startKernel.mockReturnValueOnce(freshReply.promise) + const fresh = session.runCells(FILE, [{ key: 'fresh', code: 'fresh cell' }], null) + if (freshReadyFirst) { + freshReply.resolve({ status: 'ready' }) + await fresh + finishFreshCell() + } + const expected = structuredClone(getSession(FILE)) + discovered.resolve({ + workspace: [{ path: '/obsolete/bin/python', name: 'obsolete' }], + path: [] + }) + await old + + expect(getSession(FILE)).toEqual(expected) + expect(store.getState().environments[FILE]).toEqual(ENVIRONMENT) + expect(notebookApi.startKernel).toHaveBeenCalledOnce() + expect(toast.error).not.toHaveBeenCalled() + if (!freshReadyFirst) { + freshReply.resolve({ status: 'ready' }) + await fresh + finishFreshCell() + } + expect(notebookApi.execute).toHaveBeenCalledOnce() + expect(notebookApi.execute).toHaveBeenCalledWith({ filePath: FILE, code: 'fresh cell' }) + } + ) +}) diff --git a/src/renderer/src/components/editor/ipynb-kernel-session.ts b/src/renderer/src/components/editor/ipynb-kernel-session.ts index a75f406138a..ea768b78dfe 100644 --- a/src/renderer/src/components/editor/ipynb-kernel-session.ts +++ b/src/renderer/src/components/editor/ipynb-kernel-session.ts @@ -20,6 +20,7 @@ import { } from './ipynb-kernel-store' const INTERRUPT_STALL_MS = 10_000 +const sessionLifetimes = new Map() /** Reports a failure in the first queued cell (a toast when nothing was queued) and drops the queue. */ function failQueue(filePath: string, message: string, detail = ''): void { @@ -59,10 +60,14 @@ function isOpen(filePath: string): boolean { return filePath in store.getState().sessions } +const isCurrentSession = (filePath: string, lifetime: symbol): boolean => + isOpen(filePath) && sessionLifetimes.get(filePath) === lifetime + /** Picks the nearest Python for a notebook that has none: a workspace env, else one on PATH. */ async function discoverEnvironment( filePath: string, - rootPath: string | null + rootPath: string | null, + lifetime: symbol ): Promise { const found = await window.api.notebook.listPythonEnvironments({ filePath, @@ -70,7 +75,7 @@ async function discoverEnvironment( runWorkspaceInterpreters: true }) const recommended = found.workspace[0] ?? found.path[0] - if (recommended && isOpen(filePath)) { + if (recommended && isCurrentSession(filePath, lifetime)) { setEnvironment(filePath, recommended) } return recommended @@ -81,20 +86,23 @@ async function start(filePath: string, rootPath: string | null = null): Promise< if (!getSession(filePath).trusted) { return } + const lifetime = sessionLifetimes.get(filePath) ?? Symbol() + sessionLifetimes.set(filePath, lifetime) // Why 'starting' before discovery: a second run meanwhile must queue, not start another kernel. updateSession(filePath, () => ({ status: 'starting' })) let result: KernelStartResult | null try { const environment = - store.getState().environments[filePath] ?? (await discoverEnvironment(filePath, rootPath)) + store.getState().environments[filePath] ?? + (await discoverEnvironment(filePath, rootPath, lifetime)) result = - environment && isOpen(filePath) + environment && isCurrentSession(filePath, lifetime) ? await window.api.notebook.startKernel({ filePath, python: environment.path }) : null } catch (error) { result = { status: 'failed', detail: error instanceof Error ? error.message : String(error) } } - if (!isOpen(filePath)) { + if (!isCurrentSession(filePath, lifetime)) { return } if (!result) { @@ -323,6 +331,7 @@ useAppStore.subscribe((state, previous) => { } for (const filePath of Object.keys(store.getState().sessions)) { if (!state.openFiles.some((file) => file.filePath === filePath)) { + sessionLifetimes.delete(filePath) void window.api.notebook.shutdownKernel({ filePath }) store.setState(({ sessions }) => { const { [filePath]: _closed, ...rest } = sessions