Simplify starting how we start local Jupyter, remote Jupyter and raw kernels

src/kernels/common/baseJupyterSession.ts

@@ -96,13 +96,17 @@ export abstract class BaseJupyterSession implements IBaseKernelConnectionSession
     private _wrappedKernel?: KernelConnectionWrapper;
     private _isDisposed?: boolean;
     private readonly _disposed = new EventEmitter<void>();
+    private readonly didShutdown = new EventEmitter<void>();
     protected readonly disposables: IDisposable[] = [];
     public get disposed() {
         return this._isDisposed === true;
     }
     public get onDidDispose() {
         return this._disposed.event;
     }
+    public get onDidShutdown() {
+        return this.didShutdown.event;
+    }
     protected get session(): ISessionWithSocket | undefined {
         return this._session;
     }
@@ -525,6 +529,8 @@ export abstract class BaseJupyterSession implements IBaseKernelConnectionSession
             this.restartSessionPromise = undefined;
             this.onStatusChangedEvent.fire('dead');
             this._disposed.fire();
+            this.didShutdown.fire();
+            this.didShutdown.dispose();
             this._disposed.dispose();
             this.onStatusChangedEvent.dispose();
             this.previousAnyMessageHandler?.dispose();

src/kernels/jupyter/finder/remoteKernelFinder.ts

@@ -6,7 +6,6 @@ import { getKernelId } from '../../helpers';
 import {
     BaseKernelConnectionMetadata,
     IJupyterKernelSpec,
-    IJupyterServerConnector,
     IKernelProvider,
     IJupyterConnection,
     isRemoteConnection,
@@ -27,7 +26,6 @@ import { traceError, traceWarning, traceInfoIfCI, traceVerbose } from '../../../
 import { IPythonExtensionChecker } from '../../../platform/api/types';
 import { computeServerId } from '../jupyterUtils';
 import { createPromiseFromCancellation } from '../../../platform/common/cancellation';
-import { DisplayOptions } from '../../displayOptions';
 import { isArray } from '../../../platform/common/utils/sysTypes';
 import { areObjectsWithUrisTheSame, noop } from '../../../platform/common/utils/misc';
 import { IApplicationEnvironment } from '../../../platform/common/application/types';
@@ -37,6 +35,10 @@ import { ContributedKernelFinderKind } from '../../internalTypes';
 import { disposeAllDisposables } from '../../../platform/common/helpers';
 import { PromiseMonitor } from '../../../platform/common/utils/promises';
 import { getDisplayPath } from '../../../platform/common/platform/fs-paths';
+import { JupyterConnection } from '../connection/jupyterConnection';
+import { KernelProgressReporter } from '../../../platform/progress/kernelProgressReporter';
+import { DataScience } from '../../../platform/common/utils/localize';
+import { isUnitTestExecution } from '../../../platform/common/constants';
 
 // Even after shutting down a kernel, the server API still returns the old information.
 // Re-query after 2 seconds to ensure we don't get stale information.
@@ -90,14 +92,14 @@ export class RemoteKernelFinder implements IRemoteKernelFinder, IDisposable {
         readonly cacheKey: string,
         private jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
         private extensionChecker: IPythonExtensionChecker,
-        private readonly jupyterServerConnector: IJupyterServerConnector,
         private readonly globalState: Memento,
         private readonly env: IApplicationEnvironment,
         private readonly cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator,
         kernelFinder: KernelFinder,
         private readonly kernelProvider: IKernelProvider,
         private readonly extensions: IExtensions,
-        readonly serverUri: IJupyterServerUriEntry
+        readonly serverUri: IJupyterServerUriEntry,
+        private readonly jupyterConnection: JupyterConnection
     ) {
         // When we register, add a disposable to clean ourselves up from the main kernel finder list
         // Unlike the Local kernel finder universal remote kernel finders will be added on the fly
@@ -200,7 +202,7 @@ export class RemoteKernelFinder implements IRemoteKernelFinder, IDisposable {
             } else {
                 try {
                     const kernelsWithoutCachePromise = (async () => {
-                        const connInfo = await this.getRemoteConnectionInfo(undefined, displayProgress);
+                        const connInfo = await this.getRemoteConnectionInfo(displayProgress);
                         return connInfo ? this.listKernelsFromConnection(connInfo) : Promise.resolve([]);
                     })();
 
@@ -227,7 +229,7 @@ export class RemoteKernelFinder implements IRemoteKernelFinder, IDisposable {
 
             try {
                 const kernelsWithoutCachePromise = (async () => {
-                    const connInfo = await this.getRemoteConnectionInfo(updateCacheCancellationToken.token, false);
+                    const connInfo = await this.getRemoteConnectionInfo(false);
                     return connInfo ? this.listKernelsFromConnection(connInfo) : Promise.resolve([]);
                 })();
 
@@ -260,18 +262,16 @@ export class RemoteKernelFinder implements IRemoteKernelFinder, IDisposable {
         return this.cache;
     }
 
-    private async getRemoteConnectionInfo(
-        cancelToken?: CancellationToken,
-        displayProgress: boolean = true
-    ): Promise<IJupyterConnection | undefined> {
-        const ui = new DisplayOptions(!displayProgress);
-        return this.jupyterServerConnector.connect({
-            resource: undefined,
-            ui,
-            localJupyter: false,
-            token: cancelToken,
-            serverId: this.serverUri.serverId
-        });
+    private async getRemoteConnectionInfo(displayProgress: boolean = true): Promise<IJupyterConnection | undefined> {
+        const disposables: IDisposable[] = [];
+        if (displayProgress) {
+            disposables.push(KernelProgressReporter.createProgressReporter(undefined, DataScience.connectingToJupyter));
+        }
+        return this.jupyterConnection
+            .createConnectionInfo({
+                serverId: this.serverUri.serverId
+            })
+            .finally(() => disposeAllDisposables(disposables));
     }
 
     private async getFromCache(cancelToken?: CancellationToken): Promise<RemoteKernelConnectionMetadata[]> {
@@ -435,19 +435,30 @@ export class RemoteKernelFinder implements IRemoteKernelFinder, IDisposable {
                     clearTimeout(this.cacheLoggingTimeout);
                 }
                 // Reduce the logging, as this can get written a lot,
-                this.cacheLoggingTimeout = setTimeout(() => {
-                    traceVerbose(
-                        `Updating cache with Remote kernels ${values
-                            .map((k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`)
-                            .join(', ')}, Added = ${added
-                            .map((k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`)
-                            .join(', ')}, Updated = ${updated
-                            .map((k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`)
-                            .join(', ')}, Removed = ${removed
-                            .map((k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`)
-                            .join(', ')}`
-                    );
-                }, 15_000);
+                this.cacheLoggingTimeout = setTimeout(
+                    () => {
+                        traceVerbose(
+                            `Updating cache with Remote kernels ${values
+                                .map(
+                                    (k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`
+                                )
+                                .join(', ')}, Added = ${added
+                                .map(
+                                    (k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`
+                                )
+                                .join(', ')}, Updated = ${updated
+                                .map(
+                                    (k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`
+                                )
+                                .join(', ')}, Removed = ${removed
+                                .map(
+                                    (k) => `${k.kind}:'${k.id} (interpreter id = ${getDisplayPath(k.interpreter?.id)})'`
+                                )
+                                .join(', ')}`
+                        );
+                    },
+                    isUnitTestExecution() ? 0 : 15_000
+                );
                 this.disposables.push(new Disposable(() => clearTimeout(this.cacheLoggingTimeout)));
             }
         } catch (ex) {

src/kernels/jupyter/finder/remoteKernelFinder.unit.test.ts

@@ -33,6 +33,8 @@ import { RemoteKernelSpecsCacheKey } from '../../common/commonFinder';
 import { IExtensions } from '../../../platform/common/types';
 import { createEventHandler, TestEventHandler } from '../../../test/common';
 import { RemoteKernelFinder } from './remoteKernelFinder';
+import { JupyterConnection } from '../connection/jupyterConnection';
+import { disposeAllDisposables } from '../../../platform/common/helpers';
 
 suite(`Remote Kernel Finder`, () => {
     let disposables: Disposable[] = [];
@@ -44,6 +46,7 @@ suite(`Remote Kernel Finder`, () => {
     const dummyEvent = new EventEmitter<number>();
     let cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator;
     let kernelsChanged: TestEventHandler<void>;
+    let jupyterConnection: JupyterConnection;
     const connInfo: IJupyterConnection = {
         url: 'http://foobar',
         type: 'jupyter',
@@ -151,27 +154,26 @@ suite(`Remote Kernel Finder`, () => {
         kernelFinder = new KernelFinder(disposables);
         kernelsChanged = createEventHandler(kernelFinder, 'onDidChangeKernels');
         disposables.push(kernelsChanged);
-
+        jupyterConnection = mock<JupyterConnection>();
+        when(jupyterConnection.createConnectionInfo(anything())).thenResolve(connInfo);
         remoteKernelFinder = new RemoteKernelFinder(
             'currentremote',
             'Local Kernels',
             RemoteKernelSpecsCacheKey,
             instance(jupyterSessionManagerFactory),
             instance(extensionChecker),
-            instance(serverConnector),
             instance(memento),
             instance(env),
             instance(cachedRemoteKernelValidator),
             kernelFinder,
             instance(kernelProvider),
             instance(extensions),
-            serverEntry
+            serverEntry,
+            instance(jupyterConnection)
         );
         remoteKernelFinder.activate().then(noop, noop);
     });
-    teardown(() => {
-        disposables.forEach((d) => d.dispose());
-    });
+    teardown(() => disposeAllDisposables(disposables));
     test('Kernels found', async () => {
         when(jupyterSessionManager.getRunningKernels()).thenResolve([]);
         when(jupyterSessionManager.getRunningSessions()).thenResolve([]);

src/kernels/jupyter/finder/remoteKernelFinderController.ts

@@ -3,7 +3,7 @@
 
 import { injectable, inject, named } from 'inversify';
 import { Disposable, Memento } from 'vscode';
-import { IKernelFinder, IKernelProvider, IJupyterServerConnector } from '../../types';
+import { IKernelFinder, IKernelProvider } from '../../types';
 import { GLOBAL_MEMENTO, IDisposableRegistry, IExtensions, IMemento } from '../../../platform/common/types';
 import {
     IJupyterSessionManagerFactory,
@@ -21,6 +21,7 @@ import { ContributedKernelFinderKind } from '../../internalTypes';
 import * as localize from '../../../platform/common/utils/localize';
 import { RemoteKernelSpecsCacheKey } from '../../common/commonFinder';
 import { Settings } from '../../../platform/common/constants';
+import { JupyterConnection } from '../connection/jupyterConnection';
 
 /** Strategy design */
 interface IRemoteKernelFinderRegistrationStrategy {
@@ -34,15 +35,15 @@ class MultiServerStrategy implements IRemoteKernelFinderRegistrationStrategy {
     constructor(
         private jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
         private extensionChecker: IPythonExtensionChecker,
-        private readonly jupyterServerConnector: IJupyterServerConnector,
         private readonly serverUriStorage: IJupyterServerUriStorage,
         private readonly globalState: Memento,
         private readonly env: IApplicationEnvironment,
         private readonly cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator,
         private readonly kernelFinder: KernelFinder,
         private readonly disposables: IDisposableRegistry,
         private readonly kernelProvider: IKernelProvider,
-        private readonly extensions: IExtensions
+        private readonly extensions: IExtensions,
+        private readonly jupyterConnection: JupyterConnection
     ) {}
 
     async activate(): Promise<void> {
@@ -75,14 +76,14 @@ class MultiServerStrategy implements IRemoteKernelFinderRegistrationStrategy {
                 `${RemoteKernelSpecsCacheKey}-${serverUri.serverId}`,
                 this.jupyterSessionManagerFactory,
                 this.extensionChecker,
-                this.jupyterServerConnector,
                 this.globalState,
                 this.env,
                 this.cachedRemoteKernelValidator,
                 this.kernelFinder,
                 this.kernelProvider,
                 this.extensions,
-                serverUri
+                serverUri,
+                this.jupyterConnection
             );
             this.disposables.push(finder);
 
@@ -116,28 +117,28 @@ export class RemoteKernelFinderController implements IExtensionSyncActivationSer
         @inject(IJupyterSessionManagerFactory)
         private readonly jupyterSessionManagerFactory: IJupyterSessionManagerFactory,
         @inject(IPythonExtensionChecker) private readonly extensionChecker: IPythonExtensionChecker,
-        @inject(IJupyterServerConnector) private readonly jupyterServerConnector: IJupyterServerConnector,
         @inject(IJupyterServerUriStorage) private readonly serverUriStorage: IJupyterServerUriStorage,
         @inject(IMemento) @named(GLOBAL_MEMENTO) private readonly globalState: Memento,
         @inject(IApplicationEnvironment) private readonly env: IApplicationEnvironment,
         @inject(IJupyterRemoteCachedKernelValidator)
         private readonly cachedRemoteKernelValidator: IJupyterRemoteCachedKernelValidator,
         @inject(IKernelFinder) private readonly kernelFinder: KernelFinder,
         @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider,
-        @inject(IExtensions) private readonly extensions: IExtensions
+        @inject(IExtensions) private readonly extensions: IExtensions,
+        @inject(JupyterConnection) jupyterConnection: JupyterConnection
     ) {
         this._strategy = new MultiServerStrategy(
             this.jupyterSessionManagerFactory,
             this.extensionChecker,
-            this.jupyterServerConnector,
             this.serverUriStorage,
             this.globalState,
             this.env,
             this.cachedRemoteKernelValidator,
             this.kernelFinder,
             this._localDisposables,
             this.kernelProvider,
-            this.extensions
+            this.extensions,
+            jupyterConnection
         );
     }
 

src/kernels/jupyter/launcher/jupyterKernelConnectionSessionCreator.ts

@@ -0,0 +1,75 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import { Uri } from 'vscode';
+import { Cancellation } from '../../../platform/common/cancellation';
+import {
+    IJupyterKernelConnectionSession,
+    KernelConnectionSessionCreationOptions,
+    isLocalConnection,
+    isRemoteConnection
+} from '../../types';
+import { IJupyterSessionManager } from '../types';
+import { traceError, traceInfo } from '../../../platform/logging';
+import { IWorkspaceService } from '../../../platform/common/application/types';
+import { inject, injectable } from 'inversify';
+import { noop } from '../../../platform/common/utils/misc';
+import { SessionDisposedError } from '../../../platform/errors/sessionDisposedError';
+import { RemoteJupyterServerConnectionError } from '../../../platform/errors/remoteJupyterServerConnectionError';
+
+export type JupyterKernelConnectionSessionCreationOptions = KernelConnectionSessionCreationOptions & {
+    sessionManager: IJupyterSessionManager;
+};
+
+@injectable()
+export class JupyterKernelConnectionSessionCreator {
+    constructor(@inject(IWorkspaceService) private readonly workspaceService: IWorkspaceService) {}
+    public async create(
+        options: JupyterKernelConnectionSessionCreationOptions
+    ): Promise<IJupyterKernelConnectionSession> {
+        if (options.sessionManager.isDisposed) {
+            throw new SessionDisposedError();
+        }
+        if (isRemoteConnection(options.kernelConnection)) {
+            try {
+                await Promise.all([
+                    options.sessionManager.getRunningKernels(),
+                    options.sessionManager.getKernelSpecs()
+                ]);
+            } catch (ex) {
+                traceError(
+                    'Failed to fetch running kernels from remote server, connection may be outdated or remote server may be unreachable',
+                    ex
+                );
+                throw new RemoteJupyterServerConnectionError(
+                    options.kernelConnection.baseUrl,
+                    options.kernelConnection.serverId,
+                    ex
+                );
+            }
+        }
+
+        Cancellation.throwIfCanceled(options.token);
+        // Figure out the working directory we need for our new notebook. This is only necessary for local.
+        const workingDirectory = isLocalConnection(options.kernelConnection)
+            ? await this.workspaceService.computeWorkingDirectory(options.resource)
+            : '';
+        Cancellation.throwIfCanceled(options.token);
+        // Start a session (or use the existing one if allowed)
+        const session = await options.sessionManager.startNew(
+            options.resource,
+            options.kernelConnection,
+            Uri.file(workingDirectory),
+            options.ui,
+            options.token,
+            options.creator
+        );
+        if (options.token.isCancellationRequested) {
+            // Even if this is a remote kernel, we should shut this down as it's not needed.
+            session.shutdown().catch(noop);
+        }
+        Cancellation.throwIfCanceled(options.token);
+        traceInfo(`Started session for kernel ${options.kernelConnection.kind}:${options.kernelConnection.id}`);
+        return session;
+    }
+}