Skip to content

Commit

Permalink
Move the hasRunningNotebookSession context out of the notebook sess…
Browse files Browse the repository at this point in the history
…ion service (#5688)

Another step toward removing the notebook session service entirely
(#2671).

### QA Notes

This shouldn't change any behavior, although I did fix some edge cases
where the context was not set correctly.

Opening a notebook, restarting it, closing and reopening it should all
work. The restart button should be deactivated when the session is
exited.
  • Loading branch information
seeM authored Dec 11, 2024
1 parent 55ddd14 commit 70ed8a3
Show file tree
Hide file tree
Showing 11 changed files with 396 additions and 122 deletions.
2 changes: 1 addition & 1 deletion .vscode-test.js
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,7 @@ const extensions = [
{
label: 'positron-notebook-controllers',
workspaceFolder: 'extensions/positron-notebook-controllers/test-workspace',
mocha: { timeout: 60_000 }
mocha: { timeout: 5_000 }
},
{
label: 'positron-supervisor',
Expand Down
13 changes: 12 additions & 1 deletion extensions/positron-notebook-controllers/src/commands.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@

import * as vscode from 'vscode';
import { NotebookSessionService } from './notebookSessionService';
import { getNotebookSession } from './utils';
import { getNotebookSession, isActiveNotebookEditorUri } from './utils';
import { setHasRunningNotebookSessionContext } from './extension';

export function registerCommands(context: vscode.ExtensionContext, notebookSessionService: NotebookSessionService): void {
context.subscriptions.push(vscode.commands.registerCommand('positron.restartKernel', async () => {
Expand All @@ -21,12 +22,22 @@ export function registerCommands(context: vscode.ExtensionContext, notebookSessi
throw new Error('No session found for active notebook. This command should only be available when a session is running.');
}

// Disable the hasRunningNotebookSession context before restarting.
if (isActiveNotebookEditorUri(notebook.uri)) {
await setHasRunningNotebookSessionContext(false);
}

// Restart the session with a progress bar.
try {
await vscode.window.withProgress({
location: vscode.ProgressLocation.Notification,
title: vscode.l10n.t("Restarting {0} interpreter for '{1}'", session.runtimeMetadata.runtimeName, notebook.uri.path),
}, () => notebookSessionService.restartRuntimeSession(notebook.uri));

// Enable the hasRunningNotebookSession context.
if (isActiveNotebookEditorUri(notebook.uri)) {
await setHasRunningNotebookSessionContext(true);
}
} catch (error) {
vscode.window.showErrorMessage(
vscode.l10n.t("Restarting {0} interpreter for '{1}' failed. Reason: {2}",
Expand Down
45 changes: 32 additions & 13 deletions extensions/positron-notebook-controllers/src/extension.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,17 +10,29 @@ import { NotebookSessionService } from './notebookSessionService';
import { registerCommands } from './commands';
import { JUPYTER_NOTEBOOK_TYPE } from './constants';
import { registerExecutionInfoStatusBar } from './statusBar';
import { getNotebookSession } from './utils';
import { getNotebookSession, isActiveNotebookEditorUri } from './utils';

export const log = vscode.window.createOutputChannel('Positron Notebook Controllers', { log: true });

const _onDidSetHasRunningNotebookSessionContext = new vscode.EventEmitter<boolean>();

/**
* An event that fires when the hasRunningNotebookSessionContext is set.
* Currently only for testing purposes.
*/
export const onDidSetHasRunningNotebookSessionContext = _onDidSetHasRunningNotebookSessionContext.event;

export async function activate(context: vscode.ExtensionContext): Promise<void> {
context.subscriptions.push(_onDidSetHasRunningNotebookSessionContext);

const notebookSessionService = new NotebookSessionService();
context.subscriptions.push(notebookSessionService);

// Shutdown any running sessions when a notebook is closed.
context.subscriptions.push(vscode.workspace.onDidCloseNotebookDocument(async (notebook) => {
log.debug(`Notebook closed: ${notebook.uri.path}`);
if (isActiveNotebookEditorUri(notebook.uri)) {
await setHasRunningNotebookSessionContext(false);
}
await notebookSessionService.shutdownRuntimeSession(notebook.uri);
}));

Expand Down Expand Up @@ -53,17 +65,24 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>

// Set the hasRunningNotebookSession context when the active notebook editor changes.
context.subscriptions.push(vscode.window.onDidChangeActiveNotebookEditor(async (editor) => {
const session = editor && await getNotebookSession(editor.notebook.uri);
setHasRunningNotebookSessionContext(Boolean(session));
}));

// Set the hasRunningNotebookSession context when a session is started/shutdown for the active notebook.
context.subscriptions.push(notebookSessionService.onDidChangeNotebookSession((e) => {
if (e.notebookUri.toString() === vscode.window.activeNotebookEditor?.notebook.uri.toString()) {
setHasRunningNotebookSessionContext(Boolean(e.session));
if (editor) {
// Changed to a notebook editor.
const notebookUri = editor.notebook.uri;
const session = await getNotebookSession(notebookUri);
await setHasRunningNotebookSessionContext(Boolean(session));
} else {
// Changed to a non-notebook editor.
await setHasRunningNotebookSessionContext(false);
}
}));

// Set the hasRunningNotebookSession context for the current active notebook editor.
if (vscode.window.activeNotebookEditor) {
const notebookUri = vscode.window.activeNotebookEditor.notebook.uri;
const session = await getNotebookSession(notebookUri);
await setHasRunningNotebookSessionContext(Boolean(session));
}

// Register kernel source action providers for the kernel selection quickpick.
context.subscriptions.push(vscode.notebooks.registerKernelSourceActionProvider(JUPYTER_NOTEBOOK_TYPE, {
provideNotebookKernelSourceActions: () => {
Expand All @@ -85,9 +104,9 @@ export async function activate(context: vscode.ExtensionContext): Promise<void>
registerExecutionInfoStatusBar(context.subscriptions, manager);
}

function setHasRunningNotebookSessionContext(value: boolean): void {
log.debug(`Setting 'positron.hasRunningNotebookSession' context to: ${value}`);
vscode.commands.executeCommand(
export function setHasRunningNotebookSessionContext(value: boolean): Thenable<unknown> {
_onDidSetHasRunningNotebookSessionContext.fire(value);
return vscode.commands.executeCommand(
'setContext',
'positron.hasRunningNotebookSession',
value,
Expand Down
24 changes: 19 additions & 5 deletions extensions/positron-notebook-controllers/src/notebookController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,9 +6,9 @@ import * as vscode from 'vscode';
import * as positron from 'positron';
import { NotebookSessionService } from './notebookSessionService';
import { JUPYTER_NOTEBOOK_TYPE } from './constants';
import { log } from './extension';
import { log, setHasRunningNotebookSessionContext } from './extension';
import { ResourceMap } from './map';
import { getNotebookSession } from './utils';
import { getNotebookSession, isActiveNotebookEditorUri } from './utils';

/** The type of a Jupyter notebook cell output. */
enum NotebookCellOutputType {
Expand Down Expand Up @@ -110,7 +110,7 @@ export class NotebookController implements vscode.Disposable {
this.selectRuntimeSession(e.notebook),
]);
} else {
await this._notebookSessionService.shutdownRuntimeSession(e.notebook.uri);
await this.shutdownRuntimeSession(e.notebook);
}
}));
}
Expand All @@ -122,12 +122,20 @@ export class NotebookController implements vscode.Disposable {

private async selectRuntimeSession(notebook: vscode.NotebookDocument): Promise<void> {
// If there's an existing session from another runtime, shut it down.
await this._notebookSessionService.shutdownRuntimeSession(notebook.uri);
await this.shutdownRuntimeSession(notebook);

// Start the new session.
await this.startRuntimeSession(notebook);
}

private async shutdownRuntimeSession(notebook: vscode.NotebookDocument): Promise<void> {
if (isActiveNotebookEditorUri(notebook.uri)) {
await setHasRunningNotebookSessionContext(false);
}

await this._notebookSessionService.shutdownRuntimeSession(notebook.uri);
}

/**
* Start a runtime session for a notebook.
*
Expand All @@ -136,7 +144,13 @@ export class NotebookController implements vscode.Disposable {
*/
private async startRuntimeSession(notebook: vscode.NotebookDocument): Promise<positron.LanguageRuntimeSession> {
try {
return await this._notebookSessionService.startRuntimeSession(notebook.uri, this._runtimeMetadata.runtimeId);
const session = await this._notebookSessionService.startRuntimeSession(notebook.uri, this._runtimeMetadata.runtimeId);

if (isActiveNotebookEditorUri(notebook.uri)) {
await setHasRunningNotebookSessionContext(true);
}

return session;
} catch (err) {
const retry = vscode.l10n.t('Retry');
const selection = await vscode.window.showErrorMessage(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,6 @@ import { log } from './extension';
import { ResourceMap } from './map';
import { getNotebookSession } from './utils';

export interface INotebookSessionDidChangeEvent {
/** The URI of the notebook corresponding to the session. */
readonly notebookUri: vscode.Uri;

/** The session that was set for the notebook, or undefined if it was deleted. */
readonly session?: positron.LanguageRuntimeSession;
}

/**
* The notebook session service is the main interface for interacting with
* runtime sessions; it manages the set of active sessions and provides
Expand All @@ -27,9 +19,7 @@ export interface INotebookSessionDidChangeEvent {
* required into the runtime session service and expose what's needed via the Positron Extensions
* API.
*/
export class NotebookSessionService implements vscode.Disposable {
private readonly _disposables = new Array<vscode.Disposable>();

export class NotebookSessionService {
/**
* A map of sessions currently starting, keyed by notebook URI. Values are promises that resolve
* when the session has started and is ready to execute code.
Expand All @@ -48,17 +38,6 @@ export class NotebookSessionService implements vscode.Disposable {
*/
private readonly _restartingSessionsByNotebookUri = new ResourceMap<Promise<positron.LanguageRuntimeSession>>();

/** The event emitter for the onDidChangeNotebookSession event. */
private readonly _onDidChangeNotebookSession = this._register(new vscode.EventEmitter<INotebookSessionDidChangeEvent>);

/** An event that fires when a session is set/unset for a notebook. */
readonly onDidChangeNotebookSession = this._onDidChangeNotebookSession.event;

private _register<T extends vscode.Disposable>(disposable: T): T {
this._disposables.push(disposable);
return disposable;
}

/**
* Wait for a notebook session to complete a shutdown sequence.
*
Expand Down Expand Up @@ -99,7 +78,6 @@ export class NotebookSessionService implements vscode.Disposable {
try {
const session = await this.doStartRuntimeSession(notebookUri, runtimeId);
this._startingSessionsByNotebookUri.delete(notebookUri);
this._onDidChangeNotebookSession.fire({ notebookUri, session });
log.info(`Session ${session.metadata.sessionId} is started`);
return session;
} catch (err) {
Expand All @@ -121,26 +99,15 @@ export class NotebookSessionService implements vscode.Disposable {
await shuttingDownSessionPromise;
} catch (err) {
log.error(`Waiting for notebook runtime to shutdown before starting failed. Reason ${err}`);
throw err;
// Try to start a new session anyway.
}
}

// Start the session.
try {
const session = await positron.runtime.startLanguageRuntime(
runtimeId,
path.basename(notebookUri.path), // Use the notebook's file name as the session name.
notebookUri);
log.info(
`Starting session for language runtime ${session.metadata.sessionId} `
+ `(language: ${session.runtimeMetadata.languageName}, name: ${session.runtimeMetadata.runtimeName}, `
+ `version: ${session.runtimeMetadata.runtimeVersion}, notebook: ${notebookUri.path})`
);
return session;
} catch (err) {
log.error(`Starting session for language runtime ${runtimeId} failed. Reason: ${err}`);
throw err;
}
return positron.runtime.startLanguageRuntime(
runtimeId,
path.basename(notebookUri.path), // Use the notebook's file name as the session name.
notebookUri);
}

/**
Expand All @@ -161,7 +128,6 @@ export class NotebookSessionService implements vscode.Disposable {
try {
await this.doShutdownRuntimeSession(notebookUri);
this._shuttingDownSessionsByNotebookUri.delete(notebookUri);
this._onDidChangeNotebookSession.fire({ notebookUri, session: undefined });
} catch (err) {
this._startingSessionsByNotebookUri.delete(notebookUri);
throw err;
Expand Down Expand Up @@ -255,7 +221,6 @@ export class NotebookSessionService implements vscode.Disposable {
try {
const session = await this.doRestartRuntimeSession(notebookUri);
this._restartingSessionsByNotebookUri.delete(notebookUri);
this._onDidChangeNotebookSession.fire({ notebookUri, session });
log.info(`Session ${session.metadata.sessionId} is restarted`);
return session;
} catch (err) {
Expand All @@ -276,10 +241,6 @@ export class NotebookSessionService implements vscode.Disposable {
throw new Error(`Tried to restart runtime for notebook without a running runtime: ${notebookUri.path}`);
}

// Remove the session from the map of active notebooks in case it's accessed while we're
// restarting.
this._onDidChangeNotebookSession.fire({ notebookUri, session: undefined });

// If the notebook's session is still shutting down, wait for it to finish.
const shuttingDownSessionPromise = this._shuttingDownSessionsByNotebookUri.get(notebookUri);
if (shuttingDownSessionPromise) {
Expand Down Expand Up @@ -322,8 +283,4 @@ export class NotebookSessionService implements vscode.Disposable {

return session;
}

dispose() {
this._disposables.forEach(d => d.dispose());
}
}
61 changes: 61 additions & 0 deletions extensions/positron-notebook-controllers/src/test/commands.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
/*---------------------------------------------------------------------------------------------
* Copyright (C) 2024 Posit Software, PBC. All rights reserved.
* Licensed under the Elastic License 2.0. See LICENSE.txt for license information.
*--------------------------------------------------------------------------------------------*/

import * as assert from 'assert';
import * as positron from 'positron';
import * as sinon from 'sinon';
import * as vscode from 'vscode';
import { closeAllEditors, openTestJupyterNotebookDocument } from './utils';
import { TestLanguageRuntimeSession } from './testLanguageRuntimeSession';
import { onDidSetHasRunningNotebookSessionContext } from '../extension';

suite('commands', () => {
let disposables: vscode.Disposable[];
let session: TestLanguageRuntimeSession;

setup(() => {
disposables = [];
session = new TestLanguageRuntimeSession();
disposables.push(session);
});

teardown(async () => {
await closeAllEditors();
vscode.Disposable.from(...disposables).dispose();
sinon.restore();
});

// TODO: This test is currently flaky. It should become more stable once the restartKernel
// calls directly to Positron and not via the notebook session service.
test.skip('restart disables then enables hasRunningNotebookSession context', async () => {
// Simulate a running session for the notebook.
sinon.stub(positron.runtime, 'getNotebookSession').resolves(session as any);

// Simulate a successful restart.
sinon.stub(positron.runtime, 'restartSession').callsFake(async () => {
session.setRuntimeState(positron.RuntimeState.Ready);
});

// Open a test Jupyter notebook.
await openTestJupyterNotebookDocument();

// Capture the first two hasRunningNotebookSession context values.
const promise = new Promise<boolean[]>(resolve => {
const values: boolean[] = [];
disposables.push(onDidSetHasRunningNotebookSessionContext(value => {
values.push(value);
if (values.length === 2) {
resolve(values);
}
}));
});

// Restart.
await vscode.commands.executeCommand('positron.restartKernel');

// Assert that the context is first set to false, then true.
assert.deepStrictEqual(await promise, [false, true]);
});
});
Loading

0 comments on commit 70ed8a3

Please sign in to comment.