diff --git a/extension/package.json b/extension/package.json index d80760a575..421697a9ca 100644 --- a/extension/package.json +++ b/extension/package.json @@ -1345,7 +1345,7 @@ { "id": "dvc.views.studio", "name": "Studio", - "when": "!dvc.studio.connected" + "when": "true" }, { "id": "dvc.views.experimentsColumnsTree", @@ -1394,6 +1394,11 @@ "contents": "[$(plug) Connect](command:dvc.showConnect)", "when": "!dvc.studio.connected" }, + { + "view": "dvc.views.studio", + "contents": "[$(settings-gear) Open Settings](command:dvc.showConnect)", + "when": "dvc.studio.connected" + }, { "view": "dvc.views.welcome", "contents": "New to the extension?\n[Show Walkthrough](command:dvc.getStarted)\n\nThe extension is currently unable to initialize.\n[Show Setup](command:dvc.showSetup)", diff --git a/extension/src/connect/index.ts b/extension/src/connect/index.ts index f726e07d0d..90c936e7cd 100644 --- a/extension/src/connect/index.ts +++ b/extension/src/connect/index.ts @@ -1,7 +1,7 @@ -import { commands, ExtensionContext, SecretStorage } from 'vscode' +import { commands, ExtensionContext, SecretStorage, workspace } from 'vscode' import { validateTokenInput } from './inputBox' import { STUDIO_ACCESS_TOKEN_KEY, isStudioAccessToken } from './token' -import { STUDIO_URL } from './webview/contract' +import { ConnectData, STUDIO_URL } from './webview/contract' import { Resource } from '../resourceLocator' import { ViewKey } from '../webview/constants' import { MessageFromWebview, MessageFromWebviewType } from '../webview/contract' @@ -12,15 +12,15 @@ import { Title } from '../vscode/title' import { openUrl } from '../vscode/external' import { ContextKey, setContextValue } from '../vscode/context' import { RegisteredCommands } from '../commands/external' -import { Modal } from '../vscode/modal' import { GLOBAL_WEBVIEW_DVCROOT } from '../webview/factory' -import { ConfigKey, getConfigValue } from '../vscode/config' +import { ConfigKey, getConfigValue, setConfigValue } from '../vscode/config' -export class Connect extends BaseRepository { +export class Connect extends BaseRepository { public readonly viewKey = ViewKey.CONNECT private readonly secrets: SecretStorage private studioAccessToken: string | undefined = undefined + private studioIsConnected = false constructor(context: ExtensionContext, webviewIcon: Resource) { super(GLOBAL_WEBVIEW_DVCROOT, webviewIcon) @@ -36,7 +36,7 @@ export class Connect extends BaseRepository { void this.getSecret(STUDIO_ACCESS_TOKEN_KEY).then( async studioAccessToken => { this.studioAccessToken = studioAccessToken - await this.setContext() + await this.updateIsStudioConnected() this.deferred.resolve() } ) @@ -46,13 +46,24 @@ export class Connect extends BaseRepository { if (e.key !== STUDIO_ACCESS_TOKEN_KEY) { return } + this.studioAccessToken = await this.getSecret(STUDIO_ACCESS_TOKEN_KEY) - return this.setContext() + return this.updateIsStudioConnected() + }) + ) + + this.dispose.track( + workspace.onDidChangeConfiguration(e => { + if (e.affectsConfiguration(ConfigKey.STUDIO_SHARE_EXPERIMENTS_LIVE)) { + this.sendWebviewMessage() + } }) ) } - public sendInitialWebviewData(): void {} + public sendInitialWebviewData() { + return this.sendWebviewMessage() + } public removeStudioAccessToken() { return this.removeSecret(STUDIO_ACCESS_TOKEN_KEY) @@ -84,6 +95,13 @@ export class Connect extends BaseRepository { return this.studioAccessToken } + private sendWebviewMessage() { + void this.getWebview()?.show({ + isStudioConnected: this.studioIsConnected, + shareLiveToStudio: getConfigValue(ConfigKey.STUDIO_SHARE_EXPERIMENTS_LIVE) + }) + } + private handleMessageFromWebview(message: MessageFromWebview) { switch (message.type) { case MessageFromWebviewType.OPEN_STUDIO: @@ -94,6 +112,15 @@ export class Connect extends BaseRepository { return commands.executeCommand( RegisteredCommands.ADD_STUDIO_ACCESS_TOKEN ) + case MessageFromWebviewType.REMOVE_STUDIO_TOKEN: + return commands.executeCommand( + RegisteredCommands.REMOVE_STUDIO_ACCESS_TOKEN + ) + case MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE: + return setConfigValue( + ConfigKey.STUDIO_SHARE_EXPERIMENTS_LIVE, + message.payload + ) default: Logger.error(`Unexpected message: ${JSON.stringify(message)}`) } @@ -107,19 +134,16 @@ export class Connect extends BaseRepository { return openUrl(`${STUDIO_URL}/user/_/profile?section=accessToken`) } - private setContext() { + private updateIsStudioConnected() { const storedToken = this.getStudioAccessToken() - if (isStudioAccessToken(storedToken)) { - if (this.deferred.state === 'resolved') { - void Modal.showInformation( - 'Studio is now connected. Use the "Share to Studio" command from an experiment\'s context menu to share experiments.' - ) - } - this.webview?.dispose() - return setContextValue(ContextKey.STUDIO_CONNECTED, true) - } + const isConnected = isStudioAccessToken(storedToken) + return this.setStudioIsConnected(isConnected) + } - return setContextValue(ContextKey.STUDIO_CONNECTED, false) + private setStudioIsConnected(isConnected: boolean) { + this.studioIsConnected = isConnected + this.sendWebviewMessage() + return setContextValue(ContextKey.STUDIO_CONNECTED, isConnected) } private getSecret(key: string) { diff --git a/extension/src/connect/webview/contract.ts b/extension/src/connect/webview/contract.ts index 862ff04481..1cc71b06fd 100644 --- a/extension/src/connect/webview/contract.ts +++ b/extension/src/connect/webview/contract.ts @@ -1 +1,6 @@ export const STUDIO_URL = 'https://studio.iterative.ai' + +export type ConnectData = { + isStudioConnected: boolean + shareLiveToStudio: boolean +} diff --git a/extension/src/test/suite/connect/index.test.ts b/extension/src/test/suite/connect/index.test.ts index a0d8523a07..4abe912257 100644 --- a/extension/src/test/suite/connect/index.test.ts +++ b/extension/src/test/suite/connect/index.test.ts @@ -6,9 +6,11 @@ import { ExtensionContext, SecretStorage, Uri, + WorkspaceConfiguration, commands, env, - window + window, + workspace } from 'vscode' import { Disposable } from '../../../extension' import { @@ -22,6 +24,7 @@ import { WEBVIEW_TEST_TIMEOUT } from '../timeouts' import { ContextKey } from '../../../vscode/context' import { STUDIO_ACCESS_TOKEN_KEY } from '../../../connect/token' import { RegisteredCommands } from '../../../commands/external' +import { ConfigKey } from '../../../vscode/config' suite('Connect Test Suite', () => { const disposable = Disposable.fn() @@ -154,6 +157,27 @@ suite('Connect Test Suite', () => { ) }).timeout(WEBVIEW_TEST_TIMEOUT) + it('should handle a message to set dvc.studio.shareExperimentsLive', async () => { + const { connect, mockMessageReceived } = await buildConnect() + await connect.isReady() + + const mockUpdate = stub() + + stub(workspace, 'getConfiguration').returns({ + update: mockUpdate + } as unknown as WorkspaceConfiguration) + + mockMessageReceived.fire({ + payload: false, + type: MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE + }) + + expect(mockUpdate).to.be.calledWithExactly( + ConfigKey.STUDIO_SHARE_EXPERIMENTS_LIVE, + false + ) + }) + it('should be able to delete the Studio access token from secrets storage', async () => { const mockDelete = stub() stub( diff --git a/extension/src/test/suite/util.ts b/extension/src/test/suite/util.ts index d6bab180d6..90a29e4eda 100644 --- a/extension/src/test/suite/util.ts +++ b/extension/src/test/suite/util.ts @@ -35,6 +35,7 @@ import { DvcExecutor } from '../../cli/dvc/executor' import { GitReader } from '../../cli/git/reader' import { SetupData } from '../../setup/webview/contract' import { DvcViewer } from '../../cli/dvc/viewer' +import { ConnectData } from '../../connect/webview/contract' export const mockDisposable = { dispose: stub() @@ -244,7 +245,7 @@ export const buildDependencies = ( } export const getMessageReceivedEmitter = ( - webview: BaseWebview + webview: BaseWebview ): EventEmitter => (webview as any).messageReceived export const getInputBoxEvent = (mockInputValue: string) => { diff --git a/extension/src/webview/contract.ts b/extension/src/webview/contract.ts index 223b7b53b1..10866d5bc0 100644 --- a/extension/src/webview/contract.ts +++ b/extension/src/webview/contract.ts @@ -1,3 +1,4 @@ +import { ConnectData } from '../connect/webview/contract' import { SortDefinition } from '../experiments/model/sortBy' import { TableData } from '../experiments/webview/contract' import { @@ -8,7 +9,7 @@ import { } from '../plots/webview/contract' import { SetupData } from '../setup/webview/contract' -export type WebviewData = TableData | PlotsData | SetupData | {} +export type WebviewData = TableData | PlotsData | SetupData | ConnectData export enum MessageFromWebviewType { INITIALIZED = 'initialized', @@ -49,11 +50,13 @@ export enum MessageFromWebviewType { SELECT_PYTHON_INTERPRETER = 'select-python-interpreter', SET_EXPERIMENTS_FOR_PLOTS = 'set-experiments-for-plots', SET_EXPERIMENTS_AND_OPEN_PLOTS = 'set-experiments-and-open-plots', + SET_STUDIO_SHARE_EXPERIMENTS_LIVE = 'set-studio-share-experiments-live', SHARE_EXPERIMENT_AS_BRANCH = 'share-experiment-as-branch', SHARE_EXPERIMENT_AS_COMMIT = 'share-experiment-as-commit', TOGGLE_METRIC = 'toggle-metric', TOGGLE_PLOTS_SECTION = 'toggle-plots-section', REMOVE_CUSTOM_PLOTS = 'remove-custom-plots', + REMOVE_STUDIO_TOKEN = 'remove-studio-token', MODIFY_EXPERIMENT_PARAMS_AND_QUEUE = 'modify-experiment-params-and-queue', MODIFY_EXPERIMENT_PARAMS_AND_RUN = 'modify-experiment-params-and-run', MODIFY_EXPERIMENT_PARAMS_RESET_AND_RUN = 'modify-experiment-params-reset-and-run', @@ -163,6 +166,7 @@ export type MessageFromWebview = | { type: MessageFromWebviewType.REMOVE_CUSTOM_PLOTS } + | { type: MessageFromWebviewType.REMOVE_STUDIO_TOKEN } | { type: MessageFromWebviewType.REORDER_PLOTS_COMPARISON payload: string[] @@ -209,6 +213,10 @@ export type MessageFromWebview = type: MessageFromWebviewType.SET_EXPERIMENTS_AND_OPEN_PLOTS payload: string[] } + | { + type: MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE + payload: boolean + } | { type: MessageFromWebviewType.SHARE_EXPERIMENT_AS_BRANCH payload: string diff --git a/webview/src/connect/components/App.test.tsx b/webview/src/connect/components/App.test.tsx index bebb930f17..d4bdf8039f 100644 --- a/webview/src/connect/components/App.test.tsx +++ b/webview/src/connect/components/App.test.tsx @@ -1,6 +1,10 @@ import React from 'react' import { cleanup, fireEvent, render, screen } from '@testing-library/react' -import { MessageFromWebviewType } from 'dvc/src/webview/contract' +import { + MessageFromWebviewType, + MessageToWebviewType +} from 'dvc/src/webview/contract' +import { ConnectData } from 'dvc/src/connect/webview/contract' import { App } from './App' import { vsCodeApi } from '../../shared/api' @@ -18,47 +22,79 @@ afterEach(() => { cleanup() }) -const renderApp = () => { - return render() +const setData = ({ isStudioConnected, shareLiveToStudio }: ConnectData) => { + fireEvent( + window, + new MessageEvent('message', { + data: { + data: { + isStudioConnected, + shareLiveToStudio + }, + type: MessageToWebviewType.SET_DATA + } + }) + ) +} + +const renderApp = (isStudioConnected = false, shareLiveToStudio = false) => { + render() + setData({ isStudioConnected, shareLiveToStudio }) } describe('App', () => { - it('should show three buttons which walk the user through saving a token', async () => { - renderApp() - const buttons = await screen.findAllByRole('button') - expect(buttons).toHaveLength(3) - }) + describe('Studio not connected', () => { + it('should show three buttons which walk the user through saving a token', async () => { + renderApp() + const buttons = await screen.findAllByRole('button') + expect(buttons).toHaveLength(3) + }) - it('should show a button which opens Studio', () => { - renderApp() - mockPostMessage.mockClear() - const button = screen.getByText('Sign In') - fireEvent.click(button) - expect(mockPostMessage).toHaveBeenCalledTimes(1) - expect(mockPostMessage).toHaveBeenCalledWith({ - type: MessageFromWebviewType.OPEN_STUDIO + it('should show a button which opens Studio', () => { + renderApp() + mockPostMessage.mockClear() + const button = screen.getByText('Sign In') + fireEvent.click(button) + expect(mockPostMessage).toHaveBeenCalledTimes(1) + expect(mockPostMessage).toHaveBeenCalledWith({ + type: MessageFromWebviewType.OPEN_STUDIO + }) + }) + + it("should show a button which opens the user's Studio profile", () => { + renderApp() + mockPostMessage.mockClear() + const button = screen.getByText('Get Token') + fireEvent.click(button) + expect(mockPostMessage).toHaveBeenCalledTimes(1) + expect(mockPostMessage).toHaveBeenCalledWith({ + type: MessageFromWebviewType.OPEN_STUDIO_PROFILE + }) }) - }) - it("should show a button which opens the user's Studio profile", () => { - renderApp() - mockPostMessage.mockClear() - const button = screen.getByText('Get Token') - fireEvent.click(button) - expect(mockPostMessage).toHaveBeenCalledTimes(1) - expect(mockPostMessage).toHaveBeenCalledWith({ - type: MessageFromWebviewType.OPEN_STUDIO_PROFILE + it("should show a button which allows the user's to save their Studio access token", () => { + renderApp() + mockPostMessage.mockClear() + const button = screen.getByText('Save') + fireEvent.click(button) + expect(mockPostMessage).toHaveBeenCalledTimes(1) + expect(mockPostMessage).toHaveBeenCalledWith({ + type: MessageFromWebviewType.SAVE_STUDIO_TOKEN + }) }) }) - it("should show a button which allows the user's to save their Studio access token", () => { - renderApp() - mockPostMessage.mockClear() - const button = screen.getByText('Save') - fireEvent.click(button) - expect(mockPostMessage).toHaveBeenCalledTimes(1) - expect(mockPostMessage).toHaveBeenCalledWith({ - type: MessageFromWebviewType.SAVE_STUDIO_TOKEN + describe('Studio connected', () => { + it('should render a checkbox which can be used to update dvc.studio.shareExperimentsLive', () => { + const shareExperimentsLive = false + renderApp(true, shareExperimentsLive) + mockPostMessage.mockClear() + const checkbox = screen.getByRole('checkbox') + fireEvent.click(checkbox) + expect(mockPostMessage).toHaveBeenCalledWith({ + payload: !shareExperimentsLive, + type: MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE + }) }) }) }) diff --git a/webview/src/connect/components/App.tsx b/webview/src/connect/components/App.tsx index bef22f02c0..29717f20af 100644 --- a/webview/src/connect/components/App.tsx +++ b/webview/src/connect/components/App.tsx @@ -1,9 +1,40 @@ -import React from 'react' +import React, { useCallback, useState } from 'react' +import { + MessageFromWebviewType, + MessageToWebview +} from 'dvc/src/webview/contract' +import { ConnectData } from 'dvc/src/connect/webview/contract' import { Studio } from './Studio' import { useVsCodeMessaging } from '../../shared/hooks/useVsCodeMessaging' +import { sendMessage } from '../../shared/vscode' export const App: React.FC = () => { - useVsCodeMessaging() + const [isStudioConnected, setIsStudioConnected] = useState(false) + const [shareLiveToStudio, setShareLiveToStudioValue] = + useState(false) + useVsCodeMessaging( + useCallback( + ({ data }: { data: MessageToWebview }) => { + setIsStudioConnected(data.data.isStudioConnected) + setShareLiveToStudioValue(data.data.shareLiveToStudio) + }, + [setIsStudioConnected, setShareLiveToStudioValue] + ) + ) - return + const setShareLiveToStudio = (shouldShareLive: boolean) => { + setShareLiveToStudioValue(shouldShareLive) + sendMessage({ + payload: shouldShareLive, + type: MessageFromWebviewType.SET_STUDIO_SHARE_EXPERIMENTS_LIVE + }) + } + + return ( + + ) } diff --git a/webview/src/connect/components/Studio.tsx b/webview/src/connect/components/Studio.tsx index be0230f14d..c7754637e9 100644 --- a/webview/src/connect/components/Studio.tsx +++ b/webview/src/connect/components/Studio.tsx @@ -1,10 +1,16 @@ import React from 'react' -import { STUDIO_URL } from 'dvc/src/connect/webview/contract' -import { openStudio, openStudioProfile, saveStudioToken } from './messages' +import { VSCodeCheckbox } from '@vscode/webview-ui-toolkit/react' +import { ConnectData, STUDIO_URL } from 'dvc/src/connect/webview/contract' +import { + openStudio, + openStudioProfile, + saveStudioToken, + removeStudioToken +} from './messages' import { EmptyState } from '../../shared/components/emptyState/EmptyState' import { Button } from '../../shared/components/button/Button' -export const Studio: React.FC = () => { +const Connect: React.FC = () => { return (
@@ -47,3 +53,57 @@ export const Studio: React.FC = () => { ) } + +const Settings: React.FC<{ + shareLiveToStudio: boolean + setShareLiveToStudio: (shareLiveToStudio: boolean) => void +}> = ({ shareLiveToStudio, setShareLiveToStudio }) => { + return ( + +
+

Studio Settings

+

+ Experiment metrics and plots logged with DVCLive
+ can be{' '} + + automatically shared to Studio + + . +

+

+ setShareLiveToStudio(!shareLiveToStudio)} + checked={shareLiveToStudio} + > + Share New Experiments Live + +

+
+
+ ) +} + +export const Studio: React.FC< + ConnectData & { setShareLiveToStudio: (shareLiveToStudio: boolean) => void } +> = ({ isStudioConnected, shareLiveToStudio, setShareLiveToStudio }) => { + return isStudioConnected ? ( + + ) : ( + + ) +} diff --git a/webview/src/connect/components/messages.ts b/webview/src/connect/components/messages.ts index d37c831606..4639e05766 100644 --- a/webview/src/connect/components/messages.ts +++ b/webview/src/connect/components/messages.ts @@ -9,3 +9,6 @@ export const openStudioProfile = () => export const saveStudioToken = () => sendMessage({ type: MessageFromWebviewType.SAVE_STUDIO_TOKEN }) + +export const removeStudioToken = () => + sendMessage({ type: MessageFromWebviewType.REMOVE_STUDIO_TOKEN }) diff --git a/webview/src/stories/Connect.stories.tsx b/webview/src/stories/Connect.stories.tsx index 3d1ce50b06..b26747adb9 100644 --- a/webview/src/stories/Connect.stories.tsx +++ b/webview/src/stories/Connect.stories.tsx @@ -1,22 +1,38 @@ import { Story, Meta } from '@storybook/react/types-6-0' -import React from 'react' +import React, { useState } from 'react' import { DISABLE_CHROMATIC_SNAPSHOTS } from './util' import './test-vscode-styles.scss' import '../shared/style.scss' -import { App } from '../connect/components/App' +import { Studio } from '../connect/components/Studio' export default { args: { data: {} }, - component: App, + component: Studio, parameters: DISABLE_CHROMATIC_SNAPSHOTS, title: 'Connect' } as Meta -const Template: Story = () => { - return +const Template: Story = ({ + isStudioConnected, + shareLiveToStudio: initialShareLiveToStudio +}) => { + const [shareLiveToStudio, setShareLiveToStudio] = useState( + !!initialShareLiveToStudio + ) + return ( + + ) } -export const Studio = Template.bind({}) +export const ConnectToStudio = Template.bind({}) +ConnectToStudio.args = { isStudioConnected: false } + +export const StudioSettings = Template.bind({}) +StudioSettings.args = { isStudioConnected: true, shareLiveToStudio: true }