Add AAD auth token to the http header sent to jupyter server

[zesluo]

Feature: Notebook Editor, Interactive Window, Python Editor cells

Description

Hello team! We are currenty trying to support 'Synapse Pyspark Interactive' in VSCODE extension. As the picture shown, it will include a base URL and a token. We want to add the synapse AAD token into it, so it can be sent to jupyter server.

Previously, we used sparkMagic to do this. However,this time, the situation for Synapse is different as the livy for Synapse is not very standard. So we hope to add AAD token to the http header. Thanks a lot!

Microsoft Data Science for VS Code Engineering Team: @rchiodo, @IanMatthewHuff, @DavidKutu, @DonJayamanne, @greazer, @joyceerhl

[zesluo]

Probably, we can add the AAD token here? @rchiodo , @IanMatthewHuff , @DonJayamanne , @greazer , @joyceerhl

[rchiodo]

I believe for this to work, the AAD header would have to be provided here: https://github.com/microsoft/vscode-python/blob/e129bdf8818b98bd653d0235055da8f1deb21dda/src/client/datascience/jupyter/jupyterWebSocket.ts#L28

That's the location where we add extra headers to the request.

[IanMatthewHuff]

If you are trying to add it to the http headers I think you would also need it up a level here: https://github.com/microsoft/vscode-python/blob/e129bdf8818b98bd653d0235055da8f1deb21dda/src/client/datascience/jupyter/jupyterSessionManager.ts#L224 It would need to be on the requestInit object to apply to http request as well as on the websockets.

[rchiodo]

This should be possible now using the registerRemoteServerProvider API

https://github.com/microsoft/vscode-python/blob/97cd4568bbd83810a77e2708141690752b07accb/src/client/api.ts#L97

[rchiodo]

See this code for an example: https://github.com/microsoft/vscode-python/blob/97cd4568bbd83810a77e2708141690752b07accb/src/test/datascience/extensionapi/exampleextension/ms-ai-tools-test/src/serverPicker.ts#L14

[rchiodo]

Sid also posted a video of it working here: https://github.com/microsoft/vscode-python/issues/10993#issuecomment-646340590 (although that will only be visible to Microsoft people)

[DonJayamanne]

I agree with @IanMatthewHuff I think the better way is to expose the ServerConnection.ISettings. As we are dealing with Jupyter, we might as well use strongly typed jupyter services API. The ISettings interface contains all of the information required to connect to any kind of server. It can contain tokens, headers & even custom request/fetch implementations required to retrieve data from a remote server over http(s). I.e. we let the provider give us all of the information, as opposed to us building our own API for exposing headers & tokens & the like.

Using ISettings we could either use Jupyter API to fetch data or use raw request/fetch to retrieve data manually.

[rchiodo]

I believe there was some question on timeout for the token? Maybe the example code doesn't include the addition of the expiration.

You should use these type definitions (found here):

export interface IJupyterServerUri {
    baseUrl: string;
    token: string;
    // eslint-disable-next-line @typescript-eslint/no-explicit-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 {
    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>;
}

You implement the IJupyterUriProvider. The IJupyterServerUri you return from the getServerUri call should include an expiration date. When that date is up we'll reask for the new URI which should include the new auth headers.

[DavidKutu]

Hi @zesluo, does this solve the issue?

[DavidKutu]

@zesluo feel free to reopen this or create a new issue if its not fixed.