Keep jupyter connections open and allow for expirations (#13267)

* Idea for refresh -expiration date

* Add expiration functional test

* Fixup after merge

* Add news entry

* Fix sonar problem.

* Fix unit test compilation

* Fix other unit test compile failure

* Dispose timeouts

* Try to get test to pass

Author: rchiodo

news/2 Fixes/12987.md

@@ -0,0 +1 @@
+Allow custom Jupyter server URI providers to have an expiration on their authorization headers.

src/client/datascience/jupyter/jupyterExecution.ts

@@ -17,18 +17,20 @@ import { IServiceContainer } from '../../ioc/types';
 import { PythonInterpreter } from '../../pythonEnvironments/info';
 import { captureTelemetry, sendTelemetryEvent } from '../../telemetry';
 import { JupyterSessionStartError } from '../baseJupyterSession';
-import { Commands, Telemetry } from '../constants';
+import { Commands, Identifiers, Telemetry } from '../constants';
 import { reportAction } from '../progress/decorator';
 import { ReportableAction } from '../progress/types';
 import {
     IJupyterConnection,
     IJupyterExecution,
+    IJupyterServerUri,
     IJupyterSessionManagerFactory,
     IJupyterSubCommandExecutionService,
     IJupyterUriProviderRegistration,
     INotebookServer,
     INotebookServerLaunchInfo,
-    INotebookServerOptions
+    INotebookServerOptions,
+    JupyterServerUriHandle
 } from '../types';
 import { JupyterSelfCertsError } from './jupyterSelfCertsError';
 import { createRemoteConnectionInfo } from './jupyterUtils';
@@ -44,6 +46,8 @@ export class JupyterExecutionBase implements IJupyterExecution {
     private disposed: boolean = false;
     private readonly jupyterInterpreterService: IJupyterSubCommandExecutionService;
     private readonly jupyterPickerRegistration: IJupyterUriProviderRegistration;
+    private uriToJupyterServerUri = new Map<string, IJupyterServerUri>();
+    private pendingTimeouts: (NodeJS.Timeout | number)[] = [];
 
     constructor(
         _liveShare: ILiveShareApi,
@@ -72,6 +76,10 @@ export class JupyterExecutionBase implements IJupyterExecution {
                     // When config changes happen, recreate our commands.
                     this.onSettingsChanged();
                 }
+                if (e.affectsConfiguration('python.dataScience.jupyterServerURI', undefined)) {
+                    // When server URI changes, clear our pending URI timeouts
+                    this.clearTimeouts();
+                }
             });
             this.disposableRegistry.push(disposable);
         }
@@ -83,6 +91,7 @@ export class JupyterExecutionBase implements IJupyterExecution {
 
     public dispose(): Promise<void> {
         this.disposed = true;
+        this.clearTimeouts();
         return Promise.resolve();
     }
 
@@ -367,8 +376,11 @@ export class JupyterExecutionBase implements IJupyterExecution {
                 throw new Error(localize.DataScience.jupyterNotebookFailure().format(''));
             }
         } else {
+            // Prepare our map of server URIs
+            await this.updateServerUri(options.uri);
+
             // If we have a URI spec up a connection info for it
-            return createRemoteConnectionInfo(options.uri, this.jupyterPickerRegistration);
+            return createRemoteConnectionInfo(options.uri, this.getServerUri.bind(this));
         }
     }
 
@@ -385,4 +397,45 @@ export class JupyterExecutionBase implements IJupyterExecution {
         // Clear our usableJupyterInterpreter so that we recompute our values
         this.usablePythonInterpreter = undefined;
     }
+
+    private extractJupyterServerHandleAndId(uri: string): { handle: JupyterServerUriHandle; id: string } | undefined {
+        const url: URL = new URL(uri);
+
+        // Id has to be there too.
+        const id = url.searchParams.get(Identifiers.REMOTE_URI_ID_PARAM);
+        const uriHandle = url.searchParams.get(Identifiers.REMOTE_URI_HANDLE_PARAM);
+        return id && uriHandle ? { handle: uriHandle, id } : undefined;
+    }
+
+    private clearTimeouts() {
+        // tslint:disable-next-line: no-any
+        this.pendingTimeouts.forEach((t) => clearTimeout(t as any));
+        this.pendingTimeouts = [];
+    }
+
+    private getServerUri(uri: string): IJupyterServerUri | undefined {
+        const idAndHandle = this.extractJupyterServerHandleAndId(uri);
+        if (idAndHandle) {
+            return this.uriToJupyterServerUri.get(uri);
+        }
+    }
+
+    private async updateServerUri(uri: string): Promise<void> {
+        const idAndHandle = this.extractJupyterServerHandleAndId(uri);
+        if (idAndHandle) {
+            const serverUri = await this.jupyterPickerRegistration.getJupyterServerUri(
+                idAndHandle.id,
+                idAndHandle.handle
+            );
+            this.uriToJupyterServerUri.set(uri, serverUri);
+            // See if there's an expiration date
+            if (serverUri.expiration) {
+                const timeoutInMS = serverUri.expiration.getTime() - Date.now();
+                // Week seems long enough (in case the expiration is ridiculous)
+                if (timeoutInMS > 0 && timeoutInMS < 604800000) {
+                    this.pendingTimeouts.push(setTimeout(() => this.updateServerUri(uri).ignoreErrors(), timeoutInMS));
+                }
+            }
+        }
+    }
 }

src/client/datascience/jupyter/jupyterRequest.ts

@@ -5,13 +5,14 @@ import * as nodeFetch from 'node-fetch';
 // Function for creating node Request object that prevents jupyterlab services from writing its own
 // authorization header.
 // tslint:disable: no-any
-export function createAuthorizingRequest(authorizationHeader: any) {
+export function createAuthorizingRequest(getAuthHeader: () => any) {
     class AuthorizingRequest extends nodeFetch.Request {
         constructor(input: nodeFetch.RequestInfo, init?: nodeFetch.RequestInit) {
             super(input, init);
 
             // Add all of the authorization parts onto the headers.
             const origHeaders = this.headers;
+            const authorizationHeader = getAuthHeader();
             const keys = Object.keys(authorizationHeader);
             keys.forEach((k) => origHeaders.append(k, authorizationHeader[k].toString()));
             origHeaders.append('Content-Type', 'application/json');

src/client/datascience/jupyter/jupyterSessionManager.ts

@@ -263,12 +263,12 @@ export class JupyterSessionManager implements IJupyterSessionManager {
 
         // If authorization header is provided, then we need to prevent jupyterlab services from
         // writing the authorization header.
-        if (connInfo.authorizationHeader) {
-            requestCtor = createAuthorizingRequest(connInfo.authorizationHeader);
+        if (connInfo.getAuthHeader) {
+            requestCtor = createAuthorizingRequest(connInfo.getAuthHeader);
         }
 
         // If no token is specified prompt for a password
-        if ((connInfo.token === '' || connInfo.token === 'null') && !connInfo.authorizationHeader) {
+        if ((connInfo.token === '' || connInfo.token === 'null') && !connInfo.getAuthHeader) {
             if (this.failOnPassword) {
                 throw new Error('Password request not allowed.');
             }
@@ -314,7 +314,7 @@ export class JupyterSessionManager implements IJupyterSessionManager {
             WebSocket: createJupyterWebSocket(
                 cookieString,
                 allowUnauthorized,
-                connInfo.authorizationHeader
+                connInfo.getAuthHeader
                 // tslint:disable-next-line:no-any
             ) as any,
             // Redefine fetch to our node-modules so it picks up the correct version.

src/client/datascience/jupyter/jupyterUtils.ts

@@ -9,9 +9,8 @@ import { Uri } from 'vscode';
 import { IWorkspaceService } from '../../common/application/types';
 import { noop } from '../../common/utils/misc';
 import { SystemVariables } from '../../common/variables/systemVariables';
-import { Identifiers } from '../constants';
 import { getJupyterConnectionDisplayName } from '../jupyter/jupyterConnection';
-import { IJupyterConnection, IJupyterUriProviderRegistration } from '../types';
+import { IJupyterConnection, IJupyterServerUri } from '../types';
 
 export function expandWorkingDir(
     workingDir: string | undefined,
@@ -27,10 +26,10 @@ export function expandWorkingDir(
     return path.dirname(launchingFile);
 }
 
-export async function createRemoteConnectionInfo(
+export function createRemoteConnectionInfo(
     uri: string,
-    providerRegistration: IJupyterUriProviderRegistration
-): Promise<IJupyterConnection> {
+    getJupyterServerUri: (uri: string) => IJupyterServerUri | undefined
+): IJupyterConnection {
     let url: URL;
     try {
         url = new URL(uri);
@@ -39,14 +38,10 @@ export async function createRemoteConnectionInfo(
         throw err;
     }
 
-    const id = url.searchParams.get(Identifiers.REMOTE_URI_ID_PARAM);
-    const uriHandle = url.searchParams.get(Identifiers.REMOTE_URI_HANDLE_PARAM);
-    const serverUri = id && uriHandle ? await providerRegistration.getJupyterServerUri(id, uriHandle) : undefined;
-    const baseUrl = serverUri && serverUri.baseUrl ? serverUri.baseUrl : `${url.protocol}//${url.host}${url.pathname}`;
-    const token = serverUri && serverUri.token ? serverUri.token : `${url.searchParams.get('token')}`;
-    const hostName = serverUri && serverUri.baseUrl ? new URL(serverUri.baseUrl).hostname : url.hostname;
-    const displayName =
-        serverUri && serverUri.displayName ? serverUri.displayName : getJupyterConnectionDisplayName(token, baseUrl);
+    const serverUri = getJupyterServerUri(uri);
+    const baseUrl = serverUri ? serverUri.baseUrl : `${url.protocol}//${url.host}${url.pathname}`;
+    const token = serverUri ? serverUri.token : `${url.searchParams.get('token')}`;
+    const hostName = serverUri ? new URL(serverUri.baseUrl).hostname : url.hostname;
 
     return {
         type: 'jupyter',
@@ -56,11 +51,14 @@ export async function createRemoteConnectionInfo(
         localLaunch: false,
         localProcExitCode: undefined,
         valid: true,
-        displayName,
+        displayName:
+            serverUri && serverUri.displayName
+                ? serverUri.displayName
+                : getJupyterConnectionDisplayName(token, baseUrl),
         disconnected: (_l) => {
             return { dispose: noop };
         },
         dispose: noop,
-        authorizationHeader: serverUri ? serverUri.authorizationHeader : undefined
+        getAuthHeader: serverUri ? () => getJupyterServerUri(uri)?.authorizationHeader : undefined
     };
 }

src/client/datascience/jupyter/jupyterWebSocket.ts

@@ -3,6 +3,7 @@
 'use strict';
 import * as WebSocketWS from 'ws';
 import { traceError } from '../../common/logger';
+import { noop } from '../../common/utils/misc';
 import { KernelSocketWrapper } from '../kernelSocketWrapper';
 import { IKernelSocket } from '../types';
 
@@ -11,9 +12,10 @@ export const JupyterWebSockets = new Map<string, WebSocketWS & IKernelSocket>();
 
 // We need to override the websocket that jupyter lab services uses to put in our cookie information
 // Do this as a function so that we can pass in variables the the socket will have local access to
-export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?: boolean, authorizationHeader?: any) {
+export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?: boolean, getAuthHeaders?: () => any) {
     class JupyterWebSocket extends KernelSocketWrapper(WebSocketWS) {
         private kernelId: string | undefined;
+        private timer: NodeJS.Timeout | number;
 
         constructor(url: string, protocols?: string | string[] | undefined) {
             let co: WebSocketWS.ClientOptions = {};
@@ -26,7 +28,11 @@ export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?
             if (cookieString) {
                 co_headers = { Cookie: cookieString };
             }
-            if (authorizationHeader) {
+
+            // Auth headers have to be refetched every time we create a connection. They may have expired
+            // since the last connection.
+            if (getAuthHeaders) {
+                const authorizationHeader = getAuthHeaders();
                 co_headers = co_headers ? { ...co_headers, ...authorizationHeader } : authorizationHeader;
             }
             if (co_headers) {
@@ -43,11 +49,15 @@ export function createJupyterWebSocket(cookieString?: string, allowUnauthorized?
             if (this.kernelId) {
                 JupyterWebSockets.set(this.kernelId, this);
                 this.on('close', () => {
+                    clearInterval(this.timer as any);
                     JupyterWebSockets.delete(this.kernelId!);
                 });
             } else {
                 traceError('KernelId not extracted from Kernel WebSocket URL');
             }
+
+            // Ping the websocket connection every 30 seconds to make sure it stays alive
+            this.timer = setInterval(() => this.ping(noop), 30_000);
         }
     }
     return JupyterWebSocket;

src/client/datascience/jupyterUriProviderRegistration.ts

@@ -77,7 +77,7 @@ export class JupyterUriProviderRegistration implements IJupyterUriProviderRegist
             const frames = stack.split('\n').map((f) => {
                 const result = /\((.*)\)/.exec(f);
                 if (result) {
-                    return result[1];
+                    return result[1].toLowerCase();
                 }
             });
             for (const frame of frames) {

src/client/datascience/types.ts

@@ -75,7 +75,7 @@ export interface IJupyterConnection extends Disposable {
     readonly hostName: string;
     localProcExitCode: number | undefined;
     // tslint:disable-next-line: no-any
-    authorizationHeader?: any; // Snould be a json object
+    getAuthHeader?(): any; // Snould be a json object
 }
 
 export type INotebookProviderConnection = IRawConnection | IJupyterConnection;
@@ -1333,13 +1333,14 @@ export interface IJupyterServerUri {
     token: string;
     // tslint:disable-next-line: no-any
     authorizationHeader: any; // JSON object for authorization header.
+    expiration?: Date; // Date/time when header expires and should be refreshed.
     displayName: string;
 }
 
 export type JupyterServerUriHandle = string;
 
 export interface IJupyterUriProvider {
-    id: string; // Should be a unique string (like a guid)
+    readonly id: string; // Should be a unique string (like a guid)
     getQuickPickEntryItems(): QuickPickItem[];
     handleQuickPick(item: QuickPickItem, backEnabled: boolean): Promise<JupyterServerUriHandle | 'back' | undefined>;
     getServerUri(handle: JupyterServerUriHandle): Promise<IJupyterServerUri>;

src/test/datascience/extensionapi/exampleextension/ms-ai-tools-test/src/serverPicker.ts

@@ -76,10 +76,18 @@ export class RemoteServerPickerExample implements IJupyterUriProvider {
                     // some other stuff
                     if (stdout) {
                         const output = JSON.parse(stdout.toString());
+                        const currentDate = new Date();
                         resolve({
                             baseUrl: Compute_ServerUri,
                             token: '', //output.accessToken,
-                            authorizationHeader: { Authorization: `Bearer ${output.accessToken}` }
+                            authorizationHeader: { Authorization: `Bearer ${output.accessToken}` },
+                            expiration: new Date(
+                                currentDate.getFullYear(),
+                                currentDate.getMonth(),
+                                undefined,
+                                currentDate.getHours(),
+                                currentDate.getMinutes() + 1 // Expire after one minute
+                            )
                         });
                     } else {
                         reject('Unable to get az token');

src/test/datascience/extensionapi/exampleextension/ms-ai-tools-test/src/typings/python.d.ts

@@ -52,6 +52,7 @@ export interface IJupyterServerUri {
     token: string;
     // tslint:disable-next-line: no-any
     authorizationHeader: any; // JSON object for authorization header.
+    expiration?: Date; // Date/time when header expires and should be refreshed.
 }
 
 export type JupyterServerUriHandle = string;