From 4097ca7f07a776ec5a2a9748e950eb695610fdcf Mon Sep 17 00:00:00 2001 From: julieg18 Date: Mon, 30 Oct 2023 12:34:40 -0500 Subject: [PATCH] wip demo --- extension/src/setup/index.ts | 84 +++++++++++++++++-- extension/src/setup/token.ts | 32 +++++++ extension/src/setup/webview/contract.ts | 1 + extension/src/setup/webview/messages.ts | 12 ++- extension/src/webview/contract.ts | 4 + webview/src/setup/components/App.test.tsx | 3 +- webview/src/setup/components/App.tsx | 6 +- .../src/setup/components/studio/Connect.tsx | 65 +++++++------- webview/src/setup/state/dvcSlice.ts | 1 + webview/src/setup/state/studioSlice.ts | 15 +++- webview/src/setup/util/messages.ts | 10 +-- webview/src/stories/Setup.stories.tsx | 3 +- 12 files changed, 187 insertions(+), 49 deletions(-) diff --git a/extension/src/setup/index.ts b/extension/src/setup/index.ts index 51cb74b011..7d022aefb5 100644 --- a/extension/src/setup/index.ts +++ b/extension/src/setup/index.ts @@ -4,6 +4,7 @@ import { Disposable, Disposer } from '@hediet/std/disposable' import isEmpty from 'lodash.isempty' import { DvcCliDetails, + STUDIO_URL, SetupSection, SetupData as TSetupData } from './webview/contract' @@ -16,7 +17,7 @@ import { WebviewMessages } from './webview/messages' import { validateTokenInput } from './inputBox' import { findPythonBinForInstall } from './autoInstall' import { run, runWithRecheck, runWorkspace } from './runner' -import { isStudioAccessToken } from './token' +import { isStudioAccessToken, pollForStudioToken } from './token' import { PYTHON_EXTENSION_ACTION, pickFocusedProjects, @@ -69,6 +70,7 @@ import { isActivePythonEnvGlobal, selectPythonInterpreter } from '../extensions/python' +import { openUrl } from '../vscode/external' export class Setup extends BaseRepository @@ -118,7 +120,9 @@ export class Setup private dotFolderWatcher?: Disposer private studioAccessToken: string | undefined = undefined + private studioAuthLink: string | undefined = undefined private studioIsConnected = false + private studioUserCode: string | null = null private shareLiveToStudio: boolean | undefined = undefined private focusedSection: SetupSection | undefined = undefined @@ -357,8 +361,7 @@ export class Setup return } - await this.accessConfig(cwd, Flag.GLOBAL, ConfigKey.STUDIO_TOKEN, token) - return this.updateStudioAndSend() + return this.saveStudioAccessTokenInConfig(cwd, token) } public getStudioAccessToken() { @@ -369,6 +372,11 @@ export class Setup return this.sendDataToWebview() } + private async saveStudioAccessTokenInConfig(cwd: string, token: string) { + await this.accessConfig(cwd, Flag.GLOBAL, ConfigKey.STUDIO_TOKEN, token) + return this.updateStudioAndSend() + } + private async getDvcCliDetails(): Promise { await this.config.isReady() const dvcPath = this.config.getCliPath() @@ -445,7 +453,8 @@ export class Setup pythonBinPath: getBinDisplayText(pythonBinPath), remoteList, sectionCollapsed: collectSectionCollapsed(this.focusedSection), - shareLiveToStudio: !!this.shareLiveToStudio + shareLiveToStudio: !!this.shareLiveToStudio, + studioUserCode: this.getStudioUserCode() }) this.focusedSection = undefined } @@ -456,8 +465,11 @@ export class Setup () => this.initializeGit(), (offline: boolean) => this.updateStudioOffline(offline), () => this.isPythonExtensionUsed(), - () => this.updatePythonEnvironment() + () => this.updatePythonEnvironment(), + () => this.requestStudioAuth(), + () => this.openStudioAuthLink() ) + // add both new actions to above this.dispose.track( this.onDidReceivedWebviewMessage(message => webviewMessages.handleMessageFromWebview(message) @@ -805,6 +817,68 @@ export class Setup ) } + private getStudioUserCode() { + return this.studioUserCode + } + + private async requestStudioAuth() { + const response = await fetch(`${STUDIO_URL}/api/device-login`, { + body: JSON.stringify({ + client_name: 'vscode' + }), + headers: { + 'Content-Type': 'application/json' + }, + method: 'POST' + }) + + const { + token_uri: tokenUri, + verification_uri: verificationUri, + user_code: userCode, + device_code: deviceCode + } = (await response.json()) as { + token_uri: string + verification_uri: string + user_code: string + device_code: string + } + this.studioAuthLink = verificationUri + this.studioUserCode = userCode + + await this.sendDataToWebview() + void this.requestStudioToken(deviceCode, tokenUri) + } + + private async requestStudioToken( + studioDeviceCode: string, + studioTokenRequestUri: string + ) { + const token = await pollForStudioToken( + studioTokenRequestUri, + studioDeviceCode + ) + + this.studioAccessToken = token + this.studioUserCode = null + this.studioAuthLink = undefined + + const cwd = this.getCwd() + + if (!cwd) { + return + } + + return this.saveStudioAccessTokenInConfig(cwd, token) + } + + private openStudioAuthLink() { + if (!this.studioAuthLink) { + return + } + void openUrl(this.studioAuthLink) + } + private getCwd() { if (!this.getCliCompatible()) { return diff --git a/extension/src/setup/token.ts b/extension/src/setup/token.ts index d41d13800f..899ab123e0 100644 --- a/extension/src/setup/token.ts +++ b/extension/src/setup/token.ts @@ -4,3 +4,35 @@ export const isStudioAccessToken = (text?: string): boolean => { } return text.startsWith('isat_') && text.length >= 53 } + +// chose the simplest way to do this, in reality we need a way to +// offer a straightforward way to stop the polling either because the +// user cancels or possibly because were getting errors +export const pollForStudioToken = async ( + tokenUri: string, + deviceCode: string +): Promise => { + const response = await fetch(tokenUri, { + body: JSON.stringify({ + code: deviceCode + }), + headers: { + 'Content-Type': 'application/json' + }, + method: 'POST' + }) + + if (response.status === 400) { + await new Promise(resolve => setTimeout(resolve, 2000)) + return pollForStudioToken(tokenUri, deviceCode) + } + if (response.status !== 200) { + await new Promise(resolve => setTimeout(resolve, 2000)) + return pollForStudioToken(tokenUri, deviceCode) + } + + const { access_token: accessToken } = (await response.json()) as { + access_token: string + } + return accessToken +} diff --git a/extension/src/setup/webview/contract.ts b/extension/src/setup/webview/contract.ts index 5598b4f5ff..60a6d06b60 100644 --- a/extension/src/setup/webview/contract.ts +++ b/extension/src/setup/webview/contract.ts @@ -23,6 +23,7 @@ export type SetupData = { remoteList: RemoteList sectionCollapsed: typeof DEFAULT_SECTION_COLLAPSED | undefined shareLiveToStudio: boolean + studioUserCode: string | null isAboveLatestTestedVersion: boolean | undefined } diff --git a/extension/src/setup/webview/messages.ts b/extension/src/setup/webview/messages.ts index 82e9aa189d..f77e2be461 100644 --- a/extension/src/setup/webview/messages.ts +++ b/extension/src/setup/webview/messages.ts @@ -21,19 +21,25 @@ export class WebviewMessages { private readonly updateStudioOffline: (offline: boolean) => Promise private readonly isPythonExtensionUsed: () => Promise private readonly updatePythonEnv: () => Promise + private readonly requestStudioAuth: () => Promise + private readonly openStudioAuthLink: () => void constructor( getWebview: () => BaseWebview | undefined, initializeGit: () => void, updateStudioOffline: (shareLive: boolean) => Promise, isPythonExtensionUsed: () => Promise, - updatePythonEnv: () => Promise + updatePythonEnv: () => Promise, + requestStudioAuth: () => Promise, + openStudioAuthLink: () => void ) { this.getWebview = getWebview this.initializeGit = initializeGit this.updateStudioOffline = updateStudioOffline this.isPythonExtensionUsed = isPythonExtensionUsed this.updatePythonEnv = updatePythonEnv + this.requestStudioAuth = requestStudioAuth + this.openStudioAuthLink = openStudioAuthLink } public sendWebviewMessage(data: SetupData) { @@ -68,6 +74,8 @@ export class WebviewMessages { return this.openStudio() case MessageFromWebviewType.OPEN_STUDIO_PROFILE: return this.openStudioProfile() + case MessageFromWebviewType.OPEN_STUDIO_AUTH_LINK: + return this.openStudioAuthLink() case MessageFromWebviewType.SAVE_STUDIO_TOKEN: return commands.executeCommand( RegisteredCommands.ADD_STUDIO_ACCESS_TOKEN @@ -78,6 +86,8 @@ export class WebviewMessages { ) case MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE: return this.updateStudioOffline(message.payload) + case MessageFromWebviewType.REQUEST_STUDIO_AUTH: + return this.requestStudioAuth() case MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW: return commands.executeCommand(RegisteredCommands.EXPERIMENT_SHOW) case MessageFromWebviewType.REMOTE_ADD: diff --git a/extension/src/webview/contract.ts b/extension/src/webview/contract.ts index ca19cb9893..684947e150 100644 --- a/extension/src/webview/contract.ts +++ b/extension/src/webview/contract.ts @@ -32,6 +32,7 @@ export enum MessageFromWebviewType { OPEN_PLOTS_WEBVIEW = 'open-plots-webview', OPEN_STUDIO = 'open-studio', OPEN_STUDIO_PROFILE = 'open-studio-profile', + OPEN_STUDIO_AUTH_LINK = 'open-studio-auth-link', PUSH_EXPERIMENT = 'push-experiment', REMOVE_COLUMN_FILTERS = 'remove-column-filters', REMOVE_COLUMN_SORT = 'remove-column-sort', @@ -47,6 +48,7 @@ export enum MessageFromWebviewType { RESET_COMMITS = 'reset-commits', RESIZE_COLUMN = 'resize-column', RESIZE_PLOTS = 'resize-plots', + REQUEST_STUDIO_AUTH = 'request-studio-authentication', SAVE_STUDIO_TOKEN = 'save-studio-token', SET_COMPARISON_MULTI_PLOT_VALUE = 'update-comparison-multi-plot-value', SET_SMOOTH_PLOT_VALUE = 'update-smooth-plot-value', @@ -292,7 +294,9 @@ export type MessageFromWebview = | { type: MessageFromWebviewType.SETUP_WORKSPACE } | { type: MessageFromWebviewType.OPEN_STUDIO } | { type: MessageFromWebviewType.OPEN_STUDIO_PROFILE } + | { type: MessageFromWebviewType.OPEN_STUDIO_AUTH_LINK } | { type: MessageFromWebviewType.SAVE_STUDIO_TOKEN } + | { type: MessageFromWebviewType.REQUEST_STUDIO_AUTH } | { type: MessageFromWebviewType.ADD_CONFIGURATION } | { type: MessageFromWebviewType.ZOOM_PLOT; payload?: string } | { type: MessageFromWebviewType.OPEN_EXPERIMENTS_WEBVIEW } diff --git a/webview/src/setup/components/App.test.tsx b/webview/src/setup/components/App.test.tsx index 3e816c8306..1d744de475 100644 --- a/webview/src/setup/components/App.test.tsx +++ b/webview/src/setup/components/App.test.tsx @@ -41,7 +41,8 @@ const DEFAULT_DATA = { pythonBinPath: undefined, remoteList: undefined, sectionCollapsed: undefined, - shareLiveToStudio: false + shareLiveToStudio: false, + studioUserCode: null } const renderApp = (overrideData: Partial = {}) => { diff --git a/webview/src/setup/components/App.tsx b/webview/src/setup/components/App.tsx index 23b5ddd134..114b214995 100644 --- a/webview/src/setup/components/App.tsx +++ b/webview/src/setup/components/App.tsx @@ -33,7 +33,8 @@ import { import { updateRemoteList } from '../state/remoteSlice' import { updateIsStudioConnected, - updateShareLiveToStudio + updateShareLiveToStudio, + updateStudioUserCode } from '../state/studioSlice' import { setStudioShareExperimentsLive } from '../util/messages' @@ -119,6 +120,9 @@ export const feedStore = ( case 'shareLiveToStudio': dispatch(updateShareLiveToStudio(data.data.shareLiveToStudio)) continue + case 'studioUserCode': + dispatch(updateStudioUserCode(data.data.studioUserCode)) + continue case 'remoteList': dispatch(updateRemoteList(data.data.remoteList)) continue diff --git a/webview/src/setup/components/studio/Connect.tsx b/webview/src/setup/components/studio/Connect.tsx index 9502707010..1ca0504e06 100644 --- a/webview/src/setup/components/studio/Connect.tsx +++ b/webview/src/setup/components/studio/Connect.tsx @@ -1,46 +1,49 @@ import React from 'react' +import { useSelector } from 'react-redux' import { STUDIO_URL } from 'dvc/src/setup/webview/contract' -import { - openStudio, - openStudioProfile, - saveStudioToken -} from '../../util/messages' +import { requestStudioAuth, openStudioAuthLink } from '../../util/messages' import { EmptyState } from '../../../shared/components/emptyState/EmptyState' import { Button } from '../../../shared/components/button/Button' +import { SetupState } from '../../store' export const Connect: React.FC = () => { + const { studioUserCode } = useSelector((state: SetupState) => state.studio) return (

Connect to Studio

-

- Share experiments and plots with collaborators directly from your IDE. - Start sending data with an{' '} - - access token - {' '} - generated from your Studio profile page. -

-