Expose client creation API for pylance

src/client/activation/languageClientMiddleware.ts

@@ -42,7 +42,7 @@ export class LanguageClientMiddleware extends LanguageClientMiddlewareBase {
         );
     }
 
-    protected shouldCreateHidingMiddleware(jupyterDependencyManager: IJupyterExtensionDependencyManager): boolean {
+    private shouldCreateHidingMiddleware(jupyterDependencyManager: IJupyterExtensionDependencyManager): boolean {
         return jupyterDependencyManager && jupyterDependencyManager.isJupyterExtensionInstalled;
     }
 

src/client/activation/node/analysisOptions.ts

@@ -10,7 +10,6 @@ import { IExperimentService } from '../../common/types';
 
 import { LanguageServerAnalysisOptionsBase } from '../common/analysisOptions';
 import { ILanguageServerOutputChannel } from '../types';
-import { LspNotebooksExperiment } from './lspNotebooksExperiment';
 import { traceWarn } from '../../logging';
 
 const EDITOR_CONFIG_SECTION = 'editor';
@@ -22,7 +21,6 @@ export class NodeLanguageServerAnalysisOptions extends LanguageServerAnalysisOpt
         lsOutputChannel: ILanguageServerOutputChannel,
         workspace: IWorkspaceService,
         private readonly experimentService: IExperimentService,
-        private readonly lspNotebooksExperiment: LspNotebooksExperiment,
     ) {
         super(lsOutputChannel, workspace);
     }
@@ -36,8 +34,6 @@ export class NodeLanguageServerAnalysisOptions extends LanguageServerAnalysisOpt
         return ({
             experimentationSupport: true,
             trustedWorkspaceSupport: true,
-            lspNotebooksSupport: this.lspNotebooksExperiment.isInNotebooksExperiment(),
-            lspInteractiveWindowSupport: this.lspNotebooksExperiment.isInNotebooksExperimentWithInteractiveWindowSupport(),
             autoIndentSupport: await this.isAutoIndentEnabled(),
         } as unknown) as LanguageClientOptions;
     }

src/client/activation/node/languageClientFactory.ts

@@ -11,7 +11,7 @@ import { PythonEnvironment } from '../../pythonEnvironments/info';
 import { FileBasedCancellationStrategy } from '../common/cancellationUtils';
 import { ILanguageClientFactory } from '../types';
 
-const languageClientName = 'Pylance';
+export const PYLANCE_NAME = 'Pylance';
 
 export class NodeLanguageClientFactory implements ILanguageClientFactory {
     constructor(private readonly fs: IFileSystem, private readonly extensions: IExtensions) {}
@@ -50,6 +50,6 @@ export class NodeLanguageClientFactory implements ILanguageClientFactory {
             },
         };
 
-        return new LanguageClient(PYTHON_LANGUAGE, languageClientName, serverOptions, clientOptions);
+        return new LanguageClient(PYTHON_LANGUAGE, PYLANCE_NAME, serverOptions, clientOptions);
     }
 }

src/client/activation/node/languageClientMiddleware.ts

@@ -36,40 +36,29 @@ export class NodeLanguageClientMiddleware extends LanguageClientMiddleware {
         this.jupyterExtensionIntegration = serviceContainer.get<JupyterExtensionIntegration>(
             JupyterExtensionIntegration,
         );
-        if (!this.notebookAddon && this.lspNotebooksExperiment.isInNotebooksExperimentWithInteractiveWindowSupport()) {
+        if (!this.notebookAddon) {
             this.notebookAddon = new LspInteractiveWindowMiddlewareAddon(
                 this.getClient,
                 this.jupyterExtensionIntegration,
             );
         }
     }
 
-    protected shouldCreateHidingMiddleware(jupyterDependencyManager: IJupyterExtensionDependencyManager): boolean {
-        return (
-            super.shouldCreateHidingMiddleware(jupyterDependencyManager) &&
-            !this.lspNotebooksExperiment.isInNotebooksExperiment()
-        );
-    }
-
     protected async onExtensionChange(jupyterDependencyManager: IJupyterExtensionDependencyManager): Promise<void> {
         if (jupyterDependencyManager && jupyterDependencyManager.isJupyterExtensionInstalled) {
             await this.lspNotebooksExperiment.onJupyterInstalled();
         }
 
-        if (this.lspNotebooksExperiment.isInNotebooksExperimentWithInteractiveWindowSupport()) {
-            if (!this.notebookAddon) {
-                this.notebookAddon = new LspInteractiveWindowMiddlewareAddon(
-                    this.getClient,
-                    this.jupyterExtensionIntegration,
-                );
-            }
-        } else {
-            super.onExtensionChange(jupyterDependencyManager);
+        if (!this.notebookAddon) {
+            this.notebookAddon = new LspInteractiveWindowMiddlewareAddon(
+                this.getClient,
+                this.jupyterExtensionIntegration,
+            );
         }
     }
 
     protected async getPythonPathOverride(uri: Uri | undefined): Promise<string | undefined> {
-        if (!uri || !this.lspNotebooksExperiment.isInNotebooksExperiment()) {
+        if (!uri) {
             return undefined;
         }
 

src/client/activation/node/languageServerProxy.ts

@@ -9,6 +9,7 @@ import {
     LanguageClientOptions,
 } from 'vscode-languageclient/node';
 
+import { Extension } from 'vscode';
 import { IExperimentService, IExtensions, IInterpreterPathService, Resource } from '../../common/types';
 import { IEnvironmentVariablesProvider } from '../../common/variables/types';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
@@ -20,6 +21,7 @@ import { ILanguageClientFactory, ILanguageServerProxy } from '../types';
 import { traceDecoratorError, traceDecoratorVerbose, traceError } from '../../logging';
 import { IWorkspaceService } from '../../common/application/types';
 import { PYLANCE_EXTENSION_ID } from '../../common/constants';
+import { PylanceApi } from './pylanceApi';
 
 // eslint-disable-next-line @typescript-eslint/no-namespace
 namespace InExperiment {
@@ -56,6 +58,8 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
 
     private lsVersion: string | undefined;
 
+    private pylanceApi: PylanceApi | undefined;
+
     constructor(
         private readonly factory: ILanguageClientFactory,
         private readonly experimentService: IExperimentService,
@@ -89,9 +93,16 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
         interpreter: PythonEnvironment | undefined,
         options: LanguageClientOptions,
     ): Promise<void> {
-        const extension = this.extensions.getExtension(PYLANCE_EXTENSION_ID);
+        const extension = await this.getPylanceExtension();
         this.lsVersion = extension?.packageJSON.version || '0';
 
+        const api = extension?.exports;
+        if (api && api.client && api.client.isEnabled()) {
+            this.pylanceApi = api;
+            await api.client.start();
+            return;
+        }
+
         this.cancellationStrategy = new FileBasedCancellationStrategy();
         options.connectionOptions = { cancellationStrategy: this.cancellationStrategy };
 
@@ -111,6 +122,12 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
 
     @traceDecoratorVerbose('Disposing language server')
     public async stop(): Promise<void> {
+        if (this.pylanceApi) {
+            const api = this.pylanceApi;
+            this.pylanceApi = undefined;
+            await api.client!.stop();
+        }
+
         while (this.disposables.length > 0) {
             const d = this.disposables.shift()!;
             d.dispose();
@@ -203,4 +220,17 @@ export class NodeLanguageServerProxy implements ILanguageServerProxy {
             })),
         );
     }
+
+    private async getPylanceExtension(): Promise<Extension<PylanceApi> | undefined> {
+        const extension = this.extensions.getExtension<PylanceApi>(PYLANCE_EXTENSION_ID);
+        if (!extension) {
+            return undefined;
+        }
+
+        if (!extension.isActive) {
+            await extension.activate();
+        }
+
+        return extension;
+    }
 }

src/client/activation/node/lspNotebooksExperiment.ts

@@ -1,16 +1,9 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 import { inject, injectable } from 'inversify';
-import * as semver from 'semver';
-import { Disposable, extensions } from 'vscode';
-import { IConfigurationService, IDisposableRegistry } from '../../common/types';
-import { sendTelemetryEvent } from '../../telemetry';
-import { EventName } from '../../telemetry/constants';
-import { JUPYTER_EXTENSION_ID, PYLANCE_EXTENSION_ID } from '../../common/constants';
-import { IExtensionSingleActivationService, LanguageServerType } from '../types';
-import { traceLog, traceVerbose } from '../../logging';
+import { IExtensionSingleActivationService } from '../types';
+import { traceVerbose } from '../../logging';
 import { IJupyterExtensionDependencyManager } from '../../common/application/types';
-import { ILanguageServerWatcher } from '../../languageServer/types';
 import { IServiceContainer } from '../../ioc/types';
 import { sleep } from '../../common/utils/async';
 import { JupyterExtensionIntegration } from '../../jupyter/jupyterIntegration';
@@ -19,134 +12,30 @@ import { JupyterExtensionIntegration } from '../../jupyter/jupyterIntegration';
 export class LspNotebooksExperiment implements IExtensionSingleActivationService {
     public readonly supportedWorkspaceTypes = { untrustedWorkspace: true, virtualWorkspace: true };
 
-    private pylanceExtensionChangeHandler: Disposable | undefined;
-
     private isJupyterInstalled = false;
 
-    private isInExperiment: boolean | undefined;
-
-    private supportsInteractiveWindow: boolean | undefined;
-
     constructor(
         @inject(IServiceContainer) private readonly serviceContainer: IServiceContainer,
-        @inject(IConfigurationService) private readonly configurationService: IConfigurationService,
-        @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry,
         @inject(IJupyterExtensionDependencyManager) jupyterDependencyManager: IJupyterExtensionDependencyManager,
     ) {
         this.isJupyterInstalled = jupyterDependencyManager.isJupyterExtensionInstalled;
     }
 
-    public async activate(): Promise<void> {
-        if (!LspNotebooksExperiment.isPylanceInstalled()) {
-            this.pylanceExtensionChangeHandler = extensions.onDidChange(this.pylanceExtensionsChangeHandler.bind(this));
-            this.disposables.push(this.pylanceExtensionChangeHandler);
-        }
-
-        this.updateExperimentSupport();
+    // eslint-disable-next-line class-methods-use-this
+    public activate(): Promise<void> {
+        return Promise.resolve();
     }
 
     public async onJupyterInstalled(): Promise<void> {
         if (this.isJupyterInstalled) {
             return;
         }
 
-        if (LspNotebooksExperiment.jupyterSupportsNotebooksExperiment()) {
-            await this.waitForJupyterToRegisterPythonPathFunction();
-            this.updateExperimentSupport();
-        }
+        await this.waitForJupyterToRegisterPythonPathFunction();
 
         this.isJupyterInstalled = true;
     }
 
-    public isInNotebooksExperiment(): boolean {
-        return this.isInExperiment ?? false;
-    }
-
-    public isInNotebooksExperimentWithInteractiveWindowSupport(): boolean {
-        return this.supportsInteractiveWindow ?? false;
-    }
-
-    private updateExperimentSupport(): void {
-        const wasInExperiment = this.isInExperiment;
-        const isInTreatmentGroup = true;
-        const languageServerType = this.configurationService.getSettings().languageServer;
-
-        this.isInExperiment = false;
-        if (languageServerType !== LanguageServerType.Node) {
-            traceLog(`LSP Notebooks experiment is disabled -- not using Pylance`);
-        } else if (!LspNotebooksExperiment.isJupyterInstalled()) {
-            traceLog(`LSP Notebooks experiment is disabled -- Jupyter disabled or not installed`);
-        } else if (!LspNotebooksExperiment.jupyterSupportsNotebooksExperiment()) {
-            traceLog(`LSP Notebooks experiment is disabled -- Jupyter does not support experiment`);
-        } else if (!LspNotebooksExperiment.isPylanceInstalled()) {
-            traceLog(`LSP Notebooks experiment is disabled -- Pylance disabled or not installed`);
-        } else if (!LspNotebooksExperiment.pylanceSupportsNotebooksExperiment()) {
-            traceLog(`LSP Notebooks experiment is disabled -- Pylance does not support experiment`);
-        } else if (!isInTreatmentGroup) {
-            traceLog(`LSP Notebooks experiment is disabled -- not in treatment group`);
-            // to avoid scorecard SRMs, we're also triggering the telemetry for users who meet
-            // the criteria to experience LSP notebooks, but may be in the control group.
-            sendTelemetryEvent(EventName.PYTHON_EXPERIMENTS_LSP_NOTEBOOKS);
-        } else {
-            this.isInExperiment = true;
-            traceLog(`LSP Notebooks experiment is enabled`);
-            sendTelemetryEvent(EventName.PYTHON_EXPERIMENTS_LSP_NOTEBOOKS);
-        }
-
-        this.supportsInteractiveWindow = false;
-        if (!this.isInExperiment) {
-            traceLog(`LSP Notebooks interactive window support is disabled -- not in LSP Notebooks experiment`);
-        } else if (!LspNotebooksExperiment.jupyterSupportsLspInteractiveWindow()) {
-            traceLog(`LSP Notebooks interactive window support is disabled -- Jupyter is not new enough`);
-        } else if (!LspNotebooksExperiment.pylanceSupportsLspInteractiveWindow()) {
-            traceLog(`LSP Notebooks interactive window support is disabled -- Pylance is not new enough`);
-        } else {
-            this.supportsInteractiveWindow = true;
-            traceLog(`LSP Notebooks interactive window support is enabled`);
-        }
-
-        // Our "in experiment" status can only change from false to true. That's possible if Pylance
-        // or Jupyter is installed after Python is activated. A true to false transition would require
-        // either Pylance or Jupyter to be uninstalled or downgraded after Python activated, and that
-        // would require VS Code to be reloaded before the new extension version could be used.
-        if (wasInExperiment === false && this.isInExperiment === true) {
-            const watcher = this.serviceContainer.get<ILanguageServerWatcher>(ILanguageServerWatcher);
-            if (watcher) {
-                watcher.restartLanguageServers();
-            }
-        }
-    }
-
-    private static jupyterSupportsNotebooksExperiment(): boolean {
-        const jupyterVersion = extensions.getExtension(JUPYTER_EXTENSION_ID)?.packageJSON.version;
-        return (
-            jupyterVersion && (semver.gt(jupyterVersion, '2022.5.1001411044') || semver.patch(jupyterVersion) === 100)
-        );
-    }
-
-    private static pylanceSupportsNotebooksExperiment(): boolean {
-        const pylanceVersion = extensions.getExtension(PYLANCE_EXTENSION_ID)?.packageJSON.version;
-        return (
-            pylanceVersion &&
-            (semver.gte(pylanceVersion, '2022.5.3-pre.1') || semver.prerelease(pylanceVersion)?.includes('dev'))
-        );
-    }
-
-    private static jupyterSupportsLspInteractiveWindow(): boolean {
-        const jupyterVersion = extensions.getExtension(JUPYTER_EXTENSION_ID)?.packageJSON.version;
-        return (
-            jupyterVersion && (semver.gt(jupyterVersion, '2022.7.1002041057') || semver.patch(jupyterVersion) === 100)
-        );
-    }
-
-    private static pylanceSupportsLspInteractiveWindow(): boolean {
-        const pylanceVersion = extensions.getExtension(PYLANCE_EXTENSION_ID)?.packageJSON.version;
-        return (
-            pylanceVersion &&
-            (semver.gte(pylanceVersion, '2022.7.51') || semver.prerelease(pylanceVersion)?.includes('dev'))
-        );
-    }
-
     private async waitForJupyterToRegisterPythonPathFunction(): Promise<void> {
         const jupyterExtensionIntegration = this.serviceContainer.get<JupyterExtensionIntegration>(
             JupyterExtensionIntegration,
@@ -168,21 +57,4 @@ export class LspNotebooksExperiment implements IExtensionSingleActivationService
             traceVerbose(`Timed out waiting for Jupyter to call registerJupyterPythonPathFunction`);
         }
     }
-
-    private static isPylanceInstalled(): boolean {
-        return !!extensions.getExtension(PYLANCE_EXTENSION_ID);
-    }
-
-    private static isJupyterInstalled(): boolean {
-        return !!extensions.getExtension(JUPYTER_EXTENSION_ID);
-    }
-
-    private async pylanceExtensionsChangeHandler(): Promise<void> {
-        if (LspNotebooksExperiment.isPylanceInstalled() && this.pylanceExtensionChangeHandler) {
-            this.pylanceExtensionChangeHandler.dispose();
-            this.pylanceExtensionChangeHandler = undefined;
-
-            this.updateExperimentSupport();
-        }
-    }
 }

src/client/activation/node/pylanceApi.ts

@@ -0,0 +1,29 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import {
+    CancellationToken,
+    CompletionContext,
+    CompletionItem,
+    CompletionList,
+    Position,
+    TextDocument,
+    Uri,
+} from 'vscode';
+
+export interface PylanceApi {
+    client?: {
+        isEnabled(): boolean;
+        start(): Promise<void>;
+        stop(): Promise<void>;
+    };
+    notebook?: {
+        registerJupyterPythonPathFunction(func: (uri: Uri) => Promise<string | undefined>): void;
+        registerGetNotebookUriForTextDocumentUriFunction(func: (textDocumentUri: Uri) => Uri | undefined): void;
+        getCompletionItems(
+            document: TextDocument,
+            position: Position,
+            context: CompletionContext,
+            token: CancellationToken,
+        ): Promise<CompletionItem[] | CompletionList | undefined>;
+    };
+}