From c594cb467cc5d6dd3f53e0a729c9e4589f5d38cc Mon Sep 17 00:00:00 2001 From: imjszhang Date: Wed, 22 Jul 2026 16:16:06 +0800 Subject: [PATCH 1/4] feat(notebook): import ipynb sessions Co-authored-by: Cursor --- src/main/notebook/ipc.test.ts | 10 + src/main/notebook/ipc.ts | 6 + src/main/notebook/ipynb-import.test.ts | 207 ++++++++++++++++++ src/main/notebook/ipynb-import.ts | 206 +++++++++++++++++ src/main/notebook/repository.ts | 20 ++ .../notebook/runtime-service.import.test.ts | 133 +++++++++++ src/main/notebook/runtime-service.ts | 74 ++++++- src/preload/index.d.ts | 2 + src/preload/index.ts | 4 + ...SessionNotebookDialog.interaction.test.tsx | 46 +++- .../SessionNotebookDialog.render.test.tsx | 10 +- .../pages/workspace/SessionNotebookDialog.tsx | 56 ++++- src/renderer/web/api-map.generated.ts | 1 + src/shared/notebook.ts | 17 +- 14 files changed, 781 insertions(+), 11 deletions(-) create mode 100644 src/main/notebook/ipynb-import.test.ts create mode 100644 src/main/notebook/ipynb-import.ts create mode 100644 src/main/notebook/runtime-service.import.test.ts diff --git a/src/main/notebook/ipc.test.ts b/src/main/notebook/ipc.test.ts index c1ddcae3f..26b8c456a 100644 --- a/src/main/notebook/ipc.test.ts +++ b/src/main/notebook/ipc.test.ts @@ -51,6 +51,7 @@ describe('notebook IPC handlers', () => { }), runCell: vi.fn().mockResolvedValue({ runId: 'run-2', status: 'completed' }), exportIpynb: vi.fn().mockResolvedValue({ saved: true, filePath: '/tmp/session.ipynb' }), + importIpynb: vi.fn().mockResolvedValue({ imported: true, cellCount: 1 }), beginCodeCell: vi.fn().mockResolvedValue({ cellId: 'cell-1', writeId: 'write-1' }), appendCodeCell: vi.fn().mockResolvedValue({ receivedBytes: 5 }), finishCodeCell: vi.fn().mockResolvedValue({ status: 'idle' }), @@ -94,6 +95,7 @@ describe('notebook IPC handlers', () => { workspaceCwd: '/workspace', kernel: 'python' }) + await handlers.importIpynb({ sessionId: 'session-1', workspaceCwd: '/workspace' }) expect(service.execute).toHaveBeenCalledWith({ sessionId: 'session-1', @@ -120,6 +122,10 @@ describe('notebook IPC handlers', () => { workspaceCwd: '/workspace', kernel: 'python' }) + expect(service.importIpynb).toHaveBeenCalledWith({ + sessionId: 'session-1', + workspaceCwd: '/workspace' + }) }) it('registers every notebook channel and forwards the renderer payload unchanged', async () => { @@ -132,6 +138,7 @@ describe('notebook IPC handlers', () => { runCell: vi.fn().mockResolvedValue({ runId: 'run-1', status: 'completed' }), execute: vi.fn().mockResolvedValue({ runId: 'run-2', status: 'completed' }), exportIpynb: vi.fn().mockResolvedValue({ saved: false }), + importIpynb: vi.fn().mockResolvedValue({ imported: false }), restart: vi.fn().mockResolvedValue({ sessionId: 'session-1' }), shutdown: vi.fn().mockResolvedValue({ sessionId: 'session-1', status: 'shutdown' }) } as unknown as NotebookRuntimeService @@ -147,6 +154,7 @@ describe('notebook IPC handlers', () => { 'notebook:execute', 'notebook:export-ipynb', 'notebook:export-ipynb-all', + 'notebook:import-ipynb', 'notebook:restart', 'notebook:shutdown' ]) @@ -178,6 +186,7 @@ describe('notebook IPC handlers', () => { await ipcHandlers.get('notebook:run-cell')?.(undefined, run) await ipcHandlers.get('notebook:execute')?.(undefined, execute) await ipcHandlers.get('notebook:export-ipynb')?.(undefined, session) + await ipcHandlers.get('notebook:import-ipynb')?.(undefined, session) await ipcHandlers.get('notebook:restart')?.(undefined, session) await ipcHandlers.get('notebook:shutdown')?.(undefined, session) @@ -189,6 +198,7 @@ describe('notebook IPC handlers', () => { expect(service.runCell).toHaveBeenCalledWith(publicRun) expect(service.execute).toHaveBeenCalledWith(publicExecute) expect(service.exportIpynb).toHaveBeenCalledWith(session) + expect(service.importIpynb).toHaveBeenCalledWith(session) expect(service.restart).toHaveBeenCalledWith(session) expect(service.shutdown).toHaveBeenCalledWith(session) }) diff --git a/src/main/notebook/ipc.ts b/src/main/notebook/ipc.ts index 2d5f1773f..b67518653 100644 --- a/src/main/notebook/ipc.ts +++ b/src/main/notebook/ipc.ts @@ -9,6 +9,7 @@ import type { ExportNotebookKernelRequest, ExportNotebookResult, FinishNotebookCodeCellRequest, + ImportNotebookResult, NotebookRunSummary, NotebookSessionReference, NotebookSessionRequest, @@ -34,6 +35,7 @@ type NotebookHandlers = { execute: (request: ExecuteNotebookCodeRequest) => Promise exportIpynb: (request: ExportNotebookKernelRequest) => Promise exportIpynbAll: (request: ExportNotebookAllRequest) => Promise + importIpynb: (request: NotebookSessionRequest) => Promise restart: (request: NotebookSessionRequest) => Promise shutdown: (request: NotebookSessionRequest) => ReturnType } @@ -61,6 +63,7 @@ const createNotebookHandlers = (service: NotebookRuntimeService): NotebookHandle withDataRootWrite(() => service.execute(withoutTrustedTurnContext(request))), exportIpynb: (request) => service.exportIpynb(request), exportIpynbAll: (request) => service.exportIpynbAll(request), + importIpynb: (request) => service.importIpynb(request), restart: (request) => withDataRootWrite(() => service.restart(request)), shutdown: (request) => withDataRootWrite(() => service.shutdown(request)) }) @@ -96,6 +99,9 @@ const registerNotebookIpcHandlers = (service: NotebookRuntimeService): void => { ipcMain.handle('notebook:export-ipynb-all', (_event, request: ExportNotebookAllRequest) => handlers.exportIpynbAll(request) ) + ipcMain.handle('notebook:import-ipynb', (_event, request: NotebookSessionRequest) => + handlers.importIpynb(request) + ) ipcMain.handle('notebook:restart', (_event, request: NotebookSessionRequest) => handlers.restart(request) ) diff --git a/src/main/notebook/ipynb-import.test.ts b/src/main/notebook/ipynb-import.test.ts new file mode 100644 index 000000000..e86d9ef54 --- /dev/null +++ b/src/main/notebook/ipynb-import.test.ts @@ -0,0 +1,207 @@ +import { describe, expect, it } from 'vitest' + +import type { NotebookRunDocument } from '../../shared/notebook' +import { runDocumentToIpynb } from './ipynb-export' +import { ipynbToRunRecords, type IpynbImportResult } from './ipynb-import' + +const context = { + importedAt: 1_000, + createId: (() => { + let value = 0 + return () => String(++value) + })() +} + +const importNotebook = (notebook: unknown): IpynbImportResult => + ipynbToRunRecords(notebook, { + importedAt: context.importedAt, + createId: context.createId + }) + +describe('ipynbToRunRecords', () => { + it('validates the nbformat major version and cells array', () => { + expect(() => importNotebook({ nbformat: 3, cells: [] })).toThrow('expected nbformat 4') + expect(() => importNotebook({ nbformat: 4 })).toThrow('cells must be an array') + }) + + it('imports code cells, skips markdown, and reconstructs all supported outputs', () => { + const result = importNotebook({ + nbformat: 4, + nbformat_minor: 5, + metadata: { kernelspec: { name: 'python3', language: 'python' } }, + cells: [ + { cell_type: 'markdown', source: ['# Title\n'] }, + { + cell_type: 'code', + source: ['print("hello")\n', '2 + 2'], + execution_count: 4, + metadata: { + open_science: { kernel: 'python', environment: 'analysis' } + }, + outputs: [ + { output_type: 'stream', name: 'stdout', text: ['hello\n'] }, + { output_type: 'stream', name: 'stderr', text: 'warning\n' }, + { + output_type: 'error', + ename: 'ValueError', + evalue: 'bad value', + traceback: ['line 1\n', 'line 2'] + }, + { + output_type: 'display_data', + data: { 'image/png': 'aW1hZ2U=', 'text/plain': '' }, + metadata: {} + }, + { + output_type: 'display_data', + data: { 'application/json': { answer: 42 } }, + metadata: {} + }, + { + output_type: 'execute_result', + data: { 'text/plain': ['4'] }, + metadata: {}, + execution_count: 4 + } + ] + } + ] + }) + + expect(result.skippedCellCount).toBe(1) + expect(result.runs).toHaveLength(1) + expect(result.runs[0]).toMatchObject({ + source: 'user', + inputKind: 'cell', + kernelKind: 'python', + script: 'print("hello")\n2 + 2', + status: 'imported', + startedAt: 1_000, + executionCount: 4, + environment: 'analysis', + text: { + stdout: 'hello\n', + stderr: 'warning\n', + traceback: 'line 1\nline 2' + } + }) + expect(result.runs[0].outputs).toEqual([ + { type: 'stream', name: 'stdout', text: 'hello\n' }, + { type: 'stream', name: 'stderr', text: 'warning\n' }, + { + type: 'error', + name: 'ValueError', + message: 'bad value', + traceback: 'line 1\nline 2' + }, + { type: 'display', data: { 'image/png': 'aW1hZ2U=', 'text/plain': '' } }, + { type: 'json', data: { answer: 42 } }, + { type: 'text', text: '4' } + ]) + }) + + it('uses kernelspec fallback and restores downgraded bash/repl source markers', () => { + const result = importNotebook({ + nbformat: 4, + metadata: { kernelspec: { name: 'ir', language: 'R' } }, + cells: [ + { + cell_type: 'code', + source: 'print(1)', + execution_count: null, + metadata: {}, + outputs: [] + }, + { + cell_type: 'code', + source: ['%%bash\n', 'pwd'], + execution_count: null, + metadata: { tags: ['open-science-bash'] }, + outputs: [] + }, + { + cell_type: 'code', + source: '%%javascript\nawait host.mcp()', + execution_count: null, + metadata: { open_science: { kernel: 'repl' } }, + outputs: [] + } + ] + }) + + expect(result.runs.map(({ kernelKind, script }) => ({ kernelKind, script }))).toEqual([ + { kernelKind: 'r', script: 'print(1)' }, + { kernelKind: 'bash', script: 'pwd' }, + { kernelKind: 'repl', script: 'await host.mcp()' } + ]) + }) + + it('round-trips the supported Open Science subset in both directions', async () => { + const source = { + nbformat: 4, + nbformat_minor: 5, + metadata: { kernelspec: { display_name: 'Python 3', name: 'python3', language: 'python' } }, + cells: [ + { + cell_type: 'code', + id: 'source-cell', + source: ['x = 1\n', 'x'], + execution_count: 7, + metadata: { open_science: { kernel: 'python', environment: 'analysis' } }, + outputs: [ + { + output_type: 'execute_result', + data: { 'text/plain': '1' }, + metadata: {}, + execution_count: 7 + } + ] + } + ] + } + const imported = importNotebook(source) + const document: NotebookRunDocument = { + version: 1, + projectName: 'default-project', + sessionId: 'session-1', + workspaceCwd: '/workspace', + notebookSessionRoot: '/storage/notebooks/default-project/session-1', + dataRoot: '/storage/notebooks/default-project/session-1/data', + kernel: { + language: 'python', + kernelName: 'python3', + runtimeRoot: '/storage/runtime', + lastKnownStatus: 'idle' + }, + runs: imported.runs, + updatedAt: 1_000 + } + + const exported = await runDocumentToIpynb(document) + expect(exported.cells[0]).toMatchObject({ + source: source.cells[0].source, + execution_count: 7, + outputs: source.cells[0].outputs, + metadata: { + open_science: { kernel: 'python', environment: 'analysis', status: 'imported' } + } + }) + + const reimported = importNotebook(exported) + expect( + reimported.runs.map(({ script, kernelKind, executionCount, outputs }) => ({ + script, + kernelKind, + executionCount, + outputs + })) + ).toEqual( + imported.runs.map(({ script, kernelKind, executionCount, outputs }) => ({ + script, + kernelKind, + executionCount, + outputs + })) + ) + }) +}) diff --git a/src/main/notebook/ipynb-import.ts b/src/main/notebook/ipynb-import.ts new file mode 100644 index 000000000..24226a6e7 --- /dev/null +++ b/src/main/notebook/ipynb-import.ts @@ -0,0 +1,206 @@ +import type { + NotebookKernelKind, + NotebookOutput, + NotebookRunRecord, + NotebookTextOutput +} from '../../shared/notebook' + +type IpynbImportContext = { + createId: () => string + importedAt: number +} + +type IpynbImportResult = { + runs: NotebookRunRecord[] + skippedCellCount: number +} + +type JsonObject = Record + +const isObject = (value: unknown): value is JsonObject => + typeof value === 'object' && value !== null && !Array.isArray(value) + +const textValue = (value: unknown, field: string): string => { + if (typeof value === 'string') return value + if (Array.isArray(value) && value.every((part) => typeof part === 'string')) { + return value.join('') + } + throw new Error(`Invalid .ipynb ${field}: expected a string or string array.`) +} + +const optionalObject = (value: unknown): JsonObject => (isObject(value) ? value : {}) + +const notebookKernel = (notebook: JsonObject): 'python' | 'r' => { + const metadata = optionalObject(notebook.metadata) + const kernelspec = optionalObject(metadata.kernelspec) + const name = typeof kernelspec.name === 'string' ? kernelspec.name.toLowerCase() : '' + const language = typeof kernelspec.language === 'string' ? kernelspec.language.toLowerCase() : '' + return name === 'ir' || language === 'r' ? 'r' : 'python' +} + +const cellKernel = (cell: JsonObject, fallback: 'python' | 'r'): NotebookKernelKind => { + const metadata = optionalObject(cell.metadata) + const openScience = optionalObject(metadata.open_science) + const kernel = openScience.kernel + if (kernel === 'python' || kernel === 'r' || kernel === 'repl' || kernel === 'bash') { + return kernel + } + const tags = Array.isArray(metadata.tags) ? metadata.tags : [] + if (tags.includes('open-science-bash')) return 'bash' + if (tags.includes('open-science-repl')) return 'repl' + return fallback +} + +const stripKernelMarker = (source: string, kernel: NotebookKernelKind): string => { + if (kernel === 'bash' && source.startsWith('%%bash\n')) return source.slice('%%bash\n'.length) + if (kernel === 'repl' && source.startsWith('%%javascript\n')) { + return source.slice('%%javascript\n'.length) + } + return source +} + +const mimeText = (value: unknown): string | null => { + if (typeof value === 'string') return value + if (Array.isArray(value) && value.every((part) => typeof part === 'string')) { + return value.join('') + } + return null +} + +const mapDisplayData = (dataValue: unknown, executeResult: boolean): NotebookOutput | null => { + if (!isObject(dataValue)) return null + const entries = Object.entries(dataValue) + if (entries.length === 1 && entries[0][0] === 'application/json') { + return { type: 'json', data: entries[0][1] } + } + if (executeResult && entries.length === 1 && entries[0][0] === 'text/plain') { + const text = mimeText(entries[0][1]) + return text === null ? null : { type: 'text', text } + } + + const data: Record = {} + for (const [mime, value] of entries) { + const text = mimeText(value) + if (text !== null) data[mime] = text + } + return Object.keys(data).length > 0 ? { type: 'display', data } : null +} + +const mapOutput = (value: unknown): NotebookOutput | null => { + if (!isObject(value) || typeof value.output_type !== 'string') return null + switch (value.output_type) { + case 'stream': { + const name = value.name === 'stderr' ? 'stderr' : 'stdout' + return { type: 'stream', name, text: textValue(value.text, 'stream output') } + } + case 'error': + return { + type: 'error', + name: typeof value.ename === 'string' ? value.ename : undefined, + message: typeof value.evalue === 'string' ? value.evalue : undefined, + traceback: textValue(value.traceback ?? '', 'error traceback') + } + case 'display_data': + return mapDisplayData(value.data, false) + case 'execute_result': + return mapDisplayData(value.data, true) + default: + return null + } +} + +const textProjection = (outputs: NotebookOutput[]): NotebookTextOutput => { + const stdout = outputs + .filter( + (output): output is Extract => + output.type === 'stream' && output.name === 'stdout' + ) + .map((output) => output.text) + .join('') + const stderr = outputs + .filter( + (output): output is Extract => + output.type === 'stream' && output.name === 'stderr' + ) + .map((output) => output.text) + .join('') + const traceback = outputs + .filter( + (output): output is Extract => output.type === 'error' + ) + .map((output) => output.traceback) + .join('\n') + + return { + stdout, + stderr, + traceback, + plain: [stdout, stderr].filter((text) => text.trim().length > 0) + } +} + +const readExecutionCount = (value: unknown): number | undefined => + typeof value === 'number' && Number.isInteger(value) && value >= 0 ? value : undefined + +// Parses the supported nbformat 4 subset into durable, not-yet-executed run records. ID/time +// generation is injected so the projection is deterministic in tests. +const ipynbToRunRecords = ( + notebookValue: unknown, + context: IpynbImportContext +): IpynbImportResult => { + if (!isObject(notebookValue) || notebookValue.nbformat !== 4) { + throw new Error('Unsupported .ipynb format: expected nbformat 4.') + } + if (!Array.isArray(notebookValue.cells)) { + throw new Error('Invalid .ipynb: cells must be an array.') + } + + const fallbackKernel = notebookKernel(notebookValue) + const runs: NotebookRunRecord[] = [] + let skippedCellCount = 0 + + for (const cellValue of notebookValue.cells) { + if (!isObject(cellValue) || cellValue.cell_type !== 'code') { + skippedCellCount += 1 + continue + } + const kernelKind = cellKernel(cellValue, fallbackKernel) + const source = stripKernelMarker(textValue(cellValue.source, 'cell source'), kernelKind) + const outputs = Array.isArray(cellValue.outputs) + ? cellValue.outputs + .map(mapOutput) + .filter((output): output is NotebookOutput => output !== null) + : [] + const metadata = optionalObject(cellValue.metadata) + const openScience = optionalObject(metadata.open_science) + const id = context.createId() + const environment = + (kernelKind === 'python' || kernelKind === 'r') && + typeof openScience.environment === 'string' && + openScience.environment.trim() + ? openScience.environment + : undefined + + runs.push({ + runId: `imported-run-${id}`, + cellId: `imported-cell-${id}`, + source: 'user', + inputKind: 'cell', + kernelKind, + script: source, + status: 'imported', + startedAt: context.importedAt, + executionCount: readExecutionCount(cellValue.execution_count), + text: textProjection(outputs), + outputs, + artifacts: [], + workingFiles: [], + ...(environment ? { environment } : {}) + }) + } + + return { runs, skippedCellCount } +} + +export { ipynbToRunRecords } +export type { IpynbImportContext, IpynbImportResult } diff --git a/src/main/notebook/repository.ts b/src/main/notebook/repository.ts index ded834fa9..5ec032a97 100644 --- a/src/main/notebook/repository.ts +++ b/src/main/notebook/repository.ts @@ -28,6 +28,12 @@ type AppendNotebookRunRequest = { run: NotebookRunRecord } +type AppendNotebookRunsRequest = { + projectName: string + sessionId: string + runs: NotebookRunRecord[] +} + type UpdateNotebookRunRequest = AppendNotebookRunRequest type UpdateKernelStatusRequest = { @@ -220,6 +226,19 @@ class NotebookRunRepository { })) } + // Appends an imported notebook in one queued read-modify-write turn, avoiding one run.json rewrite + // per cell while preserving the same normalization and serialization guarantees as appendRun. + async appendRuns(request: AppendNotebookRunsRequest): Promise { + return this.mutate(request.projectName, request.sessionId, (document) => ({ + ...document, + runs: [ + ...document.runs, + ...request.runs.map((run) => normalizeRun(document.notebookSessionRoot, run)) + ], + updatedAt: Date.now() + })) + } + // Replaces an existing execution record, used to turn the initial "running" entry final. async updateRun(request: UpdateNotebookRunRequest): Promise { return this.mutate(request.projectName, request.sessionId, (document) => { @@ -403,6 +422,7 @@ export { } export type { AppendNotebookRunRequest, + AppendNotebookRunsRequest, LoadNotebookRunDocumentRequest, UpdateKernelStatusRequest, UpdateNotebookRunRequest diff --git a/src/main/notebook/runtime-service.import.test.ts b/src/main/notebook/runtime-service.import.test.ts new file mode 100644 index 000000000..24438fa66 --- /dev/null +++ b/src/main/notebook/runtime-service.import.test.ts @@ -0,0 +1,133 @@ +import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +import { afterEach, describe, expect, it, vi } from 'vitest' + +import type { NotebookRunDocument } from '../../shared/notebook' +import type { NotebookRunRepository } from './repository' +import { NotebookRuntimeService } from './runtime-service' + +const roots: string[] = [] + +afterEach(async () => { + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))) +}) + +const document = (root: string): NotebookRunDocument => ({ + version: 1, + projectName: 'default-project', + sessionId: 'session-1', + workspaceCwd: '/workspace', + notebookSessionRoot: join(root, 'notebooks', 'default-project', 'session-1'), + dataRoot: join(root, 'notebooks', 'default-project', 'session-1', 'data'), + kernel: { + language: 'python', + kernelName: 'python3', + runtimeRoot: join(root, 'runtime'), + lastKnownStatus: 'idle' + }, + runs: [], + updatedAt: 1 +}) + +const executor = { + execute: vi.fn(), + shutdown: vi.fn().mockResolvedValue({ reaped: true }) +} + +describe('NotebookRuntimeService importIpynb', () => { + it('imports in one repository append and exposes data cells for rerun', async () => { + const root = await mkdtemp(join(tmpdir(), 'open-science-ipynb-import-')) + roots.push(root) + const filePath = join(root, 'source.ipynb') + await writeFile( + filePath, + JSON.stringify({ + nbformat: 4, + metadata: { kernelspec: { name: 'python3', language: 'python' } }, + cells: [ + { + cell_type: 'code', + source: 'print(1)', + execution_count: null, + metadata: {}, + outputs: [] + }, + { cell_type: 'markdown', source: '# ignored' } + ] + }) + ) + const stored = document(root) + const repository = { + loadOrCreate: vi.fn().mockResolvedValue(stored), + appendRuns: vi.fn().mockImplementation(async ({ runs }) => ({ ...stored, runs })) + } as unknown as NotebookRunRepository + const service = new NotebookRuntimeService({ + configRoot: join(root, 'config'), + dataRoot: root, + projectName: 'default-project', + repository, + executorFactory: () => executor, + pickIpynb: async () => filePath + }) + + const result = await service.importIpynb({ + sessionId: 'session-1', + workspaceCwd: '/workspace' + }) + + expect(result).toEqual({ imported: true, cellCount: 1, skippedCellCount: 1 }) + expect(repository.appendRuns).toHaveBeenCalledOnce() + const importedRuns = vi.mocked(repository.appendRuns).mock.calls[0][0].runs + expect(importedRuns[0]).toMatchObject({ + script: 'print(1)', + status: 'imported', + kernelKind: 'python' + }) + const state = await service.state({ sessionId: 'session-1', workspaceCwd: '/workspace' }) + expect(state.cells).toEqual([ + expect.objectContaining({ + id: importedRuns[0].cellId, + language: 'python', + code: 'print(1)', + status: 'idle' + }) + ]) + }) + + it('returns a cancellation result without reading or creating a session', async () => { + const repository = { + loadOrCreate: vi.fn() + } as unknown as NotebookRunRepository + const service = new NotebookRuntimeService({ + configRoot: '/config', + dataRoot: '/storage', + projectName: 'default-project', + repository, + pickIpynb: async () => null + }) + + await expect( + service.importIpynb({ sessionId: 'session-1', workspaceCwd: '/workspace' }) + ).resolves.toEqual({ imported: false }) + expect(repository.loadOrCreate).not.toHaveBeenCalled() + }) + + it('reports invalid JSON as an import error', async () => { + const root = await mkdtemp(join(tmpdir(), 'open-science-ipynb-import-')) + roots.push(root) + const filePath = join(root, 'broken.ipynb') + await writeFile(filePath, '{') + const service = new NotebookRuntimeService({ + configRoot: join(root, 'config'), + dataRoot: root, + projectName: 'default-project', + pickIpynb: async () => filePath + }) + + await expect( + service.importIpynb({ sessionId: 'session-1', workspaceCwd: '/workspace' }) + ).rejects.toThrow('Could not read .ipynb') + }) +}) diff --git a/src/main/notebook/runtime-service.ts b/src/main/notebook/runtime-service.ts index aea2144ad..e9bdc0be0 100644 --- a/src/main/notebook/runtime-service.ts +++ b/src/main/notebook/runtime-service.ts @@ -16,6 +16,7 @@ import type { ExportNotebookKernelRequest, ExportNotebookResult, FinishNotebookCodeCellRequest, + ImportNotebookResult, NotebookEnvironmentStatus, NotebookEnvironmentManifest, NotebookRunEnvironmentCapture, @@ -49,6 +50,7 @@ import { type NbformatOutput, type ResolvedArtifact } from './ipynb-export' +import { ipynbToRunRecords } from './ipynb-import' import { NotebookKernelExecutor, type NotebookKernelExecutorOptions } from './kernel-executor' import { saveIpynbAll } from './save-ipynb-all' import type { KernelProcessKind } from './kernel-executor' @@ -404,6 +406,8 @@ type NotebookRuntimeServiceOptions = { | 'markPackageMutationDirty' | 'refreshAfterPackageMutation' > + // Native file-picker seam for notebook import tests. + pickIpynb?: () => Promise } // The wire binding plus the interpreter override the executor needs. `resolvedInterpreter` is set only @@ -474,6 +478,16 @@ const saveIpynbWithDialog = async ( return { saved: true, filePath } } +const pickIpynbWithDialog = async (): Promise => { + const { dialog } = await import('electron') + const { canceled, filePaths } = await dialog.showOpenDialog({ + title: 'Import notebook', + properties: ['openFile'], + filters: [{ name: 'Jupyter Notebook', extensions: ['ipynb'] }] + }) + return canceled ? null : (filePaths[0] ?? null) +} + // Writes one .ipynb per data kernel under a user-picked directory. Used by the "Download all" path; // the per-tab path (a single .ipynb) goes through `saveIpynbWithDialog` instead. The actual // orchestration (directory picker, conflict check, partial-write cleanup) lives in save-ipynb-all @@ -2496,6 +2510,52 @@ class NotebookRuntimeService { return (this.options.saveIpynbAll ?? saveIpynbAll)(files) } + // Imports code cells as durable, not-yet-executed records and exposes data-kernel cells for rerun. + async importIpynb(request: NotebookSessionRequest): Promise { + const filePath = await (this.options.pickIpynb ?? pickIpynbWithDialog)() + if (!filePath) return { imported: false } + + let notebook: unknown + try { + notebook = JSON.parse(await readFile(filePath, 'utf8')) as unknown + } catch (error) { + const message = error instanceof Error ? error.message : String(error) + throw new Error(`Could not read .ipynb: ${message}`) + } + + const imported = ipynbToRunRecords(notebook, { + createId: randomUUID, + importedAt: Date.now() + }) + const session = await this.ensureSession(request) + if (imported.runs.length > 0) { + await this.repository.appendRuns({ + projectName: session.projectName, + sessionId: session.sessionId, + runs: imported.runs + }) + for (const run of imported.runs) { + if (run.kernelKind !== 'python' && run.kernelKind !== 'r') continue + session.cells.push({ + id: run.cellId, + language: run.kernelKind, + code: run.script, + status: 'idle', + executionCount: run.executionCount, + latestRunId: run.runId + }) + } + session.executionCount += imported.runs.length + this.notifyNotebookChanged(session) + } + + return { + imported: true, + cellCount: imported.runs.length, + skippedCellCount: imported.skippedCellCount + } + } + // Replaces the interpreter process while preserving cells and durable run history. Prefers the // executor's own in-place restart (keeps the same instance, e.g. NotebookKernelExecutor tears down // and lazily respawns its loops) and only shuts down + recreates for executors that don't support it. @@ -3702,7 +3762,19 @@ class NotebookRuntimeService { dataRoot: document.dataRoot, runtimeRoot: document.kernel.runtimeRoot, runJsonPath: getNotebookRunJsonPath(this.options.dataRoot, projectName, request.sessionId), - cells: [], + cells: document.runs + .filter( + (run) => + run.status === 'imported' && (run.kernelKind === 'python' || run.kernelKind === 'r') + ) + .map((run) => ({ + id: run.cellId, + language: run.kernelKind as NotebookLanguage, + code: run.script, + status: 'idle' as const, + executionCount: run.executionCount, + latestRunId: run.runId + })), executionCount: document.runs.length, executor: this.createExecutor(request.sessionId, projectName), executionQueues: new Map(), diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index 0d3ef6306..b0c7dba86 100644 --- a/src/preload/index.d.ts +++ b/src/preload/index.d.ts @@ -71,6 +71,7 @@ import type { ExportNotebookKernelRequest, ExportNotebookResult, FinishNotebookCodeCellRequest, + ImportNotebookResult, NotebookLanguage, NotebookRunSummary, NotebookSessionReference, @@ -560,6 +561,7 @@ interface OpenScienceAPI { execute(request: ExecuteNotebookCodeRequest): Promise exportIpynb(request: ExportNotebookKernelRequest): Promise exportIpynbAll(request: ExportNotebookAllRequest): Promise + importIpynb(request: NotebookSessionRequest): Promise restart(request: NotebookSessionRequest): Promise shutdown(request: NotebookSessionRequest): Promise<{ sessionId: string; status: 'shutdown' }> onAvailable(listener: AcpListener): RemoveListener diff --git a/src/preload/index.ts b/src/preload/index.ts index c8162179e..af6d08c2a 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -73,6 +73,7 @@ import type { ExportNotebookKernelRequest, ExportNotebookResult, FinishNotebookCodeCellRequest, + ImportNotebookResult, NotebookLanguage, NotebookRunSummary, NotebookSessionReference, @@ -619,6 +620,7 @@ type OpenScienceAPI = { execute: (request: ExecuteNotebookCodeRequest) => Promise exportIpynb: (request: ExportNotebookKernelRequest) => Promise exportIpynbAll: (request: ExportNotebookAllRequest) => Promise + importIpynb: (request: NotebookSessionRequest) => Promise restart: (request: NotebookSessionRequest) => Promise shutdown: ( request: NotebookSessionRequest @@ -1254,6 +1256,8 @@ const api: OpenScienceAPI = { ipcRenderer.invoke('notebook:export-ipynb', request) as Promise, exportIpynbAll: (request) => ipcRenderer.invoke('notebook:export-ipynb-all', request) as Promise, + importIpynb: (request) => + ipcRenderer.invoke('notebook:import-ipynb', request) as Promise, restart: (request) => ipcRenderer.invoke('notebook:restart', request) as Promise, shutdown: (request) => diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx index 6a97e08dd..1310a9d97 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx @@ -51,6 +51,7 @@ describe('SessionNotebookContent export', () => { onClose={vi.fn()} onExport={onExport} onExportAll={vi.fn()} + onImport={vi.fn()} /> ) }) @@ -72,7 +73,10 @@ describe('SessionNotebookContent export', () => { it('passes the clicked tab kernel to the export callback after switching tabs', async () => { const onExport = vi.fn().mockResolvedValue(undefined) - const mixedRuns: NotebookRunRecord[] = [run, { ...run, runId: 'r1', kernelKind: 'r', environment: 'default-r' }] + const mixedRuns: NotebookRunRecord[] = [ + run, + { ...run, runId: 'r1', kernelKind: 'r', environment: 'default-r' } + ] await act(async () => { root.render( { onClose={vi.fn()} onExport={onExport} onExportAll={vi.fn()} + onImport={vi.fn()} /> ) }) @@ -118,6 +123,7 @@ describe('SessionNotebookContent export', () => { onClose={vi.fn()} onExport={onExport} onExportAll={vi.fn()} + onImport={vi.fn()} /> ) }) @@ -146,6 +152,7 @@ describe('SessionNotebookContent export', () => { onClose={vi.fn()} onExport={onExport} onExportAll={vi.fn()} + onImport={vi.fn()} /> ) }) @@ -172,6 +179,7 @@ describe('SessionNotebookContent export', () => { onClose={vi.fn()} onExport={failingExport} onExportAll={vi.fn()} + onImport={vi.fn()} /> ) }) @@ -195,6 +203,7 @@ describe('SessionNotebookContent export', () => { onClose={vi.fn()} onExport={vi.fn()} onExportAll={vi.fn()} + onImport={vi.fn()} /> ) }) @@ -204,7 +213,10 @@ describe('SessionNotebookContent export', () => { it('invokes onExportAll for the "Download all" button on mixed sessions', async () => { const onExportAll = vi.fn().mockResolvedValue(undefined) - const mixedRuns: NotebookRunRecord[] = [run, { ...run, runId: 'r1', kernelKind: 'r', environment: 'default-r' }] + const mixedRuns: NotebookRunRecord[] = [ + run, + { ...run, runId: 'r1', kernelKind: 'r', environment: 'default-r' } + ] await act(async () => { root.render( { onClose={vi.fn()} onExport={vi.fn()} onExportAll={onExportAll} + onImport={vi.fn()} /> ) }) @@ -241,6 +254,7 @@ describe('SessionNotebookContent export', () => { onClose={vi.fn()} onExport={vi.fn()} onExportAll={onExportAll} + onImport={vi.fn()} /> ) }) @@ -251,4 +265,32 @@ describe('SessionNotebookContent export', () => { expect(allButton).toBeNull() expect(onExportAll).not.toHaveBeenCalled() }) + + it('invokes import for an empty notebook and surfaces failures', async () => { + const onImport = vi.fn().mockRejectedValue(new Error('Invalid notebook')) + await act(async () => { + root.render( + + ) + }) + + const button = container.querySelector('button[aria-label="Import .ipynb"]') + expect(button?.disabled).toBe(false) + await act(async () => { + button?.click() + await Promise.resolve() + }) + + expect(onImport).toHaveBeenCalledOnce() + expect(container.querySelector('[role="alert"]')?.textContent).toBe('Invalid notebook') + expect(button?.disabled).toBe(false) + }) }) diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx index e68b4610a..c3c689f55 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx @@ -28,7 +28,13 @@ const renderContent = (props: { error?: string }): string => renderToStaticMarkup( - + ) describe('SessionNotebookContent', () => { @@ -113,6 +119,8 @@ describe('SessionNotebookContent', () => { )?.[0] expect(populatedButton).not.toMatch(/\sdisabled(?:=|\s|>)/) expect(emptyButton).toMatch(/\sdisabled(?:=|\s|>)/) + const emptyImportButton = empty.match(/]*aria-label="Import \.ipynb"[^>]*>/)?.[0] + expect(emptyImportButton).not.toMatch(/\sdisabled(?:=|\s|>)/) }) it('hides the "Download all" button when the session has only one data kernel', () => { diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx index fcd715553..5d90cf160 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx @@ -1,5 +1,5 @@ import { useEffect, useState } from 'react' -import { Download, LoaderCircle, X } from 'lucide-react' +import { Download, LoaderCircle, Upload, X } from 'lucide-react' import { Dialog } from 'radix-ui' import { dialogOverlayClassName, dialogPanelClassName } from '@/components/ui/dialog-chrome' @@ -94,6 +94,7 @@ type SessionNotebookContentProps = { onClose: () => void onExport: (kernel: NotebookKernelKind) => Promise onExportAll: () => Promise + onImport: () => Promise } // Pure presentational body of the dialog: header summary, empty/loading/error/populated states, @@ -107,13 +108,16 @@ const SessionNotebookContent = ({ error, onClose, onExport, - onExportAll + onExportAll, + onImport }: SessionNotebookContentProps): React.JSX.Element => { const [activeKind, setActiveKind] = useState('python') const [exporting, setExporting] = useState(false) const [exportingAll, setExportingAll] = useState(false) const [exportError, setExportError] = useState() const [exportSuccess, setExportSuccess] = useState() + const [importing, setImporting] = useState(false) + const [importError, setImportError] = useState() const shortId = sessionId.slice(0, 8) const agents = runs.some((run) => run.source === 'agent') ? 1 : 0 // Only python/r runs are "cells" in the notebook sense; repl/bash are control-plane/shell runs @@ -137,8 +141,9 @@ const SessionNotebookContent = ({ ? activeKind : (KERNEL_KIND_ORDER.find((kind) => kindsWithRuns.has(kind)) ?? visibleKinds[0] ?? 'python') const visibleRuns = runs.filter((run) => resolveRunKernelKind(run) === effectiveActiveKind) - const busy = exporting || exportingAll + const busy = exporting || exportingAll || importing const exportDisabled = status !== 'ready' || runs.length === 0 || busy + const importDisabled = status !== 'ready' || busy // The main button's "current tab" = the kernel whose .ipynb will be saved. repl/bash tabs fold // into the most recent data kernel so the file still has a real kernelspec; sessions that never @@ -180,6 +185,18 @@ const SessionNotebookContent = ({ } } + const handleImport = async (): Promise => { + setImporting(true) + setImportError(undefined) + try { + await onImport() + } catch (importFailure) { + setImportError(getErrorMessage(importFailure)) + } finally { + setImporting(false) + } + } + // The "Download all" path is only useful when there's more than one data kernel to write; a // single-kernel session's secondary button would just duplicate the main button. The data-kernel // count comes from `kindsWithRuns` (control-plane kinds don't generate their own .ipynb). @@ -272,13 +289,29 @@ const SessionNotebookContent = ({

- {exportError ?? exportSuccess} + {importError ?? exportError ?? exportSuccess}

+ {/* Secondary action: only when there's more than one data kernel to write, otherwise it would just duplicate the main button. The "Download all (N)" label surfaces the count so the user knows how many files they're about to create. */} @@ -464,6 +497,17 @@ const SessionNotebookDialog = ({ } return undefined }} + onImport={async () => { + const request = { + sessionId: session.id, + projectName: session.projectId, + workspaceCwd: session.cwd ?? '' + } + const result = await window.api.notebook.importIpynb(request) + if (!result.imported) return + setRuns(await loadSessionNotebookRuns(window.api.notebook, request)) + setStatus('ready') + }} /> ) : null} diff --git a/src/renderer/web/api-map.generated.ts b/src/renderer/web/api-map.generated.ts index 1656546bc..de556e554 100644 --- a/src/renderer/web/api-map.generated.ts +++ b/src/renderer/web/api-map.generated.ts @@ -59,6 +59,7 @@ export const WEB_INVOKE_CHANNELS = { 'notebook.exportIpynbAll': 'notebook:export-ipynb-all', 'notebook.finishCodeCell': 'notebook:finish-code-cell', 'notebook.getReference': 'notebook:reference', + 'notebook.importIpynb': 'notebook:import-ipynb', 'notebook.readInputPreview': 'notebook:read-input-preview', 'notebook.restart': 'notebook:restart', 'notebook.runCell': 'notebook:run-cell', diff --git a/src/shared/notebook.ts b/src/shared/notebook.ts index 1586f7f67..2bf50993d 100644 --- a/src/shared/notebook.ts +++ b/src/shared/notebook.ts @@ -14,7 +14,14 @@ export type NotebookRunInputKind = 'cell' | 'terminal' // died (crash / force-quit) while the run was in flight — reconciled from a stale 'running' on the // next startup. 'cancelled' = the run was deliberately aborted (e.g. a force-stop disable). export type NotebookRunStatus = - 'queued' | 'running' | 'completed' | 'failed' | 'timeout' | 'interrupted' | 'cancelled' + | 'queued' + | 'running' + | 'completed' + | 'failed' + | 'timeout' + | 'interrupted' + | 'cancelled' + | 'imported' export type NotebookRunProvenanceContext = { rootFrameId: string @@ -501,6 +508,14 @@ export type ExportNotebookAllResult = files: Array<{ kernel: 'python' | 'r'; filePath: string }> } +export type ImportNotebookResult = + | { imported: false } + | { + imported: true + cellCount: number + skippedCellCount: number + } + // Starts a streamed code write into a notebook cell. export type BeginNotebookCodeCellRequest = NotebookSessionRequest & { cellId?: string From 5b4c3f31ad9b4e3cee441045b59d61f7edd60915 Mon Sep 17 00:00:00 2001 From: imjszhang Date: Fri, 31 Jul 2026 00:36:31 +0800 Subject: [PATCH 2/4] fix(notebook): align ipynb import with official data-root and UI patterns Wrap import in withDataRootWrite, use dialogSession for the picker flow, and surface import success plus an imported badge. Co-authored-by: Cursor --- src/main/notebook/ipc.ts | 2 +- src/main/notebook/ipynb-import.test.ts | 72 ++++++++++++++++++- src/main/notebook/runtime-service.ts | 4 ++ ...SessionNotebookDialog.interaction.test.tsx | 39 ++++++++++ .../SessionNotebookDialog.render.test.tsx | 12 ++++ .../pages/workspace/SessionNotebookDialog.tsx | 40 +++++++---- 6 files changed, 154 insertions(+), 15 deletions(-) diff --git a/src/main/notebook/ipc.ts b/src/main/notebook/ipc.ts index b67518653..c5c74fdd0 100644 --- a/src/main/notebook/ipc.ts +++ b/src/main/notebook/ipc.ts @@ -63,7 +63,7 @@ const createNotebookHandlers = (service: NotebookRuntimeService): NotebookHandle withDataRootWrite(() => service.execute(withoutTrustedTurnContext(request))), exportIpynb: (request) => service.exportIpynb(request), exportIpynbAll: (request) => service.exportIpynbAll(request), - importIpynb: (request) => service.importIpynb(request), + importIpynb: (request) => withDataRootWrite(() => service.importIpynb(request)), restart: (request) => withDataRootWrite(() => service.restart(request)), shutdown: (request) => withDataRootWrite(() => service.shutdown(request)) }) diff --git a/src/main/notebook/ipynb-import.test.ts b/src/main/notebook/ipynb-import.test.ts index e86d9ef54..c613b27d4 100644 --- a/src/main/notebook/ipynb-import.test.ts +++ b/src/main/notebook/ipynb-import.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from 'vitest' -import type { NotebookRunDocument } from '../../shared/notebook' -import { runDocumentToIpynb } from './ipynb-export' +import type { NotebookRunDocument, NotebookRunRecord } from '../../shared/notebook' +import { runDocumentToIpynb, runDocumentToIpynbForKernel } from './ipynb-export' import { ipynbToRunRecords, type IpynbImportResult } from './ipynb-import' const context = { @@ -204,4 +204,72 @@ describe('ipynbToRunRecords', () => { })) ) }) + + it('imports a tab-scoped python export without pulling in sibling R cells', () => { + const makeRun = (overrides: Partial): NotebookRunRecord => ({ + runId: 'run', + cellId: 'cell', + source: 'user', + kernelKind: 'python', + script: 'print("py")', + status: 'completed', + startedAt: 1, + executionCount: 1, + text: { stdout: 'py\n', stderr: '', traceback: '', plain: ['py\n'] }, + outputs: [{ type: 'stream', name: 'stdout', text: 'py\n' }], + artifacts: [], + workingFiles: [], + environment: 'default-python', + ...overrides + }) + const document: NotebookRunDocument = { + version: 1, + projectName: 'default-project', + sessionId: 'session-1', + workspaceCwd: '/workspace', + notebookSessionRoot: '/storage/notebooks/default-project/session-1', + dataRoot: '/storage/notebooks/default-project/session-1/data', + kernel: { + language: 'python', + kernelName: 'python3', + runtimeRoot: '/storage/runtime', + lastKnownStatus: 'idle' + }, + runs: [ + makeRun({ runId: 'py-1', cellId: 'py-1', script: 'print("py")' }), + makeRun({ + runId: 'bash-1', + cellId: 'bash-1', + kernelKind: 'bash', + script: 'pwd', + environment: undefined, + text: { stdout: '/tmp\n', stderr: '', traceback: '', plain: ['/tmp\n'] }, + outputs: [{ type: 'stream', name: 'stdout', text: '/tmp\n' }] + }), + makeRun({ + runId: 'r-1', + cellId: 'r-1', + kernelKind: 'r', + script: 'print("r")', + environment: 'default-r', + text: { stdout: 'r\n', stderr: '', traceback: '', plain: ['r\n'] }, + outputs: [{ type: 'stream', name: 'stdout', text: 'r\n' }] + }) + ], + updatedAt: 1_000 + } + + const pythonNotebook = runDocumentToIpynbForKernel(document, 'python') + const imported = importNotebook(pythonNotebook) + + expect(pythonNotebook.metadata.kernelspec.name).toBe('python3') + expect(imported.runs.map((run) => ({ kernelKind: run.kernelKind, script: run.script }))).toEqual( + [ + { kernelKind: 'python', script: 'print("py")' }, + { kernelKind: 'bash', script: 'pwd' } + ] + ) + expect(imported.runs.every((run) => run.status === 'imported')).toBe(true) + expect(imported.runs.some((run) => run.kernelKind === 'r')).toBe(false) + }) }) diff --git a/src/main/notebook/runtime-service.ts b/src/main/notebook/runtime-service.ts index e9bdc0be0..f8bf5cf29 100644 --- a/src/main/notebook/runtime-service.ts +++ b/src/main/notebook/runtime-service.ts @@ -2511,6 +2511,10 @@ class NotebookRuntimeService { } // Imports code cells as durable, not-yet-executed records and exposes data-kernel cells for rerun. + // Persisted `environment` is historical annotation only — runCell still resolves the session binding + // (v4). No artifact provenance / inputFiles are written here; the first real execution uses the + // existing capture path. executionCount advances by imported.runs.length so the session counter + // stays monotonic after a bulk append. async importIpynb(request: NotebookSessionRequest): Promise { const filePath = await (this.options.pickIpynb ?? pickIpynbWithDialog)() if (!filePath) return { imported: false } diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx index 1310a9d97..197ea4fd5 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx @@ -293,4 +293,43 @@ describe('SessionNotebookContent export', () => { expect(container.querySelector('[role="alert"]')?.textContent).toBe('Invalid notebook') expect(button?.disabled).toBe(false) }) + + it('shows an import success status and ignores canceled imports', async () => { + const onImport = vi.fn().mockResolvedValueOnce('Imported 2 cells (skipped 1)') + await act(async () => { + root.render( + + ) + }) + + const button = container.querySelector('button[aria-label="Import .ipynb"]') + await act(async () => { + button?.click() + await Promise.resolve() + }) + + expect(onImport).toHaveBeenCalledOnce() + expect(container.querySelector('[role="status"]')?.textContent).toBe( + 'Imported 2 cells (skipped 1)' + ) + expect(container.querySelector('[role="alert"]')).toBeNull() + + onImport.mockResolvedValueOnce(undefined) + await act(async () => { + button?.click() + await Promise.resolve() + }) + + expect(onImport).toHaveBeenCalledTimes(2) + expect(container.querySelector('[role="alert"]')).toBeNull() + expect(container.querySelector('[role="status"]')?.textContent ?? '').toBe('') + }) }) diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx index c3c689f55..8f4ede421 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx @@ -100,6 +100,18 @@ describe('SessionNotebookContent', () => { ) }) + it('renders an imported badge for not-yet-executed imported runs', () => { + const html = renderContent({ + sessionId: 's1', + runs: [makeRun({ status: 'imported', source: 'user' })], + status: 'ready' + }) + + expect(html).toContain('data-testid="session-notebook-imported-badge"') + expect(html).toContain('imported') + expect(html).not.toContain('error (line') + }) + it('enables .ipynb export for a loaded notebook and disables it when empty', () => { const populated = renderContent({ sessionId: 's1', diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx index 5d90cf160..084a96d99 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx @@ -65,6 +65,13 @@ const NotebookDialogCell = ({ ) : ( error ) + ) : run.status === 'imported' ? ( + + imported + ) : null}
{originLabel ? ( @@ -94,7 +101,7 @@ type SessionNotebookContentProps = { onClose: () => void onExport: (kernel: NotebookKernelKind) => Promise onExportAll: () => Promise - onImport: () => Promise + onImport: () => Promise } // Pure presentational body of the dialog: header summary, empty/loading/error/populated states, @@ -115,7 +122,7 @@ const SessionNotebookContent = ({ const [exporting, setExporting] = useState(false) const [exportingAll, setExportingAll] = useState(false) const [exportError, setExportError] = useState() - const [exportSuccess, setExportSuccess] = useState() + const [footerSuccess, setFooterSuccess] = useState() const [importing, setImporting] = useState(false) const [importError, setImportError] = useState() const shortId = sessionId.slice(0, 8) @@ -157,7 +164,8 @@ const SessionNotebookContent = ({ const handleExport = async (): Promise => { setExporting(true) setExportError(undefined) - setExportSuccess(undefined) + setImportError(undefined) + setFooterSuccess(undefined) try { await onExport(effectiveActiveKind) } catch (exportFailure) { @@ -173,10 +181,11 @@ const SessionNotebookContent = ({ const handleExportAll = async (): Promise => { setExportingAll(true) setExportError(undefined) - setExportSuccess(undefined) + setImportError(undefined) + setFooterSuccess(undefined) try { const message = await onExportAll() - if (message) setExportSuccess(message) + if (message) setFooterSuccess(message) } catch (exportFailure) { console.error('Failed to export notebooks by kernel:', exportFailure) setExportError(getErrorMessage(exportFailure)) @@ -188,8 +197,11 @@ const SessionNotebookContent = ({ const handleImport = async (): Promise => { setImporting(true) setImportError(undefined) + setExportError(undefined) + setFooterSuccess(undefined) try { - await onImport() + const message = await onImport() + if (message) setFooterSuccess(message) } catch (importFailure) { setImportError(getErrorMessage(importFailure)) } finally { @@ -252,7 +264,7 @@ const SessionNotebookContent = ({ data-testid={`session-notebook-tab-${kind}`} onClick={() => { setActiveKind(kind) - setExportSuccess(undefined) + setFooterSuccess(undefined) }} className={cn( 'flex items-center gap-1.5 rounded-md px-2 py-1 text-xs transition-colors', @@ -295,7 +307,7 @@ const SessionNotebookContent = ({ )} role={importError || exportError ? 'alert' : 'status'} > - {importError ?? exportError ?? exportSuccess} + {importError ?? exportError ?? footerSuccess}

{originLabel ? ( @@ -527,4 +538,4 @@ const NotebookPreview = ({ item }: NotebookPreviewProps): React.JSX.Element => { ) } -export { NotebookPreview } +export { NotebookPreview, NotebookRunCell } diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx index 8f4ede421..95780ed2b 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx @@ -1,7 +1,10 @@ import { renderToStaticMarkup } from 'react-dom/server' import { describe, expect, it, vi } from 'vitest' -import { SessionNotebookContent } from './SessionNotebookDialog' +import { + formatImportResultMessage, + SessionNotebookContent +} from './SessionNotebookDialog' import type { NotebookRunRecord } from '../../../../shared/notebook' const makeRun = (overrides: Partial = {}): NotebookRunRecord => ({ @@ -109,9 +112,27 @@ describe('SessionNotebookContent', () => { expect(html).toContain('data-testid="session-notebook-imported-badge"') expect(html).toContain('imported') + expect(html).toContain( + 'title="Snapshot from the source .ipynb — re-running appends a new run"' + ) expect(html).not.toContain('error (line') }) + it('formats import result footer messages including env notices', () => { + expect(formatImportResultMessage({ imported: false })).toBeUndefined() + expect( + formatImportResultMessage({ imported: true, cellCount: 2, skippedCellCount: 0 }) + ).toBe('Imported 2 cells') + expect( + formatImportResultMessage({ + imported: true, + cellCount: 1, + skippedCellCount: 1, + environmentNotice: { recorded: ['analysis'], bound: ['default-python'] } + }) + ).toBe('Imported 1 cell (skipped 1) — recorded env "analysis"; re-runs use default-python') + }) + it('enables .ipynb export for a loaded notebook and disables it when empty', () => { const populated = renderContent({ sessionId: 's1', diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx index 084a96d99..623acb1b7 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx @@ -9,7 +9,11 @@ import { cn } from '@/lib/utils' import type { ChatSession } from '@/stores/session-store' import { resolveDataKernelForTab } from '../../../../shared/notebook' -import type { NotebookKernelKind, NotebookRunRecord } from '../../../../shared/notebook' +import type { + ImportNotebookResult, + NotebookKernelKind, + NotebookRunRecord +} from '../../../../shared/notebook' import { NotebookCodeBlock } from './notebook-code' import { NotebookRunOutputs } from './NotebookRunOutputs' import { NotebookInputDataStrip } from './NotebookInputDataStrip' @@ -35,6 +39,23 @@ const getErrorMessage = (error: unknown): string => const pluralize = (count: number, word: string): string => `${count} ${word}${count === 1 ? '' : 's'}` +const IMPORTED_RUN_BADGE_TITLE = + 'Snapshot from the source .ipynb — re-running appends a new run' + +// Composes the Session Notebook footer status for an import result. Cancelled picks return +// undefined; successful imports surface cell counts and any recorded-vs-bound env mismatch. +const formatImportResultMessage = (result: ImportNotebookResult): string | undefined => { + if (!result.imported) return undefined + const summary = `Imported ${pluralize(result.cellCount, 'cell')}` + const withSkipped = + result.skippedCellCount > 0 ? `${summary} (skipped ${result.skippedCellCount})` : summary + const notice = result.environmentNotice + if (!notice || notice.recorded.length === 0) return withSkipped + const recorded = notice.recorded.map((name) => `"${name}"`).join(', ') + const bound = notice.bound.join(', ') + return `${withSkipped} — recorded env ${recorded}; re-runs use ${bound}` +} + // One persisted run rendered as a notebook cell: header badges, code, and split stdout/stderr. The // zero-based index is the cell number shown in [n], aligning the display with a notebook's cells. const NotebookDialogCell = ({ @@ -69,6 +90,7 @@ const NotebookDialogCell = ({ imported @@ -519,10 +541,7 @@ const SessionNotebookDialog = ({ if (!result.imported) return undefined setRuns(await loadSessionNotebookRuns(window.api.notebook, request)) setStatus('ready') - const summary = `Imported ${pluralize(result.cellCount, 'cell')}` - return result.skippedCellCount > 0 - ? `${summary} (skipped ${result.skippedCellCount})` - : summary + return formatImportResultMessage(result) }} /> ) : null} @@ -532,4 +551,9 @@ const SessionNotebookDialog = ({ ) } -export { NotebookDialogCell, SessionNotebookContent, SessionNotebookDialog } +export { + formatImportResultMessage, + NotebookDialogCell, + SessionNotebookContent, + SessionNotebookDialog +} diff --git a/src/shared/notebook.ts b/src/shared/notebook.ts index 2bf50993d..e6e3c2d97 100644 --- a/src/shared/notebook.ts +++ b/src/shared/notebook.ts @@ -514,6 +514,9 @@ export type ImportNotebookResult = imported: true cellCount: number skippedCellCount: number + // Present only when imported python/r cells recorded env names that differ from the + // session's bound env for that language. Renderer composes the footer notice from it. + environmentNotice?: { recorded: string[]; bound: string[] } } // Starts a streamed code write into a notebook cell. From ff9b3e9c731b49c71e4567ef67763562382f03fe Mon Sep 17 00:00:00 2001 From: imjszhang Date: Wed, 22 Jul 2026 16:23:37 +0800 Subject: [PATCH 4/4] feat(notebook): open sessions in JupyterLab Co-authored-by: Cursor --- src/main/notebook/ipc.test.ts | 16 ++ src/main/notebook/ipc.ts | 6 + src/main/notebook/jupyterlab.test.ts | 149 +++++++++++++ src/main/notebook/jupyterlab.ts | 198 ++++++++++++++++++ .../runtime-service.jupyterlab.test.ts | 102 +++++++++ src/main/notebook/runtime-service.ts | 72 ++++++- src/preload/index.d.ts | 2 + src/preload/index.ts | 4 + ...SessionNotebookDialog.interaction.test.tsx | 40 ++++ .../SessionNotebookDialog.render.test.tsx | 1 + .../pages/workspace/SessionNotebookDialog.tsx | 55 ++++- src/renderer/web/api-map.generated.ts | 1 + src/shared/notebook.ts | 6 + 13 files changed, 644 insertions(+), 8 deletions(-) create mode 100644 src/main/notebook/jupyterlab.test.ts create mode 100644 src/main/notebook/jupyterlab.ts create mode 100644 src/main/notebook/runtime-service.jupyterlab.test.ts diff --git a/src/main/notebook/ipc.test.ts b/src/main/notebook/ipc.test.ts index 26b8c456a..d0f2c7d1e 100644 --- a/src/main/notebook/ipc.test.ts +++ b/src/main/notebook/ipc.test.ts @@ -52,6 +52,9 @@ describe('notebook IPC handlers', () => { runCell: vi.fn().mockResolvedValue({ runId: 'run-2', status: 'completed' }), exportIpynb: vi.fn().mockResolvedValue({ saved: true, filePath: '/tmp/session.ipynb' }), importIpynb: vi.fn().mockResolvedValue({ imported: true, cellCount: 1 }), + openInJupyterLab: vi + .fn() + .mockResolvedValue({ opened: true, url: 'http://localhost:8888', alreadyRunning: false }), beginCodeCell: vi.fn().mockResolvedValue({ cellId: 'cell-1', writeId: 'write-1' }), appendCodeCell: vi.fn().mockResolvedValue({ receivedBytes: 5 }), finishCodeCell: vi.fn().mockResolvedValue({ status: 'idle' }), @@ -96,6 +99,7 @@ describe('notebook IPC handlers', () => { kernel: 'python' }) await handlers.importIpynb({ sessionId: 'session-1', workspaceCwd: '/workspace' }) + await handlers.openInJupyterLab({ sessionId: 'session-1', workspaceCwd: '/workspace' }) expect(service.execute).toHaveBeenCalledWith({ sessionId: 'session-1', @@ -126,6 +130,10 @@ describe('notebook IPC handlers', () => { sessionId: 'session-1', workspaceCwd: '/workspace' }) + expect(service.openInJupyterLab).toHaveBeenCalledWith({ + sessionId: 'session-1', + workspaceCwd: '/workspace' + }) }) it('registers every notebook channel and forwards the renderer payload unchanged', async () => { @@ -139,6 +147,11 @@ describe('notebook IPC handlers', () => { execute: vi.fn().mockResolvedValue({ runId: 'run-2', status: 'completed' }), exportIpynb: vi.fn().mockResolvedValue({ saved: false }), importIpynb: vi.fn().mockResolvedValue({ imported: false }), + openInJupyterLab: vi.fn().mockResolvedValue({ + opened: true, + url: 'http://localhost:8888', + alreadyRunning: false + }), restart: vi.fn().mockResolvedValue({ sessionId: 'session-1' }), shutdown: vi.fn().mockResolvedValue({ sessionId: 'session-1', status: 'shutdown' }) } as unknown as NotebookRuntimeService @@ -155,6 +168,7 @@ describe('notebook IPC handlers', () => { 'notebook:export-ipynb', 'notebook:export-ipynb-all', 'notebook:import-ipynb', + 'notebook:open-jupyterlab', 'notebook:restart', 'notebook:shutdown' ]) @@ -187,6 +201,7 @@ describe('notebook IPC handlers', () => { await ipcHandlers.get('notebook:execute')?.(undefined, execute) await ipcHandlers.get('notebook:export-ipynb')?.(undefined, session) await ipcHandlers.get('notebook:import-ipynb')?.(undefined, session) + await ipcHandlers.get('notebook:open-jupyterlab')?.(undefined, session) await ipcHandlers.get('notebook:restart')?.(undefined, session) await ipcHandlers.get('notebook:shutdown')?.(undefined, session) @@ -199,6 +214,7 @@ describe('notebook IPC handlers', () => { expect(service.execute).toHaveBeenCalledWith(publicExecute) expect(service.exportIpynb).toHaveBeenCalledWith(session) expect(service.importIpynb).toHaveBeenCalledWith(session) + expect(service.openInJupyterLab).toHaveBeenCalledWith(session) expect(service.restart).toHaveBeenCalledWith(session) expect(service.shutdown).toHaveBeenCalledWith(session) }) diff --git a/src/main/notebook/ipc.ts b/src/main/notebook/ipc.ts index c5c74fdd0..7a7d111ac 100644 --- a/src/main/notebook/ipc.ts +++ b/src/main/notebook/ipc.ts @@ -11,6 +11,7 @@ import type { FinishNotebookCodeCellRequest, ImportNotebookResult, NotebookRunSummary, + OpenJupyterLabResult, NotebookSessionReference, NotebookSessionRequest, NotebookSessionState, @@ -36,6 +37,7 @@ type NotebookHandlers = { exportIpynb: (request: ExportNotebookKernelRequest) => Promise exportIpynbAll: (request: ExportNotebookAllRequest) => Promise importIpynb: (request: NotebookSessionRequest) => Promise + openInJupyterLab: (request: NotebookSessionRequest) => Promise restart: (request: NotebookSessionRequest) => Promise shutdown: (request: NotebookSessionRequest) => ReturnType } @@ -64,6 +66,7 @@ const createNotebookHandlers = (service: NotebookRuntimeService): NotebookHandle exportIpynb: (request) => service.exportIpynb(request), exportIpynbAll: (request) => service.exportIpynbAll(request), importIpynb: (request) => withDataRootWrite(() => service.importIpynb(request)), + openInJupyterLab: (request) => service.openInJupyterLab(request), restart: (request) => withDataRootWrite(() => service.restart(request)), shutdown: (request) => withDataRootWrite(() => service.shutdown(request)) }) @@ -102,6 +105,9 @@ const registerNotebookIpcHandlers = (service: NotebookRuntimeService): void => { ipcMain.handle('notebook:import-ipynb', (_event, request: NotebookSessionRequest) => handlers.importIpynb(request) ) + ipcMain.handle('notebook:open-jupyterlab', (_event, request: NotebookSessionRequest) => + handlers.openInJupyterLab(request) + ) ipcMain.handle('notebook:restart', (_event, request: NotebookSessionRequest) => handlers.restart(request) ) diff --git a/src/main/notebook/jupyterlab.test.ts b/src/main/notebook/jupyterlab.test.ts new file mode 100644 index 000000000..484514aad --- /dev/null +++ b/src/main/notebook/jupyterlab.test.ts @@ -0,0 +1,149 @@ +import { EventEmitter } from 'node:events' +import type { ChildProcess } from 'node:child_process' +import { PassThrough } from 'node:stream' + +import { describe, expect, it, vi } from 'vitest' + +vi.mock('electron', () => ({ shell: { openExternal: vi.fn() } })) + +import { JupyterLabManager, type SpawnProcess } from './jupyterlab' + +type FakeChild = ChildProcess & { + stdout: PassThrough + stderr: PassThrough +} + +const fakeChild = (): FakeChild => { + const child = new EventEmitter() as FakeChild + child.stdout = new PassThrough() + child.stderr = new PassThrough() + Object.defineProperties(child, { + exitCode: { value: null, writable: true }, + signalCode: { value: null, writable: true }, + killed: { value: false, writable: true } + }) + child.kill = vi.fn(() => true) + return child +} + +const exit = (child: FakeChild, code: number): void => { + Object.defineProperty(child, 'exitCode', { value: code, writable: true }) + child.emit('exit', code, null) +} + +describe('JupyterLabManager', () => { + it('installs when the probe fails, launches, and opens the reported URL', async () => { + let installed = false + let launched: FakeChild | undefined + const spawnProcess = vi.fn((_command, args) => { + const child = fakeChild() + if (args.includes('--version')) { + queueMicrotask(() => { + exit(child, installed ? 0 : 1) + }) + } else { + launched = child + queueMicrotask(() => child.stderr.write('http://127.0.0.1:4321/lab?token=secret\n')) + } + return child + }) + const ensureInstalled = vi.fn(async () => { + installed = true + }) + const openExternal = vi.fn().mockResolvedValue(undefined) + const manager = new JupyterLabManager({ spawnProcess, openExternal }) + + const result = await manager.launch({ + sessionId: 'session-1', + command: '/env/bin/python', + notebookPath: '/session/data/session.ipynb', + rootDir: '/session/data', + cwd: '/session/data', + ensureInstalled + }) + + expect(result).toEqual({ + url: 'http://127.0.0.1:4321/lab?token=secret', + alreadyRunning: false + }) + expect(ensureInstalled).toHaveBeenCalledOnce() + expect(openExternal).toHaveBeenCalledWith(result.url) + expect(spawnProcess).toHaveBeenLastCalledWith( + '/env/bin/python', + expect.arrayContaining([ + '-m', + 'jupyterlab', + '/session/data/session.ipynb', + '--no-browser', + '--ServerApp.port=0', + '--ServerApp.root_dir=/session/data' + ]), + expect.objectContaining({ cwd: '/session/data' }) + ) + expect(launched).toBeDefined() + }) + + it('reopens an already-running session without spawning another process', async () => { + const launchChild = fakeChild() + const spawnProcess = vi.fn((_command, args) => { + if (args.includes('--version')) { + const probe = fakeChild() + queueMicrotask(() => { + exit(probe, 0) + }) + return probe + } + queueMicrotask(() => launchChild.stdout.write('http://localhost:9999/lab?token=t\n')) + return launchChild + }) + const openExternal = vi.fn().mockResolvedValue(undefined) + const manager = new JupyterLabManager({ spawnProcess, openExternal }) + const request = { + sessionId: 'session-1', + command: 'python', + notebookPath: '/data/session.ipynb', + rootDir: '/data', + cwd: '/data', + ensureInstalled: vi.fn() + } + + await manager.launch(request) + const second = await manager.launch(request) + + expect(second.alreadyRunning).toBe(true) + expect(spawnProcess).toHaveBeenCalledTimes(2) + expect(openExternal).toHaveBeenCalledTimes(2) + }) + + it('terminates tracked process trees during shutdown', async () => { + const launchChild = fakeChild() + const spawnProcess = vi.fn((_command, args) => { + if (args.includes('--version')) { + const probe = fakeChild() + queueMicrotask(() => { + exit(probe, 0) + }) + return probe + } + queueMicrotask(() => launchChild.stdout.write('http://localhost:9999/lab?token=t\n')) + return launchChild + }) + const terminate = vi.fn().mockResolvedValue({ reaped: true }) + const manager = new JupyterLabManager({ + spawnProcess, + openExternal: vi.fn().mockResolvedValue(undefined), + terminate + }) + await manager.launch({ + sessionId: 'session-1', + command: 'python', + notebookPath: '/data/session.ipynb', + rootDir: '/data', + cwd: '/data', + ensureInstalled: vi.fn() + }) + + await expect(manager.shutdownAll()).resolves.toEqual({ reaped: true }) + expect(terminate).toHaveBeenCalledWith(launchChild) + }) +}) diff --git a/src/main/notebook/jupyterlab.ts b/src/main/notebook/jupyterlab.ts new file mode 100644 index 000000000..16cf3cc69 --- /dev/null +++ b/src/main/notebook/jupyterlab.ts @@ -0,0 +1,198 @@ +import { spawn, type ChildProcess, type SpawnOptions } from 'node:child_process' + +import { shell } from 'electron' + +import { terminateProcessTree, type ProcessTreeKillResult } from '../process-tree' + +type JupyterLabLaunchRequest = { + sessionId: string + command: string + commandArgs?: string[] + notebookPath: string + rootDir: string + cwd: string + ensureInstalled: () => Promise +} + +type JupyterLabLaunchResult = { + url: string + alreadyRunning: boolean +} + +type SpawnProcess = (command: string, args: string[], options: SpawnOptions) => ChildProcess + +type JupyterLabManagerDeps = { + spawnProcess?: SpawnProcess + openExternal?: (url: string) => Promise + terminate?: (child: ChildProcess) => Promise + startupTimeoutMs?: number +} + +type RunningJupyterLab = { + child: ChildProcess + url?: string +} + +const JUPYTER_URL_PATTERN = /https?:\/\/(?:127\.0\.0\.1|localhost):\d+\/[^\s]*/i + +const probeJupyterLab = async ( + spawnProcess: SpawnProcess, + command: string, + commandArgs: string[], + cwd: string, + timeoutMs: number +): Promise => { + let child: ChildProcess + try { + child = spawnProcess(command, [...commandArgs, '-m', 'jupyterlab', '--version'], { + cwd, + windowsHide: true, + stdio: 'ignore' + }) + } catch { + return false + } + return new Promise((resolve) => { + let settled = false + const finish = (available: boolean): void => { + if (settled) return + settled = true + clearTimeout(timeout) + resolve(available) + } + const timeout = setTimeout(() => { + child.kill() + finish(false) + }, timeoutMs) + timeout.unref?.() + child.once('error', () => finish(false)) + child.once('exit', (code) => finish(code === 0)) + }) +} + +class JupyterLabManager { + private readonly running = new Map() + private readonly spawnProcess: SpawnProcess + private readonly openExternal: (url: string) => Promise + private readonly terminate: (child: ChildProcess) => Promise + private readonly startupTimeoutMs: number + + constructor(deps: JupyterLabManagerDeps = {}) { + this.spawnProcess = deps.spawnProcess ?? spawn + this.openExternal = deps.openExternal ?? ((url) => shell.openExternal(url)) + this.terminate = deps.terminate ?? ((child) => terminateProcessTree(child)) + this.startupTimeoutMs = deps.startupTimeoutMs ?? 30_000 + } + + async launch(request: JupyterLabLaunchRequest): Promise { + const existing = this.running.get(request.sessionId) + if (existing?.url && existing.child.exitCode === null) { + await this.openExternal(existing.url) + return { url: existing.url, alreadyRunning: true } + } + + const commandArgs = request.commandArgs ?? [] + const probeTimeoutMs = Math.min(this.startupTimeoutMs, 10_000) + if ( + !(await probeJupyterLab( + this.spawnProcess, + request.command, + commandArgs, + request.cwd, + probeTimeoutMs + )) + ) { + await request.ensureInstalled() + if ( + !(await probeJupyterLab( + this.spawnProcess, + request.command, + commandArgs, + request.cwd, + probeTimeoutMs + )) + ) { + throw new Error('JupyterLab installation completed but the module is still unavailable.') + } + } + + const child = this.spawnProcess( + request.command, + [ + ...commandArgs, + '-m', + 'jupyterlab', + request.notebookPath, + '--no-browser', + '--ServerApp.port=0', + `--ServerApp.root_dir=${request.rootDir}` + ], + { + cwd: request.cwd, + windowsHide: true, + stdio: ['ignore', 'pipe', 'pipe'] + } + ) + const running: RunningJupyterLab = { child } + this.running.set(request.sessionId, running) + + child.once('exit', () => { + if (this.running.get(request.sessionId)?.child === child) { + this.running.delete(request.sessionId) + } + }) + + const url = await new Promise((resolve, reject) => { + let settled = false + let output = '' + const finish = (error: Error | null, value?: string): void => { + if (settled) return + settled = true + clearTimeout(timeout) + if (error) reject(error) + else resolve(value as string) + } + const inspect = (chunk: Buffer | string): void => { + output = `${output}${String(chunk)}`.slice(-16_384) + const match = output.match(JUPYTER_URL_PATTERN) + if (match) finish(null, match[0]) + } + const timeout = setTimeout( + () => finish(new Error('Timed out waiting for JupyterLab to start.')), + this.startupTimeoutMs + ) + timeout.unref?.() + child.stdout?.on('data', inspect) + child.stderr?.on('data', inspect) + child.once('error', (error) => finish(error)) + child.once('exit', (code) => { + finish(new Error(`JupyterLab exited before startup (code ${String(code)}).`)) + }) + }).catch(async (error: unknown) => { + this.running.delete(request.sessionId) + await this.terminate(child) + throw error + }) + + running.url = url + await this.openExternal(url) + return { url, alreadyRunning: false } + } + + async shutdown(sessionId: string): Promise { + const running = this.running.get(sessionId) + if (!running) return { reaped: true } + this.running.delete(sessionId) + return this.terminate(running.child) + } + + async shutdownAll(): Promise { + const children = Array.from(this.running.values(), ({ child }) => child) + this.running.clear() + const results = await Promise.all(children.map((child) => this.terminate(child))) + return { reaped: results.every((result) => result.reaped) } + } +} + +export { JupyterLabManager } +export type { JupyterLabLaunchRequest, JupyterLabLaunchResult, JupyterLabManagerDeps, SpawnProcess } diff --git a/src/main/notebook/runtime-service.jupyterlab.test.ts b/src/main/notebook/runtime-service.jupyterlab.test.ts new file mode 100644 index 000000000..1707beb2c --- /dev/null +++ b/src/main/notebook/runtime-service.jupyterlab.test.ts @@ -0,0 +1,102 @@ +import { mkdtemp, readFile, rm } from 'node:fs/promises' +import { tmpdir } from 'node:os' +import { join } from 'node:path' + +import { afterEach, describe, expect, it, vi } from 'vitest' + +import type { NotebookRunDocument } from '../../shared/notebook' +import type { JupyterLabManager, JupyterLabLaunchRequest } from './jupyterlab' +import type { NotebookRunRepository } from './repository' +import { NotebookRuntimeService } from './runtime-service' +import { envPrefix, pythonBin, runtimeRoot } from './runtime-paths' + +const roots: string[] = [] + +afterEach(async () => { + await Promise.all(roots.splice(0).map((root) => rm(root, { recursive: true, force: true }))) +}) + +describe('NotebookRuntimeService openInJupyterLab', () => { + it('writes the current projection and launches with the bound managed Python', async () => { + const root = await mkdtemp(join(tmpdir(), 'open-science-jupyterlab-')) + roots.push(root) + const document: NotebookRunDocument = { + version: 1, + projectName: 'default-project', + sessionId: '12345678-abcd', + workspaceCwd: '/workspace', + notebookSessionRoot: join(root, 'notebooks', 'default-project', '12345678-abcd'), + dataRoot: join(root, 'notebooks', 'default-project', '12345678-abcd', 'data'), + kernel: { + language: 'python', + kernelName: 'python3', + runtimeRoot: runtimeRoot(root), + lastKnownStatus: 'idle' + }, + runs: [ + { + runId: 'run-1', + cellId: 'cell-1', + source: 'agent', + kernelKind: 'python', + script: 'print(1)', + status: 'completed', + startedAt: 1, + text: { stdout: '', stderr: '', traceback: '', plain: [] }, + outputs: [], + artifacts: [], + workingFiles: [] + } + ], + updatedAt: 2 + } + const repository = { + loadOrCreate: vi.fn().mockResolvedValue(document), + findExisting: vi.fn().mockResolvedValue(document) + } as unknown as NotebookRunRepository + let launchRequest: JupyterLabLaunchRequest | undefined + const manager = { + launch: vi.fn(async (request: JupyterLabLaunchRequest) => { + launchRequest = request + return { url: 'http://localhost:8888/lab?token=x', alreadyRunning: false } + }), + shutdown: vi.fn().mockResolvedValue({ reaped: true }), + shutdownAll: vi.fn().mockResolvedValue({ reaped: true }) + } as unknown as JupyterLabManager + const service = new NotebookRuntimeService({ + configRoot: join(root, 'config'), + dataRoot: root, + projectName: 'default-project', + repository, + executorFactory: () => ({ + execute: vi.fn(), + shutdown: vi.fn().mockResolvedValue({ reaped: true }) + }), + jupyterLabManager: manager + }) + + const result = await service.openInJupyterLab({ + sessionId: '12345678-abcd', + workspaceCwd: '/workspace' + }) + + expect(result).toEqual({ + opened: true, + url: 'http://localhost:8888/lab?token=x', + alreadyRunning: false + }) + expect(launchRequest).toMatchObject({ + sessionId: '12345678-abcd', + command: pythonBin(envPrefix(runtimeRoot(root), 'default-python')), + rootDir: document.dataRoot, + cwd: document.dataRoot + }) + const notebookPath = launchRequest?.notebookPath + expect(notebookPath).toBe(join(document.dataRoot, 'session-12345678.ipynb')) + const written = JSON.parse(await readFile(notebookPath as string, 'utf8')) as { + nbformat: number + cells: Array<{ source: string[] }> + } + expect(written).toMatchObject({ nbformat: 4, cells: [{ source: ['print(1)'] }] }) + }) +}) diff --git a/src/main/notebook/runtime-service.ts b/src/main/notebook/runtime-service.ts index 0c6fd4bb6..2243f3157 100644 --- a/src/main/notebook/runtime-service.ts +++ b/src/main/notebook/runtime-service.ts @@ -1,7 +1,7 @@ import { spawn, type ChildProcess } from 'node:child_process' import { randomUUID } from 'node:crypto' import { existsSync, realpathSync } from 'node:fs' -import { readFile, realpath, rm, writeFile } from 'node:fs/promises' +import { mkdir, readFile, realpath, rm, writeFile } from 'node:fs/promises' import { isAbsolute, join, relative, resolve, sep, win32 } from 'node:path' import type { @@ -25,6 +25,7 @@ import type { NotebookLanguage, NotebookOutput, NotebookRunDocument, + OpenJupyterLabResult, NotebookRunRecord, NotebookRunSource, NotebookRunStatus, @@ -45,12 +46,14 @@ import type { } from '../../shared/notebook-env' import type { PackageMirror } from '../../shared/mirror' import { + runDocumentToIpynb, runDocumentToIpynbByKernel, runDocumentToIpynbForKernel, type NbformatOutput, type ResolvedArtifact } from './ipynb-export' import { ipynbToRunRecords } from './ipynb-import' +import { JupyterLabManager } from './jupyterlab' import { NotebookKernelExecutor, type NotebookKernelExecutorOptions } from './kernel-executor' import { saveIpynbAll } from './save-ipynb-all' import type { KernelProcessKind } from './kernel-executor' @@ -408,6 +411,7 @@ type NotebookRuntimeServiceOptions = { > // Native file-picker seam for notebook import tests. pickIpynb?: () => Promise + jupyterLabManager?: JupyterLabManager } // The wire binding plus the interpreter override the executor needs. `resolvedInterpreter` is set only @@ -1063,6 +1067,7 @@ class EnvConcurrencyLock { // Coordinates notebook cells, shared interpreters, persisted run history, and UI notifications. class NotebookRuntimeService { private readonly repository: NotebookRunRepository + private readonly jupyterLabManager: JupyterLabManager private readonly sessions = new Map() private readonly announcedAgentSessionIds = new Set() // Serializes environment management (installs) against kernel runs on the same language's env; @@ -1157,6 +1162,7 @@ class NotebookRuntimeService { constructor(private readonly options: NotebookRuntimeServiceOptions) { this.repository = options.repository ?? new NotebookRunRepository(options.dataRoot) + this.jupyterLabManager = options.jupyterLabManager ?? new JupyterLabManager() this.mcpRpcConnectionResolver = options.getMcpRpcConnection this.packageMirrorResolver = options.getPackageMirror this.runtimeEnablementResolver = options.getRuntimeEnablement @@ -2582,6 +2588,66 @@ class NotebookRuntimeService { } } + // Materializes the current projection inside the session data root, ensures JupyterLab is available + // in the session's bound Python runtime, and opens the authenticated local URL it reports. + async openInJupyterLab(request: NotebookSessionRequest): Promise { + const session = await this.ensureSession(request) + const document = await this.repository.findExisting(session.projectName, session.sessionId) + if (!document) { + throw new Error(`Notebook session not found: ${request.sessionId}`) + } + + const binding = session.runtimeBindings.get('python') + if (binding?.status === 'unavailable') { + throw new Error('The bound Python runtime is unavailable.') + } + const environment = this.resolveRunEnv(session, 'python') + const interpreter = + binding?.source === 'external' + ? binding.resolvedInterpreter + : { + command: pythonBin(envPrefix(getRuntimeRoot(this.options.dataRoot), environment)) + } + if (!interpreter) { + throw new Error('The bound Python interpreter could not be resolved.') + } + + const artifactOutputs = await resolveNotebookArtifactOutputs( + document, + this.options.resolveArtifactPath + ) + const notebook = runDocumentToIpynb(document, { + appVersion: this.options.appVersion, + artifactOutputs + }) + await mkdir(document.dataRoot, { recursive: true }) + const notebookPath = join(document.dataRoot, `session-${request.sessionId.slice(0, 8)}.ipynb`) + await writeFile(notebookPath, `${JSON.stringify(notebook, null, 2)}\n`, 'utf8') + + const launched = await this.jupyterLabManager.launch({ + sessionId: session.sessionId, + command: interpreter.command, + commandArgs: interpreter.args, + notebookPath, + rootDir: document.dataRoot, + cwd: document.dataRoot, + ensureInstalled: async () => { + const result = await this.managePackages({ + language: 'python', + packages: ['jupyterlab'], + sessionId: session.sessionId, + workspaceCwd: request.workspaceCwd, + projectName: session.projectName + }) + if (!result.ok) { + throw new Error(result.error || result.log || 'Failed to install JupyterLab.') + } + } + }) + + return { opened: true, ...launched } + } + // Replaces the interpreter process while preserving cells and durable run history. Prefers the // executor's own in-place restart (keeps the same instance, e.g. NotebookKernelExecutor tears down // and lazily respawns its loops) and only shuts down + recreates for executors that don't support it. @@ -3389,6 +3455,7 @@ class NotebookRuntimeService { const session = this.sessions.get(request.sessionId) if (session) { + await this.jupyterLabManager.shutdown(request.sessionId) await session.executor.shutdown() this.sessions.delete(request.sessionId) } @@ -3741,8 +3808,9 @@ class NotebookRuntimeService { const results = await Promise.all( Array.from(this.sessions.values()).map((session) => session.executor.shutdown()) ) + const jupyterResult = await this.jupyterLabManager.shutdownAll() this.sessions.clear() - return { reaped: results.every((result) => result.reaped) } + return { reaped: jupyterResult.reaped && results.every((result) => result.reaped) } } // Lists sessions with a cell mid-execution, for the pre-migration active-session warning. diff --git a/src/preload/index.d.ts b/src/preload/index.d.ts index b0c7dba86..597e60062 100644 --- a/src/preload/index.d.ts +++ b/src/preload/index.d.ts @@ -74,6 +74,7 @@ import type { ImportNotebookResult, NotebookLanguage, NotebookRunSummary, + OpenJupyterLabResult, NotebookSessionReference, NotebookSessionRequest, NotebookSessionState, @@ -562,6 +563,7 @@ interface OpenScienceAPI { exportIpynb(request: ExportNotebookKernelRequest): Promise exportIpynbAll(request: ExportNotebookAllRequest): Promise importIpynb(request: NotebookSessionRequest): Promise + openInJupyterLab(request: NotebookSessionRequest): Promise restart(request: NotebookSessionRequest): Promise shutdown(request: NotebookSessionRequest): Promise<{ sessionId: string; status: 'shutdown' }> onAvailable(listener: AcpListener): RemoveListener diff --git a/src/preload/index.ts b/src/preload/index.ts index af6d08c2a..22a41b04c 100644 --- a/src/preload/index.ts +++ b/src/preload/index.ts @@ -76,6 +76,7 @@ import type { ImportNotebookResult, NotebookLanguage, NotebookRunSummary, + OpenJupyterLabResult, NotebookSessionReference, NotebookSessionRequest, NotebookSessionState, @@ -621,6 +622,7 @@ type OpenScienceAPI = { exportIpynb: (request: ExportNotebookKernelRequest) => Promise exportIpynbAll: (request: ExportNotebookAllRequest) => Promise importIpynb: (request: NotebookSessionRequest) => Promise + openInJupyterLab: (request: NotebookSessionRequest) => Promise restart: (request: NotebookSessionRequest) => Promise shutdown: ( request: NotebookSessionRequest @@ -1258,6 +1260,8 @@ const api: OpenScienceAPI = { ipcRenderer.invoke('notebook:export-ipynb-all', request) as Promise, importIpynb: (request) => ipcRenderer.invoke('notebook:import-ipynb', request) as Promise, + openInJupyterLab: (request) => + ipcRenderer.invoke('notebook:open-jupyterlab', request) as Promise, restart: (request) => ipcRenderer.invoke('notebook:restart', request) as Promise, shutdown: (request) => diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx index 197ea4fd5..43e446d12 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.interaction.test.tsx @@ -52,6 +52,7 @@ describe('SessionNotebookContent export', () => { onExport={onExport} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -87,6 +88,7 @@ describe('SessionNotebookContent export', () => { onExport={onExport} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -124,6 +126,7 @@ describe('SessionNotebookContent export', () => { onExport={onExport} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -153,6 +156,7 @@ describe('SessionNotebookContent export', () => { onExport={onExport} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -180,6 +184,7 @@ describe('SessionNotebookContent export', () => { onExport={failingExport} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -204,6 +209,7 @@ describe('SessionNotebookContent export', () => { onExport={vi.fn()} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -227,6 +233,7 @@ describe('SessionNotebookContent export', () => { onExport={vi.fn()} onExportAll={onExportAll} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -255,6 +262,7 @@ describe('SessionNotebookContent export', () => { onExport={vi.fn()} onExportAll={onExportAll} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -278,6 +286,7 @@ describe('SessionNotebookContent export', () => { onExport={vi.fn()} onExportAll={vi.fn()} onImport={onImport} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -306,6 +315,7 @@ describe('SessionNotebookContent export', () => { onExport={vi.fn()} onExportAll={vi.fn()} onImport={onImport} + onOpenJupyterLab={vi.fn()} /> ) }) @@ -332,4 +342,34 @@ describe('SessionNotebookContent export', () => { expect(container.querySelector('[role="alert"]')).toBeNull() expect(container.querySelector('[role="status"]')?.textContent ?? '').toBe('') }) + + it('opens JupyterLab and surfaces launcher failures', async () => { + const onOpenJupyterLab = vi.fn().mockRejectedValue(new Error('Install denied')) + await act(async () => { + root.render( + + ) + }) + + const button = container.querySelector( + 'button[aria-label="Open in JupyterLab"]' + ) + expect(button?.disabled).toBe(false) + await act(async () => { + button?.click() + await Promise.resolve() + }) + + expect(onOpenJupyterLab).toHaveBeenCalledOnce() + expect(container.querySelector('[role="alert"]')?.textContent).toBe('Install denied') + }) }) diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx index 95780ed2b..193fc417e 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.render.test.tsx @@ -36,6 +36,7 @@ const renderContent = (props: { onExport={vi.fn()} onExportAll={vi.fn()} onImport={vi.fn()} + onOpenJupyterLab={vi.fn()} {...props} /> ) diff --git a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx index 623acb1b7..a776892cb 100644 --- a/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx +++ b/src/renderer/src/pages/workspace/SessionNotebookDialog.tsx @@ -1,5 +1,5 @@ import { useEffect, useState } from 'react' -import { Download, LoaderCircle, Upload, X } from 'lucide-react' +import { Download, ExternalLink, LoaderCircle, Upload, X } from 'lucide-react' import { Dialog } from 'radix-ui' import { dialogOverlayClassName, dialogPanelClassName } from '@/components/ui/dialog-chrome' @@ -124,6 +124,7 @@ type SessionNotebookContentProps = { onExport: (kernel: NotebookKernelKind) => Promise onExportAll: () => Promise onImport: () => Promise + onOpenJupyterLab: () => Promise } // Pure presentational body of the dialog: header summary, empty/loading/error/populated states, @@ -138,7 +139,8 @@ const SessionNotebookContent = ({ onClose, onExport, onExportAll, - onImport + onImport, + onOpenJupyterLab }: SessionNotebookContentProps): React.JSX.Element => { const [activeKind, setActiveKind] = useState('python') const [exporting, setExporting] = useState(false) @@ -147,6 +149,8 @@ const SessionNotebookContent = ({ const [footerSuccess, setFooterSuccess] = useState() const [importing, setImporting] = useState(false) const [importError, setImportError] = useState() + const [openingJupyterLab, setOpeningJupyterLab] = useState(false) + const [jupyterLabError, setJupyterLabError] = useState() const shortId = sessionId.slice(0, 8) const agents = runs.some((run) => run.source === 'agent') ? 1 : 0 // Only python/r runs are "cells" in the notebook sense; repl/bash are control-plane/shell runs @@ -170,7 +174,7 @@ const SessionNotebookContent = ({ ? activeKind : (KERNEL_KIND_ORDER.find((kind) => kindsWithRuns.has(kind)) ?? visibleKinds[0] ?? 'python') const visibleRuns = runs.filter((run) => resolveRunKernelKind(run) === effectiveActiveKind) - const busy = exporting || exportingAll || importing + const busy = exporting || exportingAll || importing || openingJupyterLab const exportDisabled = status !== 'ready' || runs.length === 0 || busy const importDisabled = status !== 'ready' || busy @@ -187,6 +191,7 @@ const SessionNotebookContent = ({ setExporting(true) setExportError(undefined) setImportError(undefined) + setJupyterLabError(undefined) setFooterSuccess(undefined) try { await onExport(effectiveActiveKind) @@ -204,6 +209,7 @@ const SessionNotebookContent = ({ setExportingAll(true) setExportError(undefined) setImportError(undefined) + setJupyterLabError(undefined) setFooterSuccess(undefined) try { const message = await onExportAll() @@ -220,6 +226,7 @@ const SessionNotebookContent = ({ setImporting(true) setImportError(undefined) setExportError(undefined) + setJupyterLabError(undefined) setFooterSuccess(undefined) try { const message = await onImport() @@ -231,6 +238,21 @@ const SessionNotebookContent = ({ } } + const handleOpenJupyterLab = async (): Promise => { + setOpeningJupyterLab(true) + setJupyterLabError(undefined) + setImportError(undefined) + setExportError(undefined) + setFooterSuccess(undefined) + try { + await onOpenJupyterLab() + } catch (openFailure) { + setJupyterLabError(getErrorMessage(openFailure)) + } finally { + setOpeningJupyterLab(false) + } + } + // The "Download all" path is only useful when there's more than one data kernel to write; a // single-kernel session's secondary button would just duplicate the main button. The data-kernel // count comes from `kindsWithRuns` (control-plane kinds don't generate their own .ipynb). @@ -323,15 +345,29 @@ const SessionNotebookContent = ({

- {importError ?? exportError ?? footerSuccess} + {jupyterLabError ?? importError ?? exportError ?? footerSuccess}

+