From 8060ca5bfae1e7a3bd9d1816b8756f85e7a1390f Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Wed, 7 Oct 2026 22:42:01 -0700 Subject: [PATCH] perf: persist completed notebook runs in one document update (#26254) --- .../components/editor/ipynb-cell-mutations.ts | 32 ++- .../editor/useIpynbCellExecution.test.tsx | 183 ++++++++++++++++-- .../editor/useIpynbCellExecution.ts | 17 +- 3 files changed, 212 insertions(+), 20 deletions(-) diff --git a/src/renderer/src/components/editor/ipynb-cell-mutations.ts b/src/renderer/src/components/editor/ipynb-cell-mutations.ts index dc754519aef..aaf1f4d9200 100644 --- a/src/renderer/src/components/editor/ipynb-cell-mutations.ts +++ b/src/renderer/src/components/editor/ipynb-cell-mutations.ts @@ -146,6 +146,16 @@ export function updateIpynbCellRun( executionCount: number | null ): string { const root = parseNotebookRoot(content) + applyIpynbCellRun(root, index, outputs, executionCount) + return serializeNotebook(root) +} + +function applyIpynbCellRun( + root: Record, + index: number, + outputs: Record[], + executionCount: number | null +): void { const cell = ensureCell(root, index) // Like Jupyter, store stream text as lines so notebook diffs stay line-oriented. cell.outputs = outputs.map((output) => @@ -154,7 +164,27 @@ export function updateIpynbCellRun( : output ) cell.execution_count = executionCount - return serializeNotebook(root) +} + +export function updateIpynbCellRuns( + content: string, + updates: { index: number; outputs: Record[]; executionCount: number | null }[] +): string { + if (updates.length === 0) { + return content + } + const root = parseNotebookRoot(content) + let changed = false + for (const update of updates) { + try { + applyIpynbCellRun(root, update.index, update.outputs, update.executionCount) + changed = true + } catch { + // Preserve the successful prefix when a later cell no longer exists. + break + } + } + return changed ? serializeNotebook(root) : content } export function clearIpynbOutputs(content: string): string { diff --git a/src/renderer/src/components/editor/useIpynbCellExecution.test.tsx b/src/renderer/src/components/editor/useIpynbCellExecution.test.tsx index 8ec4469134d..8dfb540ff5c 100644 --- a/src/renderer/src/components/editor/useIpynbCellExecution.test.tsx +++ b/src/renderer/src/components/editor/useIpynbCellExecution.test.tsx @@ -1,22 +1,28 @@ // @vitest-environment happy-dom import { act, renderHook, waitFor } from '@testing-library/react' -import { beforeEach, describe, expect, it, vi } from 'vitest' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { parseIpynb } from './ipynb-parse' +import type { KernelFrameEvent } from '../../../../shared/notebook-kernel-types' -const { getConnectionIdMock, notebookApi, toastError } = vi.hoisted(() => { +const { getConnectionIdMock, notebookApi, toastError, frameListener } = vi.hoisted(() => { + const frameListener = { current: (_event: KernelFrameEvent): void => {} } const notebookApi = { listPythonEnvironments: vi.fn(), startKernel: vi.fn(), execute: vi.fn(), shutdownKernel: vi.fn(), - onKernelFrame: vi.fn(() => () => {}) + onKernelFrame: vi.fn((listener: (event: KernelFrameEvent) => void) => { + frameListener.current = listener + return () => {} + }) } // The kernel session subscribes to kernel frames when it loads. Object.defineProperty(window, 'api', { configurable: true, value: { notebook: notebookApi } }) return { getConnectionIdMock: vi.fn((): string | null => null), notebookApi, - toastError: vi.fn() + toastError: vi.fn(), + frameListener } }) @@ -26,6 +32,8 @@ vi.mock('sonner', () => ({ toast: { error: toastError } })) vi.mock('@/store', () => ({ useAppStore: { subscribe: () => () => {} } })) import { useIpynbCellExecution } from './useIpynbCellExecution' +import { runCells, trustNotebook } from './ipynb-kernel-session' +import { getSession } from './ipynb-kernel-store' function notebookContent(withIds: boolean): string { return JSON.stringify({ @@ -63,17 +71,21 @@ function renderExecution(filePath: string, applyContent = vi.fn(), withIds = tru return { hook, applyContent, content: () => content } } -describe('notebook cell execution', () => { - beforeEach(() => { - vi.clearAllMocks() - getConnectionIdMock.mockReturnValue(null) - notebookApi.listPythonEnvironments.mockResolvedValue({ - workspace: [{ path: '/repo/.venv/bin/python', name: '.venv', version: '3.12.1' }], - path: [] - }) - notebookApi.startKernel.mockResolvedValue({ status: 'ready' }) +beforeEach(() => { + vi.clearAllMocks() + getConnectionIdMock.mockReturnValue(null) + notebookApi.listPythonEnvironments.mockResolvedValue({ + workspace: [{ path: '/repo/.venv/bin/python', name: '.venv', version: '3.12.1' }], + path: [] }) + notebookApi.startKernel.mockResolvedValue({ status: 'ready' }) +}) +afterEach(() => { + vi.restoreAllMocks() +}) + +describe('notebook cell execution', () => { it('asks for trust before running, then runs the cell in a kernel', async () => { const { hook } = renderExecution('/repo/trust.ipynb') @@ -117,3 +129,148 @@ describe('notebook cell execution', () => { expect(cell?.key).toBe(cells[1]?.id) }) }) + +function completedNotebookContent(count: number): string { + return JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + metadata: { language_info: { name: 'python' }, custom: { keep: true } }, + cells: Array.from({ length: count }, (_, index) => ({ + id: `cell-${index}`, + cell_type: 'code', + metadata: { tags: ['keep'] }, + execution_count: null, + outputs: [], + source: [`print(${index})\n`] + })) + }) +} + +async function finishHiddenRuns(filePath: string, count: number): Promise { + trustNotebook(filePath) + await runCells( + filePath, + Array.from({ length: count }, (_, index) => ({ + key: `cell-${index}`, + code: `print(${index})` + })), + '/repo' + ) + for (let index = 0; index < count; index += 1) { + frameListener.current({ + filePath, + frame: { type: 'stream', content: { name: 'stdout', text: `α😀 ${index}\nsecond\n` } } + }) + frameListener.current({ + filePath, + frame: { + type: 'display_data', + content: { + data: { 'text/plain': `result ${index}` }, + metadata: { keep: true }, + transient: { display_id: `display-${index}` } + } + } + }) + frameListener.current({ + filePath, + frame: { type: 'done', status: 'ok', execution_count: index + 1 } + }) + } +} + +function mountFinishedNotebook(filePath: string, content: string) { + const applyContent = vi.fn() + const flushSourceDrafts = vi.fn(() => content) + const hook = renderHook(() => + useIpynbCellExecution({ + filePath, + worktreeId: 'worktree-a', + rootPath: '/repo', + flushSourceDrafts, + applyContent + }) + ) + return { hook, applyContent, flushSourceDrafts } +} + +describe('finished notebook runs retained while the viewer is hidden', () => { + it('persists all completed outputs with one serialization when the viewer remounts', async () => { + const filePath = '/repo/hidden-batch.ipynb' + const content = completedNotebookContent(20) + await finishHiddenRuns(filePath, 20) + expect(Object.values(getSession(filePath).runs).every((run) => !run.committed)).toBe(true) + const serialize = vi.spyOn(JSON, 'stringify') + const { hook, applyContent, flushSourceDrafts } = mountFinishedNotebook(filePath, content) + expect(serialize).toHaveBeenCalledTimes(1) + serialize.mockRestore() + expect(flushSourceDrafts).toHaveBeenCalledTimes(1) + expect(applyContent).toHaveBeenCalledTimes(1) + const saved: unknown = applyContent.mock.calls[0]?.[0] + expect(typeof saved).toBe('string') + if (typeof saved !== 'string') { + throw new Error('Expected saved notebook text') + } + const cells = parseIpynb(saved).cells + expect(cells).toHaveLength(20) + for (let index = 0; index < cells.length; index += 1) { + expect(cells[index]?.executionCount).toBe(index + 1) + expect(cells[index]?.outputs[0]).toEqual({ + kind: 'stream', + name: 'stdout', + text: `α😀 ${index}\nsecond\n` + }) + expect(cells[index]?.outputs[1]).toMatchObject({ + kind: 'display', + items: [{ mime: 'text/plain', value: `result ${index}` }] + }) + } + expect(saved).toContain('"custom"') + expect(saved).toContain('"tags"') + expect(saved).not.toContain('"transient"') + expect( + Object.values(getSession(filePath).runs).every( + (run) => run.committed && run.outputs.length === 0 + ) + ).toBe(true) + hook.unmount() + }) + + it('preserves the saved prefix when a later raw cell refuses the update and retires every run', async () => { + const filePath = '/repo/hidden-refused-middle.ipynb' + const content = completedNotebookContent(3) + const raw: unknown = JSON.parse(content) + if (typeof raw !== 'object' || raw === null || !('cells' in raw) || !Array.isArray(raw.cells)) { + throw new Error('Expected notebook cells') + } + raw.cells.splice(1, 0, null) + await finishHiddenRuns(filePath, 3) + const { hook, applyContent } = mountFinishedNotebook(filePath, JSON.stringify(raw)) + expect(applyContent).toHaveBeenCalledTimes(1) + const saved: unknown = applyContent.mock.calls[0]?.[0] + if (typeof saved !== 'string') { + throw new Error('Expected saved prefix') + } + const cells = parseIpynb(saved).cells + expect(cells.map((cell) => cell.executionCount)).toEqual([1, null, null]) + expect(cells.map((cell) => cell.outputs.length)).toEqual([2, 0, 0]) + expect( + Object.values(getSession(filePath).runs).every( + (run) => run.committed && run.outputs.length === 0 + ) + ).toBe(true) + hook.unmount() + }) + + it.each(['{', '{"cells":[null,{"id":"cell-0","cell_type":"code","source":[]}]}'])( + 'retires runs without changing a document that refuses its first update: %s', + async (content) => { + const filePath = `/repo/hidden-refused-first-${content}.ipynb` + await finishHiddenRuns(filePath, 1) + const { hook, applyContent } = mountFinishedNotebook(filePath, content) + expect(applyContent).not.toHaveBeenCalled() + expect(getSession(filePath).runs['cell-0']).toMatchObject({ committed: true, outputs: [] }) + hook.unmount() + } + ) +}) diff --git a/src/renderer/src/components/editor/useIpynbCellExecution.ts b/src/renderer/src/components/editor/useIpynbCellExecution.ts index 42b95a6a6ed..97bf5346d0f 100644 --- a/src/renderer/src/components/editor/useIpynbCellExecution.ts +++ b/src/renderer/src/components/editor/useIpynbCellExecution.ts @@ -2,7 +2,7 @@ import { useEffect, useState } from 'react' import { toast } from 'sonner' import { translate } from '@/i18n/i18n' import { getConnectionId } from '@/lib/connection-context' -import { clearIpynbOutputs, updateIpynbCellRun, withIpynbCellIds } from './ipynb-cell-mutations' +import { clearIpynbOutputs, updateIpynbCellRuns, withIpynbCellIds } from './ipynb-cell-mutations' import { toStoredOutputs } from './ipynb-kernel-outputs' import { forgetFinishedRuns, @@ -48,18 +48,23 @@ export function useIpynbCellExecution({ let nextContent = latestContent try { const cells = parseIpynb(latestContent).cells + const updates: { + index: number + outputs: Record[] + executionCount: number | null + }[] = [] for (const key of uncommittedRunKeys) { const run = getCellRun(filePath, key) const index = cells.findIndex((cell, cellIndex) => getIpynbCellKey(cell, cellIndex) === key) if (run && index !== -1) { - nextContent = updateIpynbCellRun( - nextContent, + updates.push({ index, - toStoredOutputs(run.outputs), - run.executionCount - ) + outputs: toStoredOutputs(run.outputs), + executionCount: run.executionCount + }) } } + nextContent = updateIpynbCellRuns(latestContent, updates) } catch { // An unparseable document has no cells to take outputs; drop them. }