Add API to get language server from external extensions

src/client/activation/common/activatorBase.ts

@@ -18,7 +18,6 @@ import {
     SignatureHelpContext,
     SymbolInformation,
     TextDocument,
-    TextDocumentContentChangeEvent,
     WorkspaceEdit
 } from 'vscode';
 import * as vscodeLanguageClient from 'vscode-languageclient/node';
@@ -72,34 +71,25 @@ export abstract class LanguageServerActivatorBase implements ILanguageServerActi
         this.manager.disconnect();
     }
 
-    public handleOpen(document: TextDocument): void {
+    public get connection() {
         const languageClient = this.getLanguageClient();
         if (languageClient) {
-            languageClient.sendNotification(
-                vscodeLanguageClient.DidOpenTextDocumentNotification.type,
-                languageClient.code2ProtocolConverter.asOpenTextDocumentParams(document)
-            );
+            // Return an object that looks like a connection
+            return {
+                sendNotification: languageClient.sendNotification.bind(languageClient),
+                sendRequest: languageClient.sendRequest.bind(languageClient),
+                sendProgress: languageClient.sendProgress.bind(languageClient),
+                onRequest: languageClient.onRequest.bind(languageClient),
+                onNotification: languageClient.onNotification.bind(languageClient),
+                onProgress: languageClient.onProgress.bind(languageClient)
+            };
         }
     }
 
-    public handleChanges(document: TextDocument, changes: TextDocumentContentChangeEvent[]): void {
+    public get capabilities() {
         const languageClient = this.getLanguageClient();
         if (languageClient) {
-            // If the language client doesn't support incremental, just send the whole document
-            if (this.textDocumentSyncKind === vscodeLanguageClient.TextDocumentSyncKind.Full) {
-                languageClient.sendNotification(
-                    vscodeLanguageClient.DidChangeTextDocumentNotification.type,
-                    languageClient.code2ProtocolConverter.asChangeTextDocumentParams(document)
-                );
-            } else {
-                languageClient.sendNotification(
-                    vscodeLanguageClient.DidChangeTextDocumentNotification.type,
-                    languageClient.code2ProtocolConverter.asChangeTextDocumentParams({
-                        document,
-                        contentChanges: changes
-                    })
-                );
-            }
+            return languageClient.initializeResult?.capabilities;
         }
     }
 
@@ -169,23 +159,6 @@ export abstract class LanguageServerActivatorBase implements ILanguageServerActi
         }
     }
 
-    private get textDocumentSyncKind(): vscodeLanguageClient.TextDocumentSyncKind {
-        const languageClient = this.getLanguageClient();
-        if (languageClient?.initializeResult?.capabilities?.textDocumentSync) {
-            const syncOptions = languageClient.initializeResult.capabilities.textDocumentSync;
-            const syncKind =
-                syncOptions !== undefined && syncOptions.hasOwnProperty('change')
-                    ? (syncOptions as vscodeLanguageClient.TextDocumentSyncOptions).change
-                    : syncOptions;
-            if (syncKind !== undefined) {
-                return syncKind as vscodeLanguageClient.TextDocumentSyncKind;
-            }
-        }
-
-        // Default is full if not provided
-        return vscodeLanguageClient.TextDocumentSyncKind.Full;
-    }
-
     private async handleProvideRenameEdits(
         document: TextDocument,
         position: Position,

src/client/activation/jedi/multiplexingActivator.ts

@@ -9,8 +9,7 @@ import {
     Position,
     ReferenceContext,
     SignatureHelpContext,
-    TextDocument,
-    TextDocumentContentChangeEvent
+    TextDocument
 } from 'vscode';
 // tslint:disable-next-line: import-name
 import { IWorkspaceService } from '../../common/application/types';
@@ -73,15 +72,15 @@ export class MultiplexingJediLanguageServerActivator implements ILanguageServerA
         return this.onDidChangeCodeLensesEmitter.event;
     }
 
-    public handleChanges(document: TextDocument, changes: TextDocumentContentChangeEvent[]) {
-        if (this.realLanguageServer && this.realLanguageServer.handleChanges) {
-            this.realLanguageServer.handleChanges(document, changes);
+    public get connection() {
+        if (this.realLanguageServer) {
+            return this.realLanguageServer.connection;
         }
     }
 
-    public handleOpen(document: TextDocument) {
-        if (this.realLanguageServer && this.realLanguageServer.handleOpen) {
-            this.realLanguageServer.handleOpen(document);
+    public get capabilities() {
+        if (this.realLanguageServer) {
+            return this.realLanguageServer.capabilities;
         }
     }
 

src/client/activation/refCountedLanguageServer.ts

@@ -17,7 +17,6 @@ import {
     SignatureHelpContext,
     SymbolInformation,
     TextDocument,
-    TextDocumentContentChangeEvent,
     WorkspaceEdit
 } from 'vscode';
 
@@ -65,12 +64,12 @@ export class RefCountedLanguageServer implements ILanguageServerActivator {
         this.impl.clearAnalysisCache ? this.impl.clearAnalysisCache() : noop();
     }
 
-    public handleChanges(document: TextDocument, changes: TextDocumentContentChangeEvent[]) {
-        this.impl.handleChanges ? this.impl.handleChanges(document, changes) : noop();
+    public get connection() {
+        return this.impl.connection;
     }
 
-    public handleOpen(document: TextDocument) {
-        this.impl.handleOpen ? this.impl.handleOpen(document) : noop();
+    public get capabilities() {
+        return this.impl.capabilities;
     }
 
     public provideRenameEdits(

src/client/activation/types.ts

@@ -13,11 +13,10 @@ import {
     HoverProvider,
     ReferenceProvider,
     RenameProvider,
-    SignatureHelpProvider,
-    TextDocument,
-    TextDocumentContentChangeEvent
+    SignatureHelpProvider
 } from 'vscode';
 import { LanguageClient, LanguageClientOptions } from 'vscode-languageclient/node';
+import * as lsp from 'vscode-languageserver-protocol';
 import { NugetPackage } from '../common/nuget/types';
 import { IDisposable, IOutputChannel, LanguageServerDownloadChannels, Resource } from '../common/types';
 import { PythonEnvironment } from '../pythonEnvironments/info';
@@ -73,17 +72,20 @@ export enum LanguageServerType {
 export const DotNetLanguageServerFolder = 'languageServer';
 export const NodeLanguageServerFolder = 'nodeLanguageServer';
 
-// tslint:disable-next-line: interface-name
-export interface DocumentHandler {
-    handleOpen(document: TextDocument): void;
-    handleChanges(document: TextDocument, changes: TextDocumentContentChangeEvent[]): void;
-}
-
 // tslint:disable-next-line: interface-name
 export interface LanguageServerCommandHandler {
     clearAnalysisCache(): void;
 }
 
+/**
+ * This interface is a subset of the vscode-protocol connection interface.
+ * It's the minimum set of functions needed in order to talk to a language server.
+ */
+export type ILanguageServerConnection = Pick<
+    lsp.ProtocolConnection,
+    'sendRequest' | 'sendNotification' | 'onProgress' | 'sendProgress' | 'onNotification' | 'onRequest'
+>;
+
 export interface ILanguageServer
     extends RenameProvider,
         DefinitionProvider,
@@ -93,9 +95,11 @@ export interface ILanguageServer
         CodeLensProvider,
         DocumentSymbolProvider,
         SignatureHelpProvider,
-        Partial<DocumentHandler>,
         Partial<LanguageServerCommandHandler>,
-        IDisposable {}
+        IDisposable {
+    readonly connection?: ILanguageServerConnection;
+    readonly capabilities?: lsp.ServerCapabilities;
+}
 
 export const ILanguageServerActivator = Symbol('ILanguageServerActivator');
 export interface ILanguageServerActivator extends ILanguageServer {

src/client/datascience/api/jupyterIntegration.ts

@@ -7,9 +7,12 @@
 
 import { inject, injectable } from 'inversify';
 import { dirname } from 'path';
-import { CancellationToken, Event, Uri } from 'vscode';
+import { CancellationToken, Disposable, Event, Uri } from 'vscode';
+import * as lsp from 'vscode-languageserver-protocol';
+import { ILanguageServerCache, ILanguageServerConnection } from '../../activation/types';
 import { InterpreterUri } from '../../common/installer/types';
 import { IExtensions, IInstaller, InstallerResponse, Product, Resource } from '../../common/types';
+import { isResource } from '../../common/utils/misc';
 import { getDebugpyPackagePath } from '../../debugger/extension/adapter/remoteLaunchers';
 import { IEnvironmentActivationService } from '../../interpreter/activation/types';
 import { IInterpreterQuickPickItem, IInterpreterSelector } from '../../interpreter/configuration/types';
@@ -18,6 +21,11 @@ import { IWindowsStoreInterpreter } from '../../interpreter/locators/types';
 import { WindowsStoreInterpreter } from '../../pythonEnvironments/discovery/locators/services/windowsStoreInterpreter';
 import { PythonEnvironment } from '../../pythonEnvironments/info';
 
+export interface ILanguageServer extends Disposable {
+    readonly connection: ILanguageServerConnection;
+    readonly capabilities: lsp.ServerCapabilities;
+}
+
 type PythonApiForJupyterExtension = {
     /**
      * IInterpreterService
@@ -57,9 +65,14 @@ type PythonApiForJupyterExtension = {
      * Returns path to where `debugpy` is. In python extension this is `/pythonFiles/lib/python`.
      */
     getDebuggerPath(): Promise<string>;
+    /**
+     * Returns a ILanguageServer that can be used for communicating with a language server process.
+     * @param resource file that determines which connection to return
+     */
+    getLanguageServer(resource?: InterpreterUri): Promise<ILanguageServer | undefined>;
 };
 
-type JupyterExtensionApi = {
+export type JupyterExtensionApi = {
     registerPythonApi(interpreterService: PythonApiForJupyterExtension): void;
 };
 
@@ -71,19 +84,11 @@ export class JupyterExtensionIntegration {
         @inject(IInterpreterSelector) private readonly interpreterSelector: IInterpreterSelector,
         @inject(WindowsStoreInterpreter) private readonly windowsStoreInterpreter: IWindowsStoreInterpreter,
         @inject(IInstaller) private readonly installer: IInstaller,
-        @inject(IEnvironmentActivationService) private readonly envActivation: IEnvironmentActivationService
+        @inject(IEnvironmentActivationService) private readonly envActivation: IEnvironmentActivationService,
+        @inject(ILanguageServerCache) private readonly languageServerCache: ILanguageServerCache
     ) {}
 
-    public async integrateWithJupyterExtension(): Promise<void> {
-        const jupyterExtension = this.extensions.getExtension<JupyterExtensionApi>('ms-ai-tools.jupyter');
-        if (!jupyterExtension) {
-            return;
-        }
-        await jupyterExtension.activate();
-        if (!jupyterExtension.isActive) {
-            return;
-        }
-        const jupyterExtensionApi = jupyterExtension.exports;
+    public registerApi(jupyterExtensionApi: JupyterExtensionApi) {
         jupyterExtensionApi.registerPythonApi({
             onDidChangeInterpreter: this.interpreterService.onDidChangeInterpreter,
             getActiveInterpreter: async (resource?: Uri) => this.interpreterService.getActiveInterpreter(resource),
@@ -104,7 +109,34 @@ export class JupyterExtensionIntegration {
                 resource?: InterpreterUri,
                 cancel?: CancellationToken
             ): Promise<InstallerResponse> => this.installer.install(product, resource, cancel),
-            getDebuggerPath: async () => dirname(getDebugpyPackagePath())
+            getDebuggerPath: async () => dirname(getDebugpyPackagePath()),
+            getLanguageServer: async (r) => {
+                const resource = isResource(r) ? r : undefined;
+                const interpreter = !isResource(r) ? r : undefined;
+                const client = await this.languageServerCache.get(resource, interpreter);
+
+                // Some langauge servers don't support the connection yet. (like Jedi until we switch to LSP)
+                if (client && client.connection && client.capabilities) {
+                    return {
+                        connection: client.connection,
+                        capabilities: client.capabilities,
+                        dispose: client.dispose
+                    };
+                }
+                return undefined;
+            }
         });
     }
+
+    public async integrateWithJupyterExtension(): Promise<void> {
+        const jupyterExtension = this.extensions.getExtension<JupyterExtensionApi>('ms-ai-tools.jupyter');
+        if (!jupyterExtension) {
+            return;
+        }
+        await jupyterExtension.activate();
+        if (!jupyterExtension.isActive) {
+            return;
+        }
+        this.registerApi(jupyterExtension.exports);
+    }
 }

src/client/datascience/interactive-common/intellisense/intellisenseProvider.ts

@@ -22,7 +22,6 @@ import {
 import { CancellationToken } from 'vscode-jsonrpc';
 import * as vscodeLanguageClient from 'vscode-languageclient/node';
 import { concatMultilineString } from '../../../../datascience-ui/common';
-import { ILanguageServer, ILanguageServerCache } from '../../../activation/types';
 import { IWorkspaceService } from '../../../common/application/types';
 import { CancellationError } from '../../../common/cancellation';
 import { traceError, traceWarning } from '../../../common/logger';
@@ -34,6 +33,7 @@ import { HiddenFileFormatString } from '../../../constants';
 import { IInterpreterService } from '../../../interpreter/contracts';
 import { PythonEnvironment } from '../../../pythonEnvironments/info';
 import { sendTelemetryWhenDone } from '../../../telemetry';
+import { JupyterExtensionIntegration } from '../../api/jupyterIntegration';
 import { Identifiers, Settings, Telemetry } from '../../constants';
 import {
     ICell,
@@ -65,6 +65,7 @@ import {
     convertToVSCodeCompletionItem
 } from './conversion';
 import { IntellisenseDocument } from './intellisenseDocument';
+import { NotebookLanguageServer } from './notebookLanguageServer';
 
 // These regexes are used to get the text from jupyter output by recognizing escape charactor \x1b
 const DocStringRegex = /\x1b\[1;31mDocstring:\x1b\[0m\s+([\s\S]*?)\r?\n\x1b\[1;31m/;
@@ -101,7 +102,7 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
     private notebookType: 'interactive' | 'native' = 'interactive';
     private potentialResource: Uri | undefined;
     private sentOpenDocument: boolean = false;
-    private languageServer: ILanguageServer | undefined;
+    private languageServer: NotebookLanguageServer | undefined;
     private resource: Resource;
     private interpreter: PythonEnvironment | undefined;
 
@@ -110,7 +111,7 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
         @inject(IDataScienceFileSystem) private fs: IDataScienceFileSystem,
         @inject(INotebookProvider) private notebookProvider: INotebookProvider,
         @inject(IInterpreterService) private interpreterService: IInterpreterService,
-        @inject(ILanguageServerCache) private languageServerCache: ILanguageServerCache,
+        @inject(JupyterExtensionIntegration) private jupyterApiProvider: JupyterExtensionIntegration,
         @inject(IJupyterVariables) @named(Identifiers.ALL_VARIABLES) private variableProvider: IJupyterVariables
     ) {}
 
@@ -198,7 +199,7 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
         return this.documentPromise.promise;
     }
 
-    protected async getLanguageServer(token: CancellationToken): Promise<ILanguageServer | undefined> {
+    protected async getLanguageServer(token: CancellationToken): Promise<NotebookLanguageServer | undefined> {
         // Resource should be our potential resource if its set. Otherwise workspace root
         const resource =
             this.potentialResource ||
@@ -225,22 +226,26 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
 
             // Get an instance of the language server (so we ref count it )
             try {
-                const languageServer = await this.languageServerCache.get(resource, interpreter);
+                const languageServer = await NotebookLanguageServer.create(
+                    this.jupyterApiProvider,
+                    resource,
+                    interpreter
+                );
 
                 // Dispose of our old language service
                 this.languageServer?.dispose();
 
                 // This new language server does not know about our document, so tell it.
                 const document = await this.getDocument();
-                if (document && languageServer.handleOpen && languageServer.handleChanges) {
+                if (document && languageServer) {
                     // If we already sent an open document, that means we need to send both the open and
                     // the new changes
                     if (this.sentOpenDocument) {
-                        languageServer.handleOpen(document);
-                        languageServer.handleChanges(document, document.getFullContentChanges());
+                        languageServer.sendOpen(document);
+                        languageServer.sendChanges(document, document.getFullContentChanges());
                     } else {
                         this.sentOpenDocument = true;
-                        languageServer.handleOpen(document);
+                        languageServer.sendOpen(document);
                     }
                 }
 
@@ -352,7 +357,7 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
         token: CancellationToken
     ): Promise<monacoEditor.languages.CompletionItem> {
         const [languageServer, document] = await Promise.all([this.getLanguageServer(token), this.getDocument()]);
-        if (languageServer && languageServer.resolveCompletionItem && document) {
+        if (languageServer && document) {
             const vscodeCompItem: CompletionItem = convertToVSCodeCompletionItem(item);
 
             // Needed by Jedi in completionSource.ts to resolve the item
@@ -378,12 +383,12 @@ export class IntellisenseProvider implements IInteractiveWindowListener {
         if (document) {
             // Broadcast an update to the language server
             const languageServer = await this.getLanguageServer(CancellationToken.None);
-            if (languageServer && languageServer.handleChanges && languageServer.handleOpen) {
+            if (languageServer) {
                 if (!this.sentOpenDocument) {
                     this.sentOpenDocument = true;
-                    return languageServer.handleOpen(document);
+                    return languageServer.sendOpen(document);
                 } else {
-                    return languageServer.handleChanges(document, changes);
+                    return languageServer.sendChanges(document, changes);
                 }
             }
         }