fix(notebook): cancel kernel starts when the notebook closes (#24858)

Co-authored-by: OrcaWin <293788423+OrcaWin@users.noreply.github.com>

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<string, (event: unknown, args: unknown) => unknown>()
@@ -22,22 +22,39 @@ import type { Store } from '../persistence'
 
 function fakeKernel() {
   let onFrame: (frame: KernelFrame) => void = () => {}
+  const exited = Promise.withResolvers<void>()
   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<string>()
+      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<string>()
+    const freshAuthorization = Promise.withResolvers<string>()
+    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<string>(), Promise.withResolvers<string>()]
+      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<string>()
+    const canonicalAuthorization = Promise.withResolvers<string>()
+    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()
+  })
 })

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<WebContents, Map<string, NotebookKernel>>()
+const startCancellations = createSenderScopedRequestCancellations()
+const startsByOwner = new WeakMap<WebContents, Map<string, Set<AbortController>>>()
+
+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<string, NotebookKernel> {
@@ -67,30 +80,55 @@ export function registerNotebookHandlers(store: Store): void {
   ipcMain.handle(
     'notebook:startKernel',
     async (event, args: { filePath: string; python: string }): Promise<KernelStartResult> => {
-      // 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<string, Set<AbortController>>()
+      startsByOwner.set(owner, starts)
+      const pending = starts.get(args.filePath) ?? new Set<AbortController>()
+      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)

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<PythonEnvironments>>(),
+  startKernel: vi.fn<() => Promise<KernelStartResult>>(),
+  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<KernelStartResult>()
+      const freshReply = Promise.withResolvers<KernelStartResult>()
+      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<PythonEnvironments>()
+      const freshReply = Promise.withResolvers<KernelStartResult>()
+      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' })
+    }
+  )
+})

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<string, symbol>()
 
 /** 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<PythonEnvironment | undefined> {
   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