From a36f4d7ccf01d05c97dee2b52020e59229900841 Mon Sep 17 00:00:00 2001 From: Wasim Lorgat Date: Thu, 28 Nov 2024 15:44:31 +0200 Subject: [PATCH] Improve runtime session service stability (#5380) The aim of this PR is to improve the stability of the runtime session service, with a slight focus on notebooks. This is another step toward https://github.com/posit-dev/positron/issues/2671. I've also added an extensive suite of unit tests for the runtime session service, for two reasons: 1. Some of the issues we're seeing with notebooks are intermittent and therefore tricky to reproduce without tests. 2. I wanted to ensure that I didn't break any existing behavior with these changes as well as more planned changes to add notebook-runtime methods to the extension API. --- .../src/notebookController.ts | 9 + .../browser/positronPreview.contribution.ts | 3 +- .../common/languageRuntimeUiClient.ts | 2 +- .../runtimeSession/common/runtimeSession.ts | 355 ++++++-- .../test/common/runtimeSession.test.ts | 840 ++++++++++++++++++ .../test/common/testLanguageRuntimeSession.ts | 42 +- .../test/common/testRuntimeSessionService.ts | 2 +- .../common/positronWorkbenchTestServices.ts | 15 +- 8 files changed, 1168 insertions(+), 100 deletions(-) create mode 100644 src/vs/workbench/services/runtimeSession/test/common/runtimeSession.test.ts diff --git a/extensions/positron-notebook-controllers/src/notebookController.ts b/extensions/positron-notebook-controllers/src/notebookController.ts index aaac1770755..04af82b2549 100644 --- a/extensions/positron-notebook-controllers/src/notebookController.ts +++ b/extensions/positron-notebook-controllers/src/notebookController.ts @@ -331,10 +331,19 @@ function executeCode( } } }).catch(error => { + // Stop listening for replies. handler.dispose(); + + // Reject the outer execution promise since we've encountered an error. reject(error); + + // Rethrow the error to stop any replies that are chained to this promise. throw error; }); + + // Avoid unhandled rejections being logged to the console. + // The actual error-handling is in the catch block above. + currentMessagePromise.catch(() => { }); }); // Execute the cell. diff --git a/src/vs/workbench/contrib/positronPreview/browser/positronPreview.contribution.ts b/src/vs/workbench/contrib/positronPreview/browser/positronPreview.contribution.ts index 3de51e015b1..d24db8db586 100644 --- a/src/vs/workbench/contrib/positronPreview/browser/positronPreview.contribution.ts +++ b/src/vs/workbench/contrib/positronPreview/browser/positronPreview.contribution.ts @@ -20,12 +20,11 @@ import { ViewContainer, IViewContainersRegistry, ViewContainerLocation, Extensio import { registerAction2 } from 'vs/platform/actions/common/actions'; import { PositronOpenUrlInViewerAction } from 'vs/workbench/contrib/positronPreview/browser/positronPreviewActions'; import { IConfigurationRegistry, Extensions as ConfigurationExtensions, ConfigurationScope, } from 'vs/platform/configuration/common/configurationRegistry'; +import { POSITRON_PREVIEW_PLOTS_IN_VIEWER } from 'vs/workbench/services/languageRuntime/common/languageRuntimeUiClient'; // The Positron preview view icon. const positronPreviewViewIcon = registerIcon('positron-preview-view-icon', Codicon.positronPreviewView, nls.localize('positronPreviewViewIcon', 'View icon of the Positron preview view.')); -export const POSITRON_PREVIEW_PLOTS_IN_VIEWER = 'positron.viewer.interactivePlotsInViewer'; - // Register the Positron preview container. const VIEW_CONTAINER: ViewContainer = Registry.as(ViewContainerExtensions.ViewContainersRegistry).registerViewContainer({ id: POSITRON_PREVIEW_VIEW_ID, diff --git a/src/vs/workbench/services/languageRuntime/common/languageRuntimeUiClient.ts b/src/vs/workbench/services/languageRuntime/common/languageRuntimeUiClient.ts index d9fd4ca94a6..83ec0da31ce 100644 --- a/src/vs/workbench/services/languageRuntime/common/languageRuntimeUiClient.ts +++ b/src/vs/workbench/services/languageRuntime/common/languageRuntimeUiClient.ts @@ -13,8 +13,8 @@ import { URI } from 'vs/base/common/uri'; import { ICommandService } from 'vs/platform/commands/common/commands'; import { ILogService } from 'vs/platform/log/common/log'; import { IConfigurationService } from 'vs/platform/configuration/common/configuration'; -import { POSITRON_PREVIEW_PLOTS_IN_VIEWER } from 'vs/workbench/contrib/positronPreview/browser/positronPreview.contribution'; +export const POSITRON_PREVIEW_PLOTS_IN_VIEWER = 'positron.viewer.interactivePlotsInViewer'; /** * The types of messages that can be sent to the backend. diff --git a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts index c5b6e4a6358..ec9d2016eab 100644 --- a/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts +++ b/src/vs/workbench/services/runtimeSession/common/runtimeSession.ts @@ -4,7 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import * as nls from 'vs/nls'; -import { DeferredPromise } from 'vs/base/common/async'; +import { DeferredPromise, disposableTimeout } from 'vs/base/common/async'; import { Emitter } from 'vs/base/common/event'; import { Disposable, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle'; import { URI } from 'vs/base/common/uri'; @@ -264,15 +264,20 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession * no console session with the given runtime identifier exists. */ getConsoleSessionForRuntime(runtimeId: string): ILanguageRuntimeSession | undefined { - const session = Array.from(this._activeSessionsBySessionId.values()).find(session => - session.session.runtimeMetadata.runtimeId === runtimeId && - session.session.metadata.sessionMode === LanguageRuntimeSessionMode.Console && - session.state !== RuntimeState.Exited); - if (session) { - return session.session; - } else { - return undefined; - } + // It's possible that there are multiple consoles for the same runtime, + // for example, if one failed to start and is uninitialized. In that case, + // we return the most recently created. + return Array.from(this._activeSessionsBySessionId.values()) + .map((info, index) => ({ info, index })) + .sort((a, b) => + b.info.session.metadata.createdTimestamp - a.info.session.metadata.createdTimestamp + // If the timestamps are the same, prefer the session that was inserted last. + || b.index - a.index) + .find(({ info }) => + info.session.runtimeMetadata.runtimeId === runtimeId && + info.session.metadata.sessionMode === LanguageRuntimeSessionMode.Console && + info.state !== RuntimeState.Exited) + ?.info.session; } /** @@ -311,8 +316,7 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession async selectRuntime(runtimeId: string, source: string): Promise { const runtime = this._languageRuntimeService.getRegisteredRuntime(runtimeId); if (!runtime) { - return Promise.reject(new Error(`Language runtime ID '${runtimeId}' ` + - `is not registered.`)); + throw new Error(`No language runtime with id '${runtimeId}' was found.`); } // Shut down any other runtime consoles for the language. @@ -365,28 +369,38 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession private async doShutdownRuntimeSession( session: ILanguageRuntimeSession, exitReason: RuntimeExitReason): Promise { + + const sessionDisposables = this._sessionDisposables.get(session.sessionId); + if (!sessionDisposables) { + throw new Error(`No disposables found for session ${session.sessionId}`); + } + // We wait for `onDidEndSession()` rather than `RuntimeState.Exited`, because the former // generates some Console output that must finish before starting up a new runtime: - const promise = new Promise(resolve => { - const disposable = session.onDidEndSession((exit) => { + const disposables = sessionDisposables.add(new DisposableStore()); + const promise = new Promise((resolve, reject) => { + disposables.add(session.onDidEndSession((exit) => { + disposables.dispose(); resolve(); - disposable.dispose(); - }); - }); - - const timeout = new Promise((_, reject) => { - setTimeout(() => { + })); + disposables.add(disposableTimeout(() => { + disposables.dispose(); reject(new Error(`Timed out waiting for runtime ` + `${formatLanguageRuntimeSession(session)} to finish exiting.`)); - }, 5000); + }, 5000)); }); // Ask the runtime to shut down. - await session.shutdown(exitReason); + try { + await session.shutdown(exitReason); + } catch (error) { + disposables.dispose(); + throw error; + } // Wait for the runtime onDidEndSession to resolve, or for the timeout to expire // (whichever comes first) - await Promise.race([promise, timeout]); + await promise; } /** @@ -419,38 +433,19 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession throw new Error(`No language runtime with id '${runtimeId}' was found.`); } - // If there is already a runtime starting for the language, throw an error. - if (sessionMode === LanguageRuntimeSessionMode.Console) { - const startingLanguageRuntime = this._startingConsolesByLanguageId.get( - languageRuntime.languageId); - if (startingLanguageRuntime) { - throw new Error(`Session for language runtime ${formatLanguageRuntimeMetadata(languageRuntime)} cannot be started because language runtime ${formatLanguageRuntimeMetadata(startingLanguageRuntime)} is already starting for the language. Request source: ${source}`); - } - - // If there is already a runtime running for the language, throw an error. - const runningLanguageRuntime = - this._consoleSessionsByLanguageId.get(languageRuntime.languageId); - if (runningLanguageRuntime) { - const metadata = runningLanguageRuntime.runtimeMetadata; - if (metadata.runtimeId === runtimeId) { - // If the runtime that is running is the one we were just asked - // to start, we're technically in good shape since the runtime - // is already running! - return runningLanguageRuntime.sessionId; - } else { - throw new Error(`A console for ` + - `${formatLanguageRuntimeMetadata(languageRuntime)} ` + - `cannot be started because a console for ` + - `${formatLanguageRuntimeMetadata(metadata)} is already running ` + - `for the ${metadata.languageName} language.`); - } - } + const runningSessionId = this.validateRuntimeSessionStart(sessionMode, languageRuntime, notebookUri, source); + if (runningSessionId) { + return runningSessionId; } // If the workspace is not trusted, defer starting the runtime until the // workspace is trusted. if (!this._workspaceTrustManagementService.isWorkspaceTrusted()) { - return this.autoStartRuntime(languageRuntime, source); + if (sessionMode === LanguageRuntimeSessionMode.Console) { + return this.autoStartRuntime(languageRuntime, source); + } else { + throw new Error(`Cannot start a ${sessionMode} session in an untrusted workspace.`); + } } // Start the runtime. @@ -471,6 +466,16 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession runtimeMetadata: ILanguageRuntimeMetadata, sessionMetadata: IRuntimeSessionMetadata): Promise { + // See if we are already starting the requested session. If we + // are, return the promise that resolves when the session is ready to + // use. This makes it possible for multiple requests to start the same + // session to be coalesced. + const startingRuntimePromise = this._startingSessionsBySessionMapKey.get( + getSessionMapKey(sessionMetadata.sessionMode, runtimeMetadata.runtimeId, sessionMetadata.notebookUri)); + if (startingRuntimePromise && !startingRuntimePromise.isSettled) { + return startingRuntimePromise.p.then(() => { }); + } + // Ensure that the runtime is registered. const languageRuntime = this._languageRuntimeService.getRegisteredRuntime( runtimeMetadata.runtimeId); @@ -480,9 +485,10 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession this._languageRuntimeService.registerRuntime(runtimeMetadata); } - // Add the runtime to the starting runtimes, if it's a console session. - if (sessionMetadata.sessionMode === LanguageRuntimeSessionMode.Console) { - this._startingConsolesByLanguageId.set(runtimeMetadata.languageId, runtimeMetadata); + const runningSessionId = this.validateRuntimeSessionStart( + sessionMetadata.sessionMode, runtimeMetadata, sessionMetadata.notebookUri); + if (runningSessionId) { + return; } // Create a promise that resolves when the runtime is ready to use. @@ -491,6 +497,13 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession sessionMetadata.sessionMode, runtimeMetadata.runtimeId, sessionMetadata.notebookUri); this._startingSessionsBySessionMapKey.set(sessionMapKey, startPromise); + // It's possible that startPromise is never awaited, so we log any errors here + // at the debug level since we still expect the error to be handled/logged elsewhere. + startPromise.p.catch((err) => this._logService.debug(`Error starting session: ${err}`)); + + this.setStartingSessionMaps( + sessionMetadata.sessionMode, runtimeMetadata, sessionMetadata.notebookUri); + // We should already have a session manager registered, since we can't // get here until the extension host has been activated. if (this._sessionManagers.length === 0) { @@ -509,8 +522,8 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession `Reconnecting to session '${sessionMetadata.sessionId}' for language runtime ` + `${formatLanguageRuntimeMetadata(runtimeMetadata)} failed. Reason: ${err}`); startPromise.error(err); - this._startingSessionsBySessionMapKey.delete(sessionMapKey); - this._startingConsolesByLanguageId.delete(runtimeMetadata.languageId); + this.clearStartingSessionMaps( + sessionMetadata.sessionMode, runtimeMetadata, sessionMetadata.notebookUri); throw err; } @@ -631,18 +644,19 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession `${formatLanguageRuntimeMetadata(metadata)} (Source: ${source}) ` + `because workspace trust has not been granted. ` + `The runtime will be started when workspace trust is granted.`); - this._workspaceTrustManagementService.onDidChangeTrust((trusted) => { + const disposable = this._register(this._workspaceTrustManagementService.onDidChangeTrust((trusted) => { if (!trusted) { // If the workspace is still not trusted, do nothing. - return ''; + return; } // If the workspace is trusted, start the runtime. + disposable.dispose(); this._logService.info(`Language runtime ` + `${formatLanguageRuntimeMetadata(metadata)} ` + `automatically starting after workspace trust was granted. ` + `Source: ${source}`); - return this.doAutoStartRuntime(metadata, source); - }); + this.doAutoStartRuntime(metadata, source); + })); } return ''; @@ -698,6 +712,38 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession metadata: ILanguageRuntimeMetadata, source: string): Promise { + // Auto-started runtimes are (currently) always console sessions. + const sessionMode = LanguageRuntimeSessionMode.Console; + const notebookUri = undefined; + + // See if we are already starting the requested session. If we + // are, return the promise that resolves when the session is ready to + // use. This makes it possible for multiple requests to start the same + // session to be coalesced. + const startingRuntimePromise = this._startingSessionsBySessionMapKey.get( + getSessionMapKey(sessionMode, metadata.runtimeId, notebookUri)); + if (startingRuntimePromise && !startingRuntimePromise.isSettled) { + return startingRuntimePromise.p; + } + + const runningSessionId = this.validateRuntimeSessionStart(sessionMode, metadata, notebookUri, source); + if (runningSessionId) { + return runningSessionId; + } + + // Before attempting to validate the runtime, add it to the set of + // starting consoles. + this._startingConsolesByLanguageId.set(metadata.languageId, metadata); + + // Create a promise that resolves when the runtime is ready to use. + const startPromise = new DeferredPromise(); + const sessionMapKey = getSessionMapKey(sessionMode, metadata.runtimeId, notebookUri); + this._startingSessionsBySessionMapKey.set(sessionMapKey, startPromise); + + // It's possible that startPromise is never awaited, so we log any errors here + // at the debug level since we still expect the error to be handled/logged elsewhere. + startPromise.p.catch(err => this._logService.debug(`Error starting runtime session: ${err}`)); + // Check to see if the runtime has already been registered with the // language runtime service. const languageRuntime = @@ -707,10 +753,6 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession // If it has not been registered, validate the metadata. if (!languageRuntime) { try { - // Before attempting to validate the runtime, add it to the set of - // starting consoles. - this._startingConsolesByLanguageId.set(metadata.languageId, metadata); - // Attempt to validate the metadata. Note that this can throw if the metadata // is invalid! const validated = await sessionManager.validateMetadata(metadata); @@ -762,9 +804,7 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession } } - // Auto-started runtimes are (currently) always console sessions. - return this.doCreateRuntimeSession(metadata, metadata.runtimeName, - LanguageRuntimeSessionMode.Console, source); + return this.doCreateRuntimeSession(metadata, metadata.runtimeName, sessionMode, source, notebookUri); } /** @@ -784,15 +824,19 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession sessionMode: LanguageRuntimeSessionMode, source: string, notebookUri?: URI): Promise { - // Add the runtime to the starting runtimes. - if (sessionMode === LanguageRuntimeSessionMode.Console) { - this._startingConsolesByLanguageId.set(runtimeMetadata.languageId, runtimeMetadata); - } + this.setStartingSessionMaps(sessionMode, runtimeMetadata, notebookUri); - // Create a promise that resolves when the runtime is ready to use. - const startPromise = new DeferredPromise(); + // Create a promise that resolves when the runtime is ready to use, if there isn't already one. const sessionMapKey = getSessionMapKey(sessionMode, runtimeMetadata.runtimeId, notebookUri); - this._startingSessionsBySessionMapKey.set(sessionMapKey, startPromise); + let startPromise = this._startingSessionsBySessionMapKey.get(sessionMapKey); + if (!startPromise || startPromise.isSettled) { + startPromise = new DeferredPromise(); + this._startingSessionsBySessionMapKey.set(sessionMapKey, startPromise); + + // It's possible that startPromise is never awaited, so we log any errors here + // at the debug level since we still expect the error to be handled/logged elsewhere. + startPromise.p.catch(err => this._logService.debug(`Error starting runtime session: ${err}`)); + } const sessionManager = await this.getManagerForRuntime(runtimeMetadata); const sessionId = this.generateNewSessionId(runtimeMetadata); @@ -814,8 +858,7 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession `Creating session for language runtime ` + `${formatLanguageRuntimeMetadata(runtimeMetadata)} failed. Reason: ${err}`); startPromise.error(err); - this._startingSessionsBySessionMapKey.delete(sessionMapKey); - this._startingConsolesByLanguageId.delete(runtimeMetadata.languageId); + this.clearStartingSessionMaps(sessionMode, runtimeMetadata, notebookUri); // Re-throw the error. throw err; @@ -854,21 +897,18 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession // Attach event handlers to the newly provisioned session. this.attachToSession(session, manager); - const sessionMapKey = getSessionMapKey( - session.metadata.sessionMode, session.runtimeMetadata.runtimeId, session.metadata.notebookUri); try { // Attempt to start, or reconnect to, the session. await session.start(); // The runtime started. Move it from the starting runtimes to the // running runtimes. - this._startingSessionsBySessionMapKey.delete(sessionMapKey); + this.clearStartingSessionMaps( + session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri); if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) { - this._startingConsolesByLanguageId.delete(session.runtimeMetadata.languageId); this._consoleSessionsByLanguageId.set(session.runtimeMetadata.languageId, session); } else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) { if (session.metadata.notebookUri) { - this._startingNotebooksByNotebookUri.delete(session.metadata.notebookUri); this._logService.info(`Notebook session for ${session.metadata.notebookUri} started: ${session.metadata.sessionId}`); this._notebookSessionsByNotebookUri.set(session.metadata.notebookUri, session); } else { @@ -885,10 +925,8 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession this._foregroundSession = session; } } catch (reason) { - - // Remove the runtime from the starting runtimes. - this._startingConsolesByLanguageId.delete(session.runtimeMetadata.languageId); - this._startingSessionsBySessionMapKey.delete(sessionMapKey); + this.clearStartingSessionMaps( + session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri); // Fire the onDidFailStartRuntime event. this._onDidFailStartRuntimeEmitter.fire(session); @@ -960,17 +998,17 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession this.foregroundSession = session; } - // If this is a console session and there isn't already a console session - // for this language, set this one as the console session. - // (This restores the console session in the case of a - // restart) + // Restore the session in the case of a restart. if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console && !this._consoleSessionsByLanguageId.has(session.runtimeMetadata.languageId)) { this._consoleSessionsByLanguageId.set(session.runtimeMetadata.languageId, session); + } else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook && + session.metadata.notebookUri && + !this._notebookSessionsByNotebookUri.has(session.metadata.notebookUri)) { + this._notebookSessionsByNotebookUri.set(session.metadata.notebookUri, session); } - // Start the UI client instance once the runtime is fully online. this.startUiClient(session); break; @@ -1042,6 +1080,132 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession })); } + /** + * Validate whether a runtime session can be started. + * + * @param sessionMode The mode of the new session. + * @param languageRuntime The metadata of the runtime to start. + * @param notebookUri The notebook URI to attach to the session, if any. + * @param source The source of the request to start the runtime, if known. + * @throws An error if the session cannot be started. + * @returns A session ID if a session is already running that matches the request, or undefined. + */ + private validateRuntimeSessionStart( + sessionMode: LanguageRuntimeSessionMode, + languageRuntime: ILanguageRuntimeMetadata, + notebookUri: URI | undefined, + source?: string, + ): string | undefined { + // If there is already a runtime starting for the language, throw an error. + if (sessionMode === LanguageRuntimeSessionMode.Console) { + const startingLanguageRuntime = this._startingConsolesByLanguageId.get( + languageRuntime.languageId); + if (startingLanguageRuntime) { + throw new Error(`Session for language runtime ` + + `${formatLanguageRuntimeMetadata(languageRuntime)} ` + + `cannot be started because language runtime ` + + `${formatLanguageRuntimeMetadata(startingLanguageRuntime)} ` + + `is already starting for the language.` + + (source ? ` Request source: ${source}` : ``)); + } + + // If there is already a runtime running for the language, throw an error. + const runningLanguageRuntime = + this._consoleSessionsByLanguageId.get(languageRuntime.languageId); + if (runningLanguageRuntime) { + const metadata = runningLanguageRuntime.runtimeMetadata; + if (metadata.runtimeId === languageRuntime.runtimeId) { + // If the runtime that is running is the one we were just asked + // to start, we're technically in good shape since the runtime + // is already running! + return runningLanguageRuntime.sessionId; + } else { + throw new Error(`A console for ` + + `${formatLanguageRuntimeMetadata(languageRuntime)} ` + + `cannot be started because a console for ` + + `${formatLanguageRuntimeMetadata(metadata)} is already running ` + + `for the ${metadata.languageName} language.` + + (source ? ` Request source: ${source}` : ``)); + } + } + } else if (sessionMode === LanguageRuntimeSessionMode.Notebook) { + // If no notebook URI is provided, throw an error. + if (!notebookUri) { + throw new Error(`A notebook URI must be provided when starting a notebook session.`); + } + + // If there is already a runtime starting for the notebook, throw an error. + const startingLanguageRuntime = this._startingNotebooksByNotebookUri.get(notebookUri); + if (startingLanguageRuntime) { + throw new Error(`Session for language runtime ` + + `${formatLanguageRuntimeMetadata(languageRuntime)} ` + + `cannot be started because language runtime ` + + `${formatLanguageRuntimeMetadata(startingLanguageRuntime)} ` + + `is already starting for the notebook ${notebookUri.toString()}.` + + (source ? ` Request source: ${source}` : ``)); + } + + // If there is already a runtime running for the notebook, throw an error. + const runningLanguageRuntime = this._notebookSessionsByNotebookUri.get(notebookUri); + if (runningLanguageRuntime) { + const metadata = runningLanguageRuntime.runtimeMetadata; + if (metadata.runtimeId === languageRuntime.runtimeId) { + // If the runtime that is running is the one we were just asked + // to start, we're technically in good shape since the runtime + // is already running! + return runningLanguageRuntime.sessionId; + } else { + throw new Error(`A notebook for ` + + `${formatLanguageRuntimeMetadata(languageRuntime)} ` + + `cannot be started because a notebook for ` + + `${formatLanguageRuntimeMetadata(metadata)} is already running ` + + `for the URI ${notebookUri.toString()}.` + + (source ? ` Request source: ${source}` : ``)); + } + } + } + + return undefined; + } + + /** + * Sets the session maps for a starting session. + * + * @param sessionMode The mode of the session. + * @param runtimeMetadata The metadata of the session's runtime. + * @param notebookUri The notebook URI attached to the session, if any. + */ + private setStartingSessionMaps( + sessionMode: LanguageRuntimeSessionMode, + runtimeMetadata: ILanguageRuntimeMetadata, + notebookUri?: URI) { + if (sessionMode === LanguageRuntimeSessionMode.Console) { + this._startingConsolesByLanguageId.set(runtimeMetadata.languageId, runtimeMetadata); + } else if (sessionMode === LanguageRuntimeSessionMode.Notebook && notebookUri) { + this._startingNotebooksByNotebookUri.set(notebookUri, runtimeMetadata); + } + } + + /** + * Clears the session maps for a starting session. + * + * @param sessionMode The mode of the session. + * @param runtimeMetadata The metadata of the session's runtime. + * @param notebookUri The notebook URI attached to the session, if any. + */ + private clearStartingSessionMaps( + sessionMode: LanguageRuntimeSessionMode, + runtimeMetadata: ILanguageRuntimeMetadata, + notebookUri?: URI) { + const sessionMapKey = getSessionMapKey(sessionMode, runtimeMetadata.runtimeId, notebookUri); + this._startingSessionsBySessionMapKey.delete(sessionMapKey); + if (sessionMode === LanguageRuntimeSessionMode.Console) { + this._startingConsolesByLanguageId.delete(runtimeMetadata.languageId); + } else if (sessionMode === LanguageRuntimeSessionMode.Notebook && notebookUri) { + this._startingNotebooksByNotebookUri.delete(notebookUri); + } + } + /** * Updates the session maps (for active consoles, notebooks, etc.), after a * session exits. @@ -1079,6 +1243,19 @@ export class RuntimeSessionService extends Disposable implements IRuntimeSession state === RuntimeState.Ready) { // The runtime looks like it could handle a restart request, so send // one over. + + // Mark the session as starting until the restart sequence completes. + this.setStartingSessionMaps( + session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri); + const disposable = this._register(session.onDidChangeRuntimeState((state) => { + if (state === RuntimeState.Ready) { + disposable.dispose(); + this.clearStartingSessionMaps( + session.metadata.sessionMode, session.runtimeMetadata, session.metadata.notebookUri); + } + })); + + // Restart the session. return session.restart(); } else if (state === RuntimeState.Uninitialized || state === RuntimeState.Exited) { diff --git a/src/vs/workbench/services/runtimeSession/test/common/runtimeSession.test.ts b/src/vs/workbench/services/runtimeSession/test/common/runtimeSession.test.ts new file mode 100644 index 00000000000..ea178cedeea --- /dev/null +++ b/src/vs/workbench/services/runtimeSession/test/common/runtimeSession.test.ts @@ -0,0 +1,840 @@ +/*--------------------------------------------------------------------------------------------- + * 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 sinon from 'sinon'; +import { Event } from 'vs/base/common/event'; +import { URI } from 'vs/base/common/uri'; +import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils'; +import { IConfigurationService } from 'vs/platform/configuration/common/configuration'; +import { TestConfigurationService } from 'vs/platform/configuration/test/common/testConfigurationService'; +import { TestInstantiationService } from 'vs/platform/instantiation/test/common/instantiationServiceMock'; +import { IWorkspaceTrustManagementService } from 'vs/platform/workspace/common/workspaceTrust'; +import { formatLanguageRuntimeMetadata, formatLanguageRuntimeSession, ILanguageRuntimeMetadata, ILanguageRuntimeService, LanguageRuntimeSessionMode, RuntimeExitReason, RuntimeState } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService'; +import { ILanguageRuntimeSession, IRuntimeSessionMetadata, IRuntimeSessionService, IRuntimeSessionWillStartEvent } from 'vs/workbench/services/runtimeSession/common/runtimeSessionService'; +import { TestLanguageRuntimeSession, waitForRuntimeState } from 'vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession'; +import { createRuntimeServices, createTestLanguageRuntimeMetadata, startTestLanguageRuntimeSession } from 'vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService'; +import { TestRuntimeSessionManager } from 'vs/workbench/test/common/positronWorkbenchTestServices'; +import { TestWorkspaceTrustManagementService } from 'vs/workbench/test/common/workbenchTestServices'; + +type IStartSessionTask = (runtimeMetadata?: ILanguageRuntimeMetadata) => Promise; + +suite('Positron - RuntimeSessionService', () => { + const disposables = ensureNoDisposablesAreLeakedInTestSuite(); + const startReason = 'Test requested to start a runtime session'; + const notebookUri = URI.file('/path/to/notebook'); + let instantiationService: TestInstantiationService; + let languageRuntimeService: ILanguageRuntimeService; + let runtimeSessionService: IRuntimeSessionService; + let configService: TestConfigurationService; + let workspaceTrustManagementService: TestWorkspaceTrustManagementService; + let manager: TestRuntimeSessionManager; + let runtime: ILanguageRuntimeMetadata; + let anotherRuntime: ILanguageRuntimeMetadata; + let sessionName: string; + let unregisteredRuntime: ILanguageRuntimeMetadata; + + setup(() => { + instantiationService = disposables.add(new TestInstantiationService()); + createRuntimeServices(instantiationService, disposables); + languageRuntimeService = instantiationService.get(ILanguageRuntimeService); + runtimeSessionService = instantiationService.get(IRuntimeSessionService); + configService = instantiationService.get(IConfigurationService) as TestConfigurationService; + workspaceTrustManagementService = instantiationService.get(IWorkspaceTrustManagementService) as TestWorkspaceTrustManagementService; + manager = TestRuntimeSessionManager.instance; + + runtime = createTestLanguageRuntimeMetadata(instantiationService, disposables); + anotherRuntime = createTestLanguageRuntimeMetadata(instantiationService, disposables); + sessionName = runtime.runtimeName; + unregisteredRuntime = { runtimeId: 'unregistered-runtime-id' } as ILanguageRuntimeMetadata; + + // Enable automatic startup. + configService.setUserConfiguration('positron.interpreters.automaticStartup', true); + + // Trust the workspace. + workspaceTrustManagementService.setWorkspaceTrust(true); + }); + + function startSession( + runtimeMetadata = runtime, + sessionMode: LanguageRuntimeSessionMode, + notebookUri?: URI, + ) { + return startTestLanguageRuntimeSession( + instantiationService, + disposables, + { + runtime: runtimeMetadata, + sessionName, + startReason, + sessionMode, + notebookUri, + }, + ); + } + + function startConsole(runtimeMetadata?: ILanguageRuntimeMetadata) { + return startSession(runtimeMetadata, LanguageRuntimeSessionMode.Console); + } + + function startNotebook(runtimeMetadata?: ILanguageRuntimeMetadata, notebookUri_ = notebookUri) { + return startSession(runtimeMetadata, LanguageRuntimeSessionMode.Notebook, notebookUri_); + } + + interface IServiceState { + hasStartingOrRunningConsole?: boolean; + consoleSession?: ILanguageRuntimeSession; + consoleSessionForLanguage?: ILanguageRuntimeSession; + consoleSessionForRuntime?: ILanguageRuntimeSession; + notebookSession?: ILanguageRuntimeSession; + notebookSessionForNotebookUri?: ILanguageRuntimeSession; + activeSessions?: ILanguageRuntimeSession[]; + } + + function assertServiceState(expectedState?: IServiceState, runtimeMetadata = runtime): void { + // Check the active sessions. + assert.deepStrictEqual(runtimeSessionService.activeSessions, expectedState?.activeSessions ?? []); + + // Check the console session state. + assert.strictEqual( + runtimeSessionService.hasStartingOrRunningConsole(runtimeMetadata.languageId), + expectedState?.hasStartingOrRunningConsole ?? false, + expectedState?.hasStartingOrRunningConsole ? + 'Expected a starting or running console session' : + 'Expected no starting or running console session', + ); + assert.strictEqual( + runtimeSessionService.getConsoleSessionForLanguage(runtimeMetadata.languageId), + expectedState?.consoleSessionForLanguage, + ); + assert.strictEqual( + runtimeSessionService.getConsoleSessionForRuntime(runtimeMetadata.runtimeId), + expectedState?.consoleSessionForRuntime, + ); + assert.strictEqual( + runtimeSessionService.getSession(expectedState?.consoleSession?.sessionId ?? ''), + expectedState?.consoleSession, + ); + + // Check the notebook session state. + assert.strictEqual( + runtimeSessionService.getSession(expectedState?.notebookSession?.sessionId ?? ''), + expectedState?.notebookSession, + ); + assert.strictEqual( + runtimeSessionService.getNotebookSessionForNotebookUri(notebookUri), + expectedState?.notebookSessionForNotebookUri, + ); + } + + function assertSingleSessionWillStart(sessionMode: LanguageRuntimeSessionMode) { + if (sessionMode === LanguageRuntimeSessionMode.Console) { + assertServiceState({ hasStartingOrRunningConsole: true }); + } else if (sessionMode === LanguageRuntimeSessionMode.Notebook) { + assertServiceState(); + } + } + + function assertHasSingleSession(session: ILanguageRuntimeSession) { + if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Console) { + assertServiceState({ + hasStartingOrRunningConsole: true, + consoleSession: session, + consoleSessionForLanguage: session, + consoleSessionForRuntime: session, + activeSessions: [session], + }, session.runtimeMetadata); + } else if (session.metadata.sessionMode === LanguageRuntimeSessionMode.Notebook) { + assertServiceState({ + notebookSession: session, + notebookSessionForNotebookUri: session, + activeSessions: [session], + }, session.runtimeMetadata); + } + } + + function assertSingleSessionIsStarting(session: ILanguageRuntimeSession) { + assertHasSingleSession(session); + assert.strictEqual(session.getRuntimeState(), RuntimeState.Starting); + } + + function assertSingleSessionIsRestarting(session: ILanguageRuntimeSession) { + assertHasSingleSession(session); + assert.strictEqual(session.getRuntimeState(), RuntimeState.Restarting); + } + + function assertSingleSessionIsReady(session: ILanguageRuntimeSession) { + assertHasSingleSession(session); + assert.strictEqual(session.getRuntimeState(), RuntimeState.Ready); + } + + async function restoreSession( + sessionMetadata: IRuntimeSessionMetadata, runtimeMetadata = runtime, + ) { + await runtimeSessionService.restoreRuntimeSession(runtimeMetadata, sessionMetadata); + + // Ensure that the session gets disposed after the test. + const session = runtimeSessionService.getSession(sessionMetadata.sessionId); + assert.ok(session instanceof TestLanguageRuntimeSession); + disposables.add(session); + + return session; + } + + function restoreConsole(runtimeMetadata = runtime) { + const sessionMetadata: IRuntimeSessionMetadata = { + sessionId: 'test-console-session-id', + sessionName, + sessionMode: LanguageRuntimeSessionMode.Console, + createdTimestamp: Date.now(), + notebookUri: undefined, + startReason, + }; + return restoreSession(sessionMetadata, runtimeMetadata); + } + + function restoreNotebook(runtimeMetadata = runtime) { + const sessionMetadata: IRuntimeSessionMetadata = { + sessionId: 'test-notebook-session-id', + sessionName, + sessionMode: LanguageRuntimeSessionMode.Notebook, + createdTimestamp: Date.now(), + notebookUri, + startReason, + }; + return restoreSession(sessionMetadata, runtimeMetadata); + } + + async function autoStartSession(runtimeMetadata = runtime) { + const sessionId = await runtimeSessionService.autoStartRuntime(runtimeMetadata, startReason); + assert.ok(sessionId); + const session = runtimeSessionService.getSession(sessionId); + assert.ok(session instanceof TestLanguageRuntimeSession); + disposables.add(session); + return session; + } + + async function selectRuntime(runtimeMetadata = runtime) { + await runtimeSessionService.selectRuntime(runtimeMetadata.runtimeId, startReason); + const session = runtimeSessionService.getConsoleSessionForRuntime(runtimeMetadata.runtimeId); + assert.ok(session instanceof TestLanguageRuntimeSession); + disposables.add(session); + return session; + } + + const data: { action: string; startConsole: IStartSessionTask; startNotebook?: IStartSessionTask }[] = [ + { action: 'start', startConsole: startConsole, startNotebook: startNotebook }, + { action: 'restore', startConsole: restoreConsole, startNotebook: restoreNotebook }, + { action: 'auto start', startConsole: autoStartSession }, + { action: 'select', startConsole: selectRuntime }, + ]; + for (const { action, startConsole, startNotebook } of data) { + + for (const mode of [LanguageRuntimeSessionMode.Console, LanguageRuntimeSessionMode.Notebook]) { + const start = mode === LanguageRuntimeSessionMode.Console ? startConsole : startNotebook; + if (!start) { + continue; + } + + test(`${action} ${mode} returns the expected session`, async () => { + const session = await start(); + + assert.strictEqual(session.getRuntimeState(), RuntimeState.Starting); + assert.strictEqual(session.metadata.sessionName, sessionName); + assert.strictEqual(session.metadata.sessionMode, mode); + assert.strictEqual(session.metadata.startReason, startReason); + assert.strictEqual(session.runtimeMetadata, runtime); + + if (mode === LanguageRuntimeSessionMode.Console) { + assert.strictEqual(session.metadata.notebookUri, undefined); + } else { + assert.strictEqual(session.metadata.notebookUri, notebookUri); + } + }); + + test(`${action} ${mode} sets the expected service state`, async () => { + // Check the initial state. + assertServiceState(); + + const promise = start(); + + // Check the state before awaiting the promise. + assertSingleSessionWillStart(mode); + + const session = await promise; + + // Check the state after awaiting the promise. + assertSingleSessionIsStarting(session); + }); + + // TODO: Should onWillStartSession only fire once? + // It currently fires twice. Before the session is started and when the session + // enters the ready state. + test(`${action} ${mode} fires onWillStartSession`, async () => { + let error: Error | undefined; + const target = sinon.spy(({ session }: IRuntimeSessionWillStartEvent) => { + try { + if (target.callCount > 1) { + return; + } + assert.strictEqual(session.getRuntimeState(), RuntimeState.Uninitialized); + + assertSingleSessionWillStart(mode); + } catch (e) { + error = e; + } + }); + disposables.add(runtimeSessionService.onWillStartSession(target)); + const session = await start(); + + sinon.assert.calledTwice(target); + // When restoring a session, the first event is fired with isNew: false. + sinon.assert.calledWith(target.getCall(0), { isNew: action !== 'restore', session }); + sinon.assert.calledWith(target.getCall(1), { isNew: true, session }); + assert.ifError(error); + }); + + test(`${action} ${mode} fires onDidStartRuntime`, async () => { + let error: Error | undefined; + const target = sinon.stub<[e: ILanguageRuntimeSession]>().callsFake(session => { + try { + assert.strictEqual(session.getRuntimeState(), RuntimeState.Starting); + + assertSingleSessionIsStarting(session); + } catch (e) { + error = e; + } + }); + disposables.add(runtimeSessionService.onDidStartRuntime(target)); + + const session = await start(); + + sinon.assert.calledOnceWithExactly(target, session); + assert.ifError(error); + }); + + test(`${action} ${mode} fires events in order`, async () => { + const willStartSession = sinon.spy(); + disposables.add(runtimeSessionService.onWillStartSession(willStartSession)); + + const didStartRuntime = sinon.spy(); + disposables.add(runtimeSessionService.onDidStartRuntime(didStartRuntime)); + + await start(); + + sinon.assert.callOrder(willStartSession, didStartRuntime); + }); + + if (mode === LanguageRuntimeSessionMode.Console) { + test(`${action} ${mode} sets foregroundSession`, async () => { + const target = sinon.spy(); + disposables.add(runtimeSessionService.onDidChangeForegroundSession(target)); + + const session = await start(); + + assert.strictEqual(runtimeSessionService.foregroundSession, session); + + await waitForRuntimeState(session, RuntimeState.Ready); + + // TODO: Feels a bit surprising that this isn't fired. It's because we set the private + // _foregroundSession property instead of the setter. When the 'ready' state is + // entered, we skip setting foregroundSession because it already matches the session. + sinon.assert.notCalled(target); + }); + } + + if (action === 'start' || action === 'select') { + test(`${action} ${mode} throws for unknown runtime`, async () => { + const runtimeId = 'unknown-runtime-id'; + await assert.rejects( + start({ runtimeId } as ILanguageRuntimeMetadata,), + new Error(`No language runtime with id '${runtimeId}' was found.`), + ); + }); + } + + const createOrRestoreMethod = action === 'restore' ? 'restoreSession' : 'createSession'; + test(`${action} ${mode} encounters ${createOrRestoreMethod}() error`, async () => { + const error = new Error('Failed to create session'); + const stub = sinon.stub(manager, createOrRestoreMethod).rejects(error); + + await assert.rejects(start(), error); + + // If we start now, without createOrRestoreMethod rejecting, it should work. + stub.restore(); + const session = await start(); + + assertSingleSessionIsStarting(session); + }); + + test(`${action} ${mode} encounters session.start() error`, async () => { + // Listen to the onWillStartSession event and stub session.start() to throw an error. + const willStartSession = sinon.spy((e: IRuntimeSessionWillStartEvent) => { + sinon.stub(e.session, 'start').rejects(new Error('Session failed to start')); + }); + const willStartSessionDisposable = runtimeSessionService.onWillStartSession(willStartSession); + + const didFailStartRuntime = sinon.spy(); + disposables.add(runtimeSessionService.onDidFailStartRuntime(didFailStartRuntime)); + + const didStartRuntime = sinon.spy(); + disposables.add(runtimeSessionService.onDidStartRuntime(didStartRuntime)); + + const session1 = await start(); + + assert.strictEqual(session1.getRuntimeState(), RuntimeState.Uninitialized); + + if (mode === LanguageRuntimeSessionMode.Console) { + assertServiceState({ + hasStartingOrRunningConsole: false, + // Note that getConsoleSessionForRuntime includes uninitialized sessions + // but getConsoleSessionForLanguage does not. + consoleSessionForLanguage: undefined, + consoleSessionForRuntime: session1, + activeSessions: [session1], + }); + } else { + assertServiceState({ activeSessions: [session1] }); + } + + sinon.assert.calledOnceWithExactly(didFailStartRuntime, session1); + sinon.assert.callOrder(willStartSession, didFailStartRuntime); + sinon.assert.notCalled(didStartRuntime); + + // If we start now, without session.start() rejecting, it should work. + willStartSessionDisposable.dispose(); + const session2 = await start(); + + assert.strictEqual(session2.getRuntimeState(), RuntimeState.Starting); + + const expectedActiveSessions = action === 'restore' ? + // Restoring a session twice overwrites the previous session in activeSessions. + [session2] : + // Other actions create a new session in activeSessions. + [session1, session2]; + + if (mode === LanguageRuntimeSessionMode.Console) { + assertServiceState({ + hasStartingOrRunningConsole: true, + consoleSession: session2, + consoleSessionForLanguage: session2, + consoleSessionForRuntime: session2, + activeSessions: expectedActiveSessions, + }); + } else { + assertServiceState({ + notebookSession: session2, + notebookSessionForNotebookUri: session2, + activeSessions: expectedActiveSessions, + }); + } + }); + + test(`${action} ${mode} throws if another runtime is starting for the language`, async () => { + let error: Error; + if (mode === LanguageRuntimeSessionMode.Console) { + error = new Error(`Session for language runtime ${formatLanguageRuntimeMetadata(anotherRuntime)} ` + + `cannot be started because language runtime ${formatLanguageRuntimeMetadata(runtime)} ` + + `is already starting for the language.` + + (action !== 'restore' ? ` Request source: ${startReason}` : '')); + } else { + error = new Error(`Session for language runtime ${formatLanguageRuntimeMetadata(anotherRuntime)} cannot ` + + `be started because language runtime ${formatLanguageRuntimeMetadata(runtime)} ` + + `is already starting for the notebook ${notebookUri.toString()}.` + + (action !== 'restore' ? ` Request source: ${startReason}` : '')); + } + + await assert.rejects( + Promise.all([ + start(), + start(anotherRuntime), + ]), + error); + }); + + // Skip for 'select' since selecting another runtime is expected in that case. + if (action !== 'select') { + test(`${action} ${mode} throws if another runtime is running for the language`, async () => { + let error: Error; + if (mode === LanguageRuntimeSessionMode.Console) { + error = new Error(`A console for ${formatLanguageRuntimeMetadata(anotherRuntime)} cannot ` + + `be started because a console for ${formatLanguageRuntimeMetadata(runtime)} ` + + `is already running for the ${runtime.languageName} language.` + + (action !== 'restore' ? ` Request source: ${startReason}` : '')); + } else { + error = new Error(`A notebook for ${formatLanguageRuntimeMetadata(anotherRuntime)} cannot ` + + `be started because a notebook for ${formatLanguageRuntimeMetadata(runtime)} ` + + `is already running for the URI ${notebookUri.toString()}.` + + (action !== 'restore' ? ` Request source: ${startReason}` : '')); + } + + await start(); + await assert.rejects( + start(anotherRuntime), + error, + ); + }); + } + + test(`${action} ${mode} successively`, async () => { + const result1 = await start(); + const result2 = await start(); + const result3 = await start(); + + assert.strictEqual(result1, result2); + assert.strictEqual(result2, result3); + + assertSingleSessionIsStarting(result1); + }); + + test(`${action} ${mode} concurrently`, async () => { + const [result1, result2, result3] = await Promise.all([start(), start(), start()]); + + assert.strictEqual(result1, result2); + assert.strictEqual(result2, result3); + + assertSingleSessionIsStarting(result1); + }); + } + + if (startNotebook) { + test(`${action} console and notebook from the same runtime concurrently`, async () => { + // Consoles and notebooks shouldn't interfere with each other, even for the same runtime. + const [consoleSession, notebookSession] = await Promise.all([ + startConsole(), + startNotebook(), + ]); + + assert.strictEqual(consoleSession.getRuntimeState(), RuntimeState.Starting); + assert.strictEqual(notebookSession.getRuntimeState(), RuntimeState.Starting); + + assertServiceState({ + hasStartingOrRunningConsole: true, + consoleSession, + consoleSessionForLanguage: consoleSession, + consoleSessionForRuntime: consoleSession, + notebookSession, + notebookSessionForNotebookUri: notebookSession, + activeSessions: [consoleSession, notebookSession], + }); + }); + } + } + + test(`start notebook without notebook uri`, async () => { + await assert.rejects( + startSession(undefined, LanguageRuntimeSessionMode.Notebook, undefined), + new Error('A notebook URI must be provided when starting a notebook session.'), + ); + }); + + test('restore console registers runtime if unregistered', async () => { + // The runtime should not yet be registered. + assert.strictEqual(languageRuntimeService.getRegisteredRuntime(unregisteredRuntime.runtimeId), undefined); + + await restoreConsole(unregisteredRuntime); + + // The runtime should now be registered. + assert.strictEqual(languageRuntimeService.getRegisteredRuntime(unregisteredRuntime.runtimeId), unregisteredRuntime); + }); + + test('auto start validates runtime if unregistered', async () => { + // The runtime should not yet be registered. + assert.strictEqual(languageRuntimeService.getRegisteredRuntime(unregisteredRuntime.runtimeId), undefined); + + // Update the validator to add extra runtime data. + const validatedMetadata: Partial = { + extraRuntimeData: { someNewKey: 'someNewValue' } + }; + manager.setValidateMetadata(async (metadata: ILanguageRuntimeMetadata) => { + return { ...metadata, ...validatedMetadata }; + }); + + await autoStartSession(unregisteredRuntime); + + // The validated metadata should now be registered. + assert.deepStrictEqual( + languageRuntimeService.getRegisteredRuntime(unregisteredRuntime.runtimeId), + { ...unregisteredRuntime, ...validatedMetadata } + ); + }); + + test('auto start throws if runtime validation errors', async () => { + // The runtime should not yet be registered. + assert.strictEqual(languageRuntimeService.getRegisteredRuntime(unregisteredRuntime.runtimeId), undefined); + + // Update the validator to throw. + const error = new Error('Failed to validate runtime metadata'); + manager.setValidateMetadata(async (_metadata: ILanguageRuntimeMetadata) => { + throw error; + }); + + await assert.rejects(autoStartSession(unregisteredRuntime), error); + + // The runtime should remain unregistered. + assert.strictEqual(languageRuntimeService.getRegisteredRuntime(unregisteredRuntime.runtimeId), undefined); + }); + + test('auto start console does nothing if automatic startup is disabled', async () => { + configService.setUserConfiguration('positron.interpreters.automaticStartup', false); + + const sessionId = await runtimeSessionService.autoStartRuntime(runtime, startReason); + + assert.strictEqual(sessionId, ''); + assertServiceState(); + }); + + for (const action of ['auto start', 'start']) { + test(`${action} console in an untrusted workspace defers until trust is granted`, async () => { + workspaceTrustManagementService.setWorkspaceTrust(false); + + let sessionId: string; + if (action === 'auto start') { + sessionId = await runtimeSessionService.autoStartRuntime(runtime, startReason); + } else { + sessionId = await runtimeSessionService.startNewRuntimeSession( + runtime.runtimeId, sessionName, LanguageRuntimeSessionMode.Console, undefined, startReason); + } + + assert.strictEqual(sessionId, ''); + assertServiceState(); + + workspaceTrustManagementService.setWorkspaceTrust(true); + + // The session should eventually start. + const session = await Event.toPromise(runtimeSessionService.onDidStartRuntime); + disposables.add(session); + + assertSingleSessionIsStarting(session); + }); + } + + test('start notebook in an untrusted workspace throws', async () => { + workspaceTrustManagementService.setWorkspaceTrust(false); + + await assert.rejects(startNotebook(), new Error('Cannot start a notebook session in an untrusted workspace.')); + }); + + test('select console while another runtime is running for the language', async () => { + const session1 = await startConsole(anotherRuntime); + await waitForRuntimeState(session1, RuntimeState.Ready); + const session2 = await selectRuntime(); + + assert.strictEqual(session1.getRuntimeState(), RuntimeState.Exited); + assert.strictEqual(session2.getRuntimeState(), RuntimeState.Starting); + + assertServiceState({ + hasStartingOrRunningConsole: true, + consoleSession: session2, + consoleSessionForLanguage: session2, + consoleSessionForRuntime: session2, + activeSessions: [session1, session2], + }); + }); + + test('select console throws if session is still starting', async () => { + await startConsole(anotherRuntime); + await assert.rejects( + selectRuntime(), + new Error('Cannot shut down kernel; it is not (yet) running. (state = starting)'), + ); + }); + + test('select console to the same runtime sets the foreground session', async () => { + const session1 = await startConsole(); + + runtimeSessionService.foregroundSession = undefined; + + const session2 = await selectRuntime(); + + assert.strictEqual(session1, session2); + assert.strictEqual(runtimeSessionService.foregroundSession, session1); + }); + + test(`select console to another runtime and first session never fires onDidEndSession`, async () => { + const session = await startConsole(); + await waitForRuntimeState(session, RuntimeState.Ready); + + // Stub onDidEndSession to never fire, causing the shutdown to time out. + sinon.stub(session, 'onDidEndSession').returns({ dispose: () => { } }); + + // Use a fake timer to avoid actually having to wait for the timeout. + const clock = sinon.useFakeTimers(); + const promise = assert.rejects(selectRuntime(anotherRuntime), new Error(`Timed out waiting for runtime ` + + `${formatLanguageRuntimeSession(session)} to finish exiting.`)); + await clock.tickAsync(10_000); + await promise; + }); + + test(`select console to another runtime encounters session.shutdown() error`, async () => { + const session = await startConsole(); + + // Stub session.shutdown() to throw an error. + const error = new Error('Session failed to shut down'); + sinon.stub(session, 'shutdown').rejects(error); + + // We also want to ensure that the timeout is not hit in this case but don't want to + // actually wait, so we use a fake timer. + const clock = sinon.useFakeTimers(); + await assert.rejects(selectRuntime(anotherRuntime), error); + await clock.tickAsync(10_000); + }); + + function restartSession(sessionId: string) { + return runtimeSessionService.restartSession(sessionId, startReason); + } + + for (const { mode, start } of [ + { mode: LanguageRuntimeSessionMode.Console, start: startConsole }, + { mode: LanguageRuntimeSessionMode.Notebook, start: startNotebook }, + ]) { + test(`restart ${mode} throws if session not found`, async () => { + const sessionId = 'unknown-session-id'; + assert.rejects( + restartSession(sessionId), + new Error(`No session with ID '${sessionId}' was found.`), + ); + }); + + for (const state of [RuntimeState.Busy, RuntimeState.Idle, RuntimeState.Ready]) { + test(`restart ${mode} in '${state}' state`, async () => { + // Start the session and wait for it to be ready. + const session = await start(); + await waitForRuntimeState(session, RuntimeState.Ready); + + // Set the state to the desired state. + if (session.getRuntimeState() !== state) { + session.setRuntimeState(state); + } + + await restartSession(session.sessionId); + + assert.strictEqual(session.getRuntimeState(), RuntimeState.Restarting); + assertSingleSessionIsRestarting(session); + + await waitForRuntimeState(session, RuntimeState.Ready); + assertSingleSessionIsReady(session); + }); + } + + for (const state of [RuntimeState.Uninitialized, RuntimeState.Exited]) { + test(`restart ${mode} in '${state}' state`, async () => { + // Get a session to the exited state. + const session = await start(); + await waitForRuntimeState(session, RuntimeState.Ready); + await session.shutdown(RuntimeExitReason.Shutdown); + await waitForRuntimeState(session, RuntimeState.Exited); + + await restartSession(session.sessionId); + + // The existing sessino should remain exited. + assert.strictEqual(session.getRuntimeState(), RuntimeState.Exited); + + // A new session should be starting. + let newSession: ILanguageRuntimeSession | undefined; + if (mode === LanguageRuntimeSessionMode.Console) { + newSession = runtimeSessionService.getConsoleSessionForRuntime(runtime.runtimeId); + } else { + newSession = runtimeSessionService.getNotebookSessionForNotebookUri(notebookUri); + } + assert.ok(newSession); + disposables.add(newSession); + + assert.strictEqual(newSession.getRuntimeState(), RuntimeState.Starting); + assert.strictEqual(newSession.metadata.sessionName, session.metadata.sessionName); + assert.strictEqual(newSession.metadata.sessionMode, session.metadata.sessionMode); + assert.strictEqual(newSession.metadata.notebookUri, session.metadata.notebookUri); + assert.strictEqual(newSession.runtimeMetadata, session.runtimeMetadata); + + if (mode === LanguageRuntimeSessionMode.Console) { + assertServiceState({ + hasStartingOrRunningConsole: true, + consoleSession: newSession, + consoleSessionForLanguage: newSession, + consoleSessionForRuntime: newSession, + activeSessions: [session, newSession], + }); + } else { + assertServiceState({ + notebookSession: newSession, + notebookSessionForNotebookUri: newSession, + activeSessions: [session, newSession], + }); + } + }); + } + + test(`restart ${mode} in 'starting' state`, async () => { + const session = await start(); + assert.strictEqual(session.getRuntimeState(), RuntimeState.Starting); + + await restartSession(session.sessionId); + + assertSingleSessionIsStarting(session); + }); + + test(`restart ${mode} in 'restarting' state`, async () => { + const session = await start(); + await waitForRuntimeState(session, RuntimeState.Ready); + + session.restart(); + assert.strictEqual(session.getRuntimeState(), RuntimeState.Restarting); + + const target = sinon.spy(session, 'restart'); + + await restartSession(session.sessionId); + + assertSingleSessionIsRestarting(session); + + sinon.assert.notCalled(target); + }); + + test(`restart ${mode} concurrently`, async () => { + const session = await start(); + await waitForRuntimeState(session, RuntimeState.Ready); + + const target = sinon.spy(session, 'restart'); + + await Promise.all([ + restartSession(session.sessionId), + restartSession(session.sessionId), + restartSession(session.sessionId), + ]); + + assertSingleSessionIsRestarting(session); + + sinon.assert.calledOnce(target); + }); + + test(`restart ${mode} successively`, async () => { + const session = await start(); + + const target = sinon.spy(session, 'restart'); + + await waitForRuntimeState(session, RuntimeState.Ready); + await restartSession(session.sessionId); + await waitForRuntimeState(session, RuntimeState.Ready); + await restartSession(session.sessionId); + await waitForRuntimeState(session, RuntimeState.Ready); + await restartSession(session.sessionId); + + assertSingleSessionIsRestarting(session); + + sinon.assert.calledThrice(target); + }); + + test(`restart ${mode} while ready -> start`, async () => { + const session = await start(); + await waitForRuntimeState(session, RuntimeState.Ready); + + await restartSession(session.sessionId); + await waitForRuntimeState(session, RuntimeState.Ready); + + const newSession = await start(); + + assertSingleSessionIsReady(newSession); + }); + } +}); diff --git a/src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts b/src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts index 863a5fdea46..7363d755cb0 100644 --- a/src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts +++ b/src/vs/workbench/services/runtimeSession/test/common/testLanguageRuntimeSession.ts @@ -11,6 +11,7 @@ import { ILanguageRuntimeSession, IRuntimeClientInstance, IRuntimeSessionMetadat import { ILanguageRuntimeClientCreatedEvent, ILanguageRuntimeExit, ILanguageRuntimeInfo, ILanguageRuntimeMessage, ILanguageRuntimeMessageClearOutput, ILanguageRuntimeMessageError, ILanguageRuntimeMessageInput, ILanguageRuntimeMessageIPyWidget, ILanguageRuntimeMessageOutput, ILanguageRuntimeMessagePrompt, ILanguageRuntimeMessageResult, ILanguageRuntimeMessageState, ILanguageRuntimeMessageStream, ILanguageRuntimeMetadata, ILanguageRuntimeStartupFailure, LanguageRuntimeMessageType, RuntimeCodeExecutionMode, RuntimeCodeFragmentStatus, RuntimeErrorBehavior, RuntimeExitReason, RuntimeOnlineState, RuntimeOutputKind, RuntimeState } from 'vs/workbench/services/languageRuntime/common/languageRuntimeService'; import { IRuntimeClientEvent } from 'vs/workbench/services/languageRuntime/common/languageRuntimeUiClient'; import { TestRuntimeClientInstance } from 'vs/workbench/services/languageRuntime/test/common/testRuntimeClientInstance'; +import { CancellationError } from 'vs/base/common/errors'; export class TestLanguageRuntimeSession extends Disposable implements ILanguageRuntimeSession { private readonly _onDidChangeRuntimeState = this._register(new Emitter()); @@ -166,10 +167,24 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR async restart(): Promise { await this.shutdown(RuntimeExitReason.Restart); - await this.start(); + + // Wait for the session to exit, then start it again. + const disposable = this._register(this.onDidChangeRuntimeState(state => { + if (state === RuntimeState.Exited) { + disposable.dispose(); + this.start(); + } + })); } async shutdown(exitReason: RuntimeExitReason): Promise { + if (this._currentState !== RuntimeState.Idle && + this._currentState !== RuntimeState.Busy && + this._currentState !== RuntimeState.Ready) { + throw new Error('Cannot shut down kernel; it is not (yet) running.' + + ` (state = ${this._currentState})`); + } + if (exitReason === RuntimeExitReason.Restart) { this._onDidChangeRuntimeState.fire(RuntimeState.Restarting); } else { @@ -200,10 +215,6 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR throw new Error('Not implemented.'); } - override dispose() { - super.dispose(); - } - // Test helpers setRuntimeState(state: RuntimeState) { @@ -366,3 +377,24 @@ export class TestLanguageRuntimeSession extends Disposable implements ILanguageR }); } } + +export async function waitForRuntimeState( + session: ILanguageRuntimeSession, + state: RuntimeState, + timeout = 10_000, +) { + return new Promise((resolve, reject) => { + const timer = setTimeout(() => { + disposable.dispose(); + reject(new CancellationError()); + }, timeout); + + const disposable = session.onDidChangeRuntimeState(newState => { + if (newState === state) { + clearTimeout(timer); + disposable.dispose(); + resolve(); + } + }); + }); +} diff --git a/src/vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService.ts b/src/vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService.ts index 9ffc0958465..d6e6c187b0f 100644 --- a/src/vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService.ts +++ b/src/vs/workbench/services/runtimeSession/test/common/testRuntimeSessionService.ts @@ -76,7 +76,7 @@ export function createTestLanguageRuntimeMetadata( disposables.add(languageRuntimeService.registerRuntime(runtime)); // Register the test runtime manager. - const manager = new TestRuntimeSessionManager(); + const manager = TestRuntimeSessionManager.instance; disposables.add(runtimeSessionService.registerSessionManager(manager)); return runtime; diff --git a/src/vs/workbench/test/common/positronWorkbenchTestServices.ts b/src/vs/workbench/test/common/positronWorkbenchTestServices.ts index 32923d7e7d3..028f06ec31e 100644 --- a/src/vs/workbench/test/common/positronWorkbenchTestServices.ts +++ b/src/vs/workbench/test/common/positronWorkbenchTestServices.ts @@ -107,6 +107,10 @@ export class TestPositronModalDialogService implements IPositronModalDialogsServ } } export class TestRuntimeSessionManager implements ILanguageRuntimeSessionManager { + public static readonly instance = new TestRuntimeSessionManager(); + + private _validateMetadata?: (metadata: ILanguageRuntimeMetadata) => Promise; + async managesRuntime(runtime: ILanguageRuntimeMetadata): Promise { return true; } @@ -119,7 +123,14 @@ export class TestRuntimeSessionManager implements ILanguageRuntimeSessionManager return new TestLanguageRuntimeSession(sessionMetadata, runtimeMetadata); } - validateMetadata(metadata: ILanguageRuntimeMetadata): Promise { - throw new Error('Method not implemented'); + async validateMetadata(metadata: ILanguageRuntimeMetadata): Promise { + if (this._validateMetadata) { + return this._validateMetadata(metadata); + } + return metadata; + } + + setValidateMetadata(handler: (metadata: ILanguageRuntimeMetadata) => Promise): void { + this._validateMetadata = handler; } }