Interactive window middleware for Pylance LSP notebooks

src/client/activation/languageClientMiddlewareBase.ts

@@ -171,6 +171,29 @@ export class LanguageClientMiddlewareBase implements Middleware {
         return this.callNext('willSaveWaitUntil', arguments);
     }
 
+    public async didOpenNotebook() {
+        return this.callNotebooksNext('didOpen', arguments);
+    }
+
+    public async didSaveNotebook() {
+        return this.callNotebooksNext('didSave', arguments);
+    }
+
+    public async didChangeNotebook() {
+        return this.callNotebooksNext('didChange', arguments);
+    }
+
+    public async didCloseNotebook() {
+        return this.callNotebooksNext('didClose', arguments);
+    }
+
+    notebooks = {
+        didOpen: this.didOpenNotebook.bind(this),
+        didSave: this.didSaveNotebook.bind(this),
+        didChange: this.didChangeNotebook.bind(this),
+        didClose: this.didCloseNotebook.bind(this),
+    };
+
     public async provideCompletionItem() {
         if (await this.connected) {
             return this.callNextAndSendTelemetry(
@@ -463,6 +486,17 @@ export class LanguageClientMiddlewareBase implements Middleware {
         return args[args.length - 1](...args);
     }
 
+    private callNotebooksNext(funcName: 'didOpen' | 'didSave' | 'didChange' | 'didClose', args: IArguments) {
+        // This function uses the last argument to call the 'next' item. If we're allowing notebook
+        // middleware, it calls into the notebook middleware first.
+        if (this.notebookAddon?.notebooks && (this.notebookAddon.notebooks as any)[funcName]) {
+            // It would be nice to use args.callee, but not supported in strict mode
+            return (this.notebookAddon.notebooks as any)[funcName](...args);
+        }
+
+        return args[args.length - 1](...args);
+    }
+
     private callNextAndSendTelemetry<T extends keyof MiddleWareMethods>(
         lspMethod: string,
         debounceMilliseconds: number,

src/client/activation/node/analysisOptions.ts

@@ -24,6 +24,7 @@ export class NodeLanguageServerAnalysisOptions extends LanguageServerAnalysisOpt
             experimentationSupport: true,
             trustedWorkspaceSupport: true,
             lspNotebooksSupport: this.lspNotebooksExperiment.isInNotebooksExperiment(),
+            lspInteractiveWindowSupport: this.lspNotebooksExperiment.isInNotebooksExperimentWithInteractiveWindowSupport(),
         } as unknown) as LanguageClientOptions;
     }
 }

src/client/activation/node/languageClientMiddleware.ts

@@ -2,11 +2,13 @@
 // Licensed under the MIT License.
 
 import { Uri } from 'vscode';
+import { LanguageClient } from 'vscode-languageclient/node';
 import { IJupyterExtensionDependencyManager } from '../../common/application/types';
 import { IServiceContainer } from '../../ioc/types';
 import { JupyterExtensionIntegration } from '../../jupyter/jupyterIntegration';
 import { traceLog } from '../../logging';
 import { LanguageClientMiddleware } from '../languageClientMiddleware';
+import { LspInteractiveWindowMiddlewareAddon } from './lspInteractiveWindowMiddlewareAddon';
 
 import { LanguageServerType } from '../types';
 
@@ -15,11 +17,27 @@ import { LspNotebooksExperiment } from './lspNotebooksExperiment';
 export class NodeLanguageClientMiddleware extends LanguageClientMiddleware {
     private readonly lspNotebooksExperiment: LspNotebooksExperiment;
 
-    public constructor(serviceContainer: IServiceContainer, serverVersion?: string) {
+    private readonly jupyterExtensionIntegration: JupyterExtensionIntegration;
+
+    public constructor(
+        serviceContainer: IServiceContainer,
+        private getClient: () => LanguageClient | undefined,
+        serverVersion?: string,
+    ) {
         super(serviceContainer, LanguageServerType.Node, serverVersion);
 
         this.lspNotebooksExperiment = serviceContainer.get<LspNotebooksExperiment>(LspNotebooksExperiment);
         this.setupHidingMiddleware(serviceContainer);
+
+        this.jupyterExtensionIntegration = serviceContainer.get<JupyterExtensionIntegration>(
+            JupyterExtensionIntegration,
+        );
+        if (!this.notebookAddon && this.lspNotebooksExperiment.isInNotebooksExperimentWithInteractiveWindowSupport()) {
+            this.notebookAddon = new LspInteractiveWindowMiddlewareAddon(
+                this.getClient,
+                this.jupyterExtensionIntegration,
+            );
+        }
     }
 
     protected shouldCreateHidingMiddleware(jupyterDependencyManager: IJupyterExtensionDependencyManager): boolean {
@@ -34,18 +52,24 @@ export class NodeLanguageClientMiddleware extends LanguageClientMiddleware {
             await this.lspNotebooksExperiment.onJupyterInstalled();
         }
 
-        super.onExtensionChange(jupyterDependencyManager);
+        if (this.lspNotebooksExperiment.isInNotebooksExperimentWithInteractiveWindowSupport()) {
+            if (!this.notebookAddon) {
+                this.notebookAddon = new LspInteractiveWindowMiddlewareAddon(
+                    this.getClient,
+                    this.jupyterExtensionIntegration,
+                );
+            }
+        } else {
+            super.onExtensionChange(jupyterDependencyManager);
+        }
     }
 
     protected async getPythonPathOverride(uri: Uri | undefined): Promise<string | undefined> {
         if (!uri || !this.lspNotebooksExperiment.isInNotebooksExperiment()) {
             return undefined;
         }
 
-        const jupyterExtensionIntegration = this.serviceContainer?.get<JupyterExtensionIntegration>(
-            JupyterExtensionIntegration,
-        );
-        const jupyterPythonPathFunction = jupyterExtensionIntegration?.getJupyterPythonPathFunction();
+        const jupyterPythonPathFunction = this.jupyterExtensionIntegration.getJupyterPythonPathFunction();
         if (!jupyterPythonPathFunction) {
             return undefined;
         }

src/client/activation/node/lspInteractiveWindowMiddlewareAddon.ts

@@ -0,0 +1,187 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+import { Disposable, NotebookCell, NotebookDocument, TextDocument, TextDocumentChangeEvent, Uri } from 'vscode';
+import { Converter } from 'vscode-languageclient/lib/common/codeConverter';
+import {
+    DidChangeNotebookDocumentNotification,
+    LanguageClient,
+    Middleware,
+    NotebookCellKind,
+    NotebookDocumentChangeEvent,
+} from 'vscode-languageclient/node';
+import * as proto from 'vscode-languageserver-protocol';
+import { JupyterExtensionIntegration } from '../../jupyter/jupyterIntegration';
+
+type TextContent = Required<Required<Required<proto.NotebookDocumentChangeEvent>['cells']>['textContent']>[0];
+
+/**
+ * Detects the input box text documents of Interactive Windows and makes them appear to be
+ * the last cell of their corresponding notebooks.
+ */
+export class LspInteractiveWindowMiddlewareAddon implements Middleware, Disposable {
+    constructor(
+        private readonly getClient: () => LanguageClient | undefined,
+        private readonly jupyterExtensionIntegration: JupyterExtensionIntegration,
+    ) {
+        // Make sure a bunch of functions are bound to this. VS code can call them without a this context
+        this.didOpen = this.didOpen.bind(this);
+        this.didChange = this.didChange.bind(this);
+        this.didClose = this.didClose.bind(this);
+    }
+
+    public dispose(): void {
+        // Nothing to dispose at the moment
+    }
+
+    // Map of document URIs to NotebookDocuments for all known notebooks.
+    private notebookDocumentMap: Map<string, NotebookDocument> = new Map<string, NotebookDocument>();
+
+    // Map of document URIs to TextDocuments that should be linked to a notebook
+    // whose didOpen we're expecting to see in the future.
+    private unlinkedTextDocumentMap: Map<string, TextDocument> = new Map<string, TextDocument>();
+
+    public async didOpen(document: TextDocument, next: (ev: TextDocument) => Promise<void>): Promise<void> {
+        const notebookUri = this.getNotebookUriForTextDocumentUri(document.uri);
+        if (!notebookUri) {
+            await next(document);
+            return;
+        }
+
+        const notebookDocument = this.notebookDocumentMap.get(notebookUri.toString());
+        if (!notebookDocument) {
+            this.unlinkedTextDocumentMap.set(notebookUri.toString(), document);
+            return;
+        }
+
+        try {
+            const result: NotebookDocumentChangeEvent = {
+                cells: {
+                    structure: {
+                        array: {
+                            start: notebookDocument.cellCount,
+                            deleteCount: 0,
+                            cells: [{ kind: NotebookCellKind.Code, document: document.uri.toString() }],
+                        },
+                        didOpen: [
+                            {
+                                uri: document.uri.toString(),
+                                languageId: document.languageId,
+                                version: document.version,
+                                text: document.getText(),
+                            },
+                        ],
+                        didClose: undefined,
+                    },
+                },
+            };
+
+            await this.getClient()?.sendNotification(DidChangeNotebookDocumentNotification.type, {
+                notebookDocument: { version: notebookDocument.version, uri: notebookUri.toString() },
+                change: result,
+            });
+        } catch (error) {
+            this.getClient()?.error('Sending DidChangeNotebookDocumentNotification failed', error);
+            throw error;
+        }
+    }
+
+    public async didChange(
+        event: TextDocumentChangeEvent,
+        next: (ev: TextDocumentChangeEvent) => Promise<void>,
+    ): Promise<void> {
+        const notebookUri = this.getNotebookUriForTextDocumentUri(event.document.uri);
+        if (!notebookUri) {
+            await next(event);
+            return;
+        }
+
+        const notebookDocument = this.notebookDocumentMap.get(notebookUri.toString());
+        if (notebookDocument) {
+            const client = this.getClient();
+            if (client) {
+                client.sendNotification(proto.DidChangeNotebookDocumentNotification.type, {
+                    notebookDocument: { uri: notebookUri.toString(), version: notebookDocument.version },
+                    change: {
+                        cells: {
+                            textContent: [
+                                LspInteractiveWindowMiddlewareAddon._asTextContentChange(
+                                    event,
+                                    client.code2ProtocolConverter,
+                                ),
+                            ],
+                        },
+                    },
+                });
+            }
+        }
+    }
+
+    private static _asTextContentChange(event: TextDocumentChangeEvent, c2pConverter: Converter): TextContent {
+        const params = c2pConverter.asChangeTextDocumentParams(event);
+        return { document: params.textDocument, changes: params.contentChanges };
+    }
+
+    public async didClose(document: TextDocument, next: (ev: TextDocument) => Promise<void>): Promise<void> {
+        const notebookUri = this.getNotebookUriForTextDocumentUri(document.uri);
+        if (!notebookUri) {
+            await next(document);
+            return;
+        }
+
+        this.unlinkedTextDocumentMap.delete(notebookUri.toString());
+    }
+
+    public async didOpenNotebook(
+        notebookDocument: NotebookDocument,
+        cells: NotebookCell[],
+        next: (notebookDocument: NotebookDocument, cells: NotebookCell[]) => Promise<void>,
+    ): Promise<void> {
+        this.notebookDocumentMap.set(notebookDocument.uri.toString(), notebookDocument);
+
+        const relatedTextDocument = this.unlinkedTextDocumentMap.get(notebookDocument.uri.toString());
+        if (relatedTextDocument) {
+            const newCells = [
+                ...cells,
+                {
+                    index: notebookDocument.cellCount,
+                    notebook: notebookDocument,
+                    kind: NotebookCellKind.Code,
+                    document: relatedTextDocument,
+                    metadata: {},
+                    outputs: [],
+                    executionSummary: undefined,
+                },
+            ];
+
+            this.unlinkedTextDocumentMap.delete(notebookDocument.uri.toString());
+
+            await next(notebookDocument, newCells);
+        } else {
+            await next(notebookDocument, cells);
+        }
+    }
+
+    public async didCloseNotebook(
+        notebookDocument: NotebookDocument,
+        cells: NotebookCell[],
+        next: (notebookDocument: NotebookDocument, cells: NotebookCell[]) => Promise<void>,
+    ): Promise<void> {
+        this.notebookDocumentMap.delete(notebookDocument.uri.toString());
+
+        await next(notebookDocument, cells);
+    }
+
+    notebooks = {
+        didOpen: this.didOpenNotebook.bind(this),
+        didClose: this.didCloseNotebook.bind(this),
+    };
+
+    private getNotebookUriForTextDocumentUri(textDocumentUri: Uri): Uri | undefined {
+        const getNotebookUriFunction = this.jupyterExtensionIntegration.getGetNotebookUriForTextDocumentUriFunction();
+        if (!getNotebookUriFunction) {
+            return undefined;
+        }
+
+        return getNotebookUriFunction(textDocumentUri);
+    }
+}

src/client/activation/node/lspNotebooksExperiment.ts

@@ -25,6 +25,8 @@ export class LspNotebooksExperiment implements IExtensionSingleActivationService
 
     private isInExperiment: boolean | undefined;
 
+    private supportsInteractiveWindow: boolean | undefined;
+
     constructor(
         @inject(IServiceContainer) private readonly serviceContainer: IServiceContainer,
         @inject(IConfigurationService) private readonly configurationService: IConfigurationService,
@@ -60,6 +62,10 @@ export class LspNotebooksExperiment implements IExtensionSingleActivationService
         return this.isInExperiment ?? false;
     }
 
+    public isInNotebooksExperimentWithInteractiveWindowSupport(): boolean {
+        return this.supportsInteractiveWindow ?? false;
+    }
+
     private updateExperimentSupport(): void {
         const wasInExperiment = this.isInExperiment;
         const isInTreatmentGroup = this.configurationService.getSettings().pylanceLspNotebooksEnabled;
@@ -87,6 +93,18 @@ export class LspNotebooksExperiment implements IExtensionSingleActivationService
             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
@@ -114,6 +132,21 @@ export class LspNotebooksExperiment implements IExtensionSingleActivationService
         );
     }
 
+    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,

src/client/activation/node/manager.ts

@@ -117,7 +117,11 @@ export class NodeLanguageServerManager implements ILanguageServerManager {
     @traceDecoratorVerbose('Starting language server')
     protected async startLanguageServer(): Promise<void> {
         const options = await this.analysisOptions.getAnalysisOptions();
-        this.middleware = new NodeLanguageClientMiddleware(this.serviceContainer, this.lsVersion);
+        this.middleware = new NodeLanguageClientMiddleware(
+            this.serviceContainer,
+            () => this.languageServerProxy.languageClient,
+            this.lsVersion,
+        );
         options.middleware = this.middleware;
 
         // Make sure the middleware is connected if we restart and we we're already connected.