From 358666e7fe869c2b037f723111dc765bedeaebe8 Mon Sep 17 00:00:00 2001 From: anthonyting <49772744+anthonyting@users.noreply.github.com> Date: Thu, 30 Jan 2025 22:34:33 -0800 Subject: [PATCH] feat(stepfunctions): Support calling TestState API from Workflow Studio #6421 ## Problem The Workflow Studio webview currently does not allow for calling the [TestState API](https://docs.aws.amazon.com/step-functions/latest/dg/test-state-isolation.html). This API is used for testing individual states in isolation, and helps with debugging when constructing a state machine. It is available in the console version of Workflow Studio. ## Solution Adding support for calling APIs from the webview using message passing. This is the added flow: 1. The webview sends a message to the extension to call sfn:TestState or iam:ListRoles 2. The extension performs the API call using its credentials and default credential region 3. The extension sends the response as a message to the webview Note: this PR is dependent on [this PR](https://github.com/aws/aws-toolkit-vscode/pull/6375) being merged first since it requires an [aws-sdk version update](https://github.com/aws/aws-toolkit-vscode/pull/6375/commits/f4e0893d70dfe99b93988ea8eeb2caa4b210d8b3). --- packages/core/src/extensionNode.ts | 2 +- .../src/shared/clients/stepFunctionsClient.ts | 6 ++ packages/core/src/shared/logger/logger.ts | 9 +- .../workflowStudio/activation.ts | 10 +- .../workflowStudio/handleMessage.ts | 25 ++++- .../src/stepFunctions/workflowStudio/types.ts | 25 +++++ .../workflowStudioApiHandler.ts | 78 ++++++++++++++ .../workflowStudio/workflowStudioEditor.ts | 17 ++- .../workflowStudioEditorProvider.ts | 15 ++- .../workflowStudioApiHandler.test.ts | 101 ++++++++++++++++++ 10 files changed, 261 insertions(+), 27 deletions(-) create mode 100644 packages/core/src/stepFunctions/workflowStudio/workflowStudioApiHandler.ts create mode 100644 packages/core/src/test/stepFunctions/workflowStudio/workflowStudioApiHandler.test.ts diff --git a/packages/core/src/extensionNode.ts b/packages/core/src/extensionNode.ts index 8b25136faab..8481a5b066a 100644 --- a/packages/core/src/extensionNode.ts +++ b/packages/core/src/extensionNode.ts @@ -196,7 +196,7 @@ export async function activate(context: vscode.ExtensionContext) { await activateStepFunctions(context, globals.awsContext, globals.outputChannel) - await activateStepFunctionsWorkflowStudio(context) + await activateStepFunctionsWorkflowStudio() await activateRedshift(extContext) diff --git a/packages/core/src/shared/clients/stepFunctionsClient.ts b/packages/core/src/shared/clients/stepFunctionsClient.ts index 53036148c87..66d45bcd58a 100644 --- a/packages/core/src/shared/clients/stepFunctionsClient.ts +++ b/packages/core/src/shared/clients/stepFunctionsClient.ts @@ -67,6 +67,12 @@ export class DefaultStepFunctionsClient { return client.updateStateMachine(params).promise() } + public async testState(params: StepFunctions.TestStateInput): Promise { + const client = await this.createSdkClient() + + return await client.testState(params).promise() + } + private async createSdkClient(): Promise { return await globals.sdkClientBuilder.createAwsService(StepFunctions, undefined, this.regionCode) } diff --git a/packages/core/src/shared/logger/logger.ts b/packages/core/src/shared/logger/logger.ts index d5bf5f13380..de74bba3061 100644 --- a/packages/core/src/shared/logger/logger.ts +++ b/packages/core/src/shared/logger/logger.ts @@ -5,7 +5,14 @@ import * as vscode from 'vscode' -export type LogTopic = 'crashMonitoring' | 'dev/beta' | 'notifications' | 'test' | 'childProcess' | 'unknown' +export type LogTopic = + | 'crashMonitoring' + | 'dev/beta' + | 'notifications' + | 'test' + | 'childProcess' + | 'unknown' + | 'stepfunctions' class ErrorLog { constructor( diff --git a/packages/core/src/stepFunctions/workflowStudio/activation.ts b/packages/core/src/stepFunctions/workflowStudio/activation.ts index 376a5a04df9..7711c78faa3 100644 --- a/packages/core/src/stepFunctions/workflowStudio/activation.ts +++ b/packages/core/src/stepFunctions/workflowStudio/activation.ts @@ -6,16 +6,16 @@ import * as vscode from 'vscode' import { WorkflowStudioEditorProvider } from './workflowStudioEditorProvider' import { Commands } from '../../shared/vscode/commands2' +import { globals } from '../../shared' /** * Activates the extension and registers all necessary components. - * @param extensionContext The extension context object. */ -export async function activate(extensionContext: vscode.ExtensionContext): Promise { - extensionContext.subscriptions.push(WorkflowStudioEditorProvider.register(extensionContext)) +export async function activate(): Promise { + globals.context.subscriptions.push(WorkflowStudioEditorProvider.register()) // Open the file with Workflow Studio editor in a new tab, or focus on the tab with WFS if it is already open - extensionContext.subscriptions.push( + globals.context.subscriptions.push( Commands.register('aws.stepfunctions.openWithWorkflowStudio', async (uri: vscode.Uri) => { await vscode.commands.executeCommand('vscode.openWith', uri, WorkflowStudioEditorProvider.viewType) }) @@ -23,7 +23,7 @@ export async function activate(extensionContext: vscode.ExtensionContext): Promi // Close the active editor and open the file with Workflow Studio (or close and switch to the existing relevant tab). // This command is expected to always be called from the active tab in the default editor mode - extensionContext.subscriptions.push( + globals.context.subscriptions.push( Commands.register('aws.stepfunctions.switchToWorkflowStudio', async (uri: vscode.Uri) => { await vscode.commands.executeCommand('workbench.action.closeActiveEditor') await vscode.commands.executeCommand('vscode.openWith', uri, WorkflowStudioEditorProvider.viewType) diff --git a/packages/core/src/stepFunctions/workflowStudio/handleMessage.ts b/packages/core/src/stepFunctions/workflowStudio/handleMessage.ts index 597dd0b578d..7b8aade4612 100644 --- a/packages/core/src/stepFunctions/workflowStudio/handleMessage.ts +++ b/packages/core/src/stepFunctions/workflowStudio/handleMessage.ts @@ -13,6 +13,7 @@ import { FileChangedMessage, FileChangeEventTrigger, SyncFileRequestMessage, + ApiCallRequestMessage, } from './types' import { submitFeedback } from '../../feedback/vue/submitFeedback' import { placeholder } from '../../shared/vscode/commands2' @@ -20,6 +21,8 @@ import * as nls from 'vscode-nls' import vscode from 'vscode' import { telemetry } from '../../shared/telemetry/telemetry' import { ToolkitError } from '../../shared/errors' +import { WorkflowStudioApiHandler } from './workflowStudioApiHandler' +import { getLogger, globals } from '../../shared' const localize = nls.loadMessageBundle() /** @@ -48,6 +51,9 @@ export async function handleMessage(message: Message, context: WebviewContext) { case Command.OPEN_FEEDBACK: void submitFeedback(placeholder, 'Workflow Studio') break + case Command.API_CALL: + apiCallMessageHandler(message as ApiCallRequestMessage, context) + break } } else if (messageType === MessageType.BROADCAST) { switch (command) { @@ -150,7 +156,9 @@ async function saveFileMessageHandler(request: SaveFileRequestMessage, context: ) ) } catch (err) { - throw ToolkitError.chain(err, 'Could not save asl file.', { code: 'SaveFailed' }) + throw ToolkitError.chain(err, 'Could not save asl file.', { + code: 'SaveFailed', + }) } }) } @@ -179,7 +187,20 @@ async function autoSyncFileMessageHandler(request: SyncFileRequestMessage, conte ) await vscode.workspace.applyEdit(edit) } catch (err) { - throw ToolkitError.chain(err, 'Could not autosave asl file.', { code: 'AutoSaveFailed' }) + throw ToolkitError.chain(err, 'Could not autosave asl file.', { + code: 'AutoSaveFailed', + }) } }) } + +/** + * Handler for making API calls from the webview and returning the response. + * @param request The request message containing the API to call and the parameters + * @param context The webview context used for returning the API response to the webview + */ +function apiCallMessageHandler(request: ApiCallRequestMessage, context: WebviewContext) { + const logger = getLogger('stepfunctions') + const apiHandler = new WorkflowStudioApiHandler(globals.awsContext.getCredentialDefaultRegion(), context) + apiHandler.performApiCall(request).catch((error) => logger.error('%s API call failed: %O', request.apiName, error)) +} diff --git a/packages/core/src/stepFunctions/workflowStudio/types.ts b/packages/core/src/stepFunctions/workflowStudio/types.ts index 668b4ace8ab..01748b8a1b0 100644 --- a/packages/core/src/stepFunctions/workflowStudio/types.ts +++ b/packages/core/src/stepFunctions/workflowStudio/types.ts @@ -2,6 +2,7 @@ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. * SPDX-License-Identifier: Apache-2.0 */ +import { IAM, StepFunctions } from 'aws-sdk' import * as vscode from 'vscode' export type WebviewContext = { @@ -39,6 +40,7 @@ export enum Command { LOAD_STAGE = 'LOAD_STAGE', OPEN_FEEDBACK = 'OPEN_FEEDBACK', CLOSE_WFS = 'CLOSE_WFS', + API_CALL = 'API_CALL', } export type FileWatchInfo = { @@ -71,3 +73,26 @@ export interface SaveFileRequestMessage extends Message { export interface SyncFileRequestMessage extends SaveFileRequestMessage { fileContents: string } + +export enum ApiAction { + IAMListRoles = 'iam:ListRoles', + SFNTestState = 'sfn:TestState', +} + +type ApiCallRequestMapping = { + [ApiAction.IAMListRoles]: IAM.ListRolesRequest + [ApiAction.SFNTestState]: StepFunctions.TestStateInput +} + +interface ApiCallRequestMessageBase extends Message { + requestId: string + apiName: ApiName + params: ApiCallRequestMapping[ApiName] +} + +/** + * The message from the webview describing what API and parameters to call. + */ +export type ApiCallRequestMessage = + | ApiCallRequestMessageBase + | ApiCallRequestMessageBase diff --git a/packages/core/src/stepFunctions/workflowStudio/workflowStudioApiHandler.ts b/packages/core/src/stepFunctions/workflowStudio/workflowStudioApiHandler.ts new file mode 100644 index 00000000000..c353a7de3ad --- /dev/null +++ b/packages/core/src/stepFunctions/workflowStudio/workflowStudioApiHandler.ts @@ -0,0 +1,78 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { IAM, StepFunctions } from 'aws-sdk' +import { DefaultIamClient } from '../../shared/clients/iamClient' +import { DefaultStepFunctionsClient } from '../../shared/clients/stepFunctionsClient' +import { ApiAction, ApiCallRequestMessage, Command, MessageType, WebviewContext } from './types' +import { telemetry } from '../../shared/telemetry/telemetry' + +export class WorkflowStudioApiHandler { + public constructor( + region: string, + private readonly context: WebviewContext, + private readonly clients = { + sfn: new DefaultStepFunctionsClient(region), + iam: new DefaultIamClient(region), + } + ) {} + + /** + * Performs the API call on behalf of the webview, and sends the sucesss or error response to the webview. + */ + public async performApiCall({ apiName, params, requestId }: ApiCallRequestMessage): Promise { + try { + let response + switch (apiName) { + case ApiAction.IAMListRoles: + response = await this.listRoles(params) + break + case ApiAction.SFNTestState: + response = await this.testState(params) + break + default: + throw new Error(`Unknown API: ${apiName}`) + } + + await this.context.panel.webview.postMessage({ + messageType: MessageType.RESPONSE, + command: Command.API_CALL, + apiName, + response, + requestId, + isSuccess: true, + }) + } catch (err) { + await this.context.panel.webview.postMessage({ + messageType: MessageType.RESPONSE, + command: Command.API_CALL, + apiName, + error: + err instanceof Error + ? { + message: err.message, + name: err.name, + stack: err.stack, + } + : { + message: String(err), + }, + requestId, + isSuccess: false, + }) + } + } + + public async testState(params: StepFunctions.TestStateInput): Promise { + telemetry.ui_click.emit({ + elementId: 'stepfunctions_testState', + }) + return this.clients.sfn.testState(params) + } + + public async listRoles(params: IAM.ListRolesRequest): Promise { + return this.clients.iam.listRoles(params) + } +} diff --git a/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditor.ts b/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditor.ts index 09ed19ccd19..c0acebc698a 100644 --- a/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditor.ts +++ b/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditor.ts @@ -13,6 +13,7 @@ import { CancellationError } from '../../shared/utilities/timeoutUtils' import { handleMessage } from './handleMessage' import { isInvalidJsonFile } from '../utils' import { setContext } from '../../shared/vscode/setContext' +import { globals } from '../../shared' /** * The main class for Workflow Studio Editor. This class handles the creation and management @@ -37,7 +38,6 @@ export class WorkflowStudioEditor { public constructor( textDocument: vscode.TextDocument, webviewPanel: vscode.WebviewPanel, - context: vscode.ExtensionContext, fileId: string, getWebviewContent: () => Promise ) { @@ -54,7 +54,7 @@ export class WorkflowStudioEditor { id: this.fileId, }) - this.setupWebviewPanel(textDocument, context) + this.setupWebviewPanel(textDocument) } public get onVisualizationDisposeEvent(): vscode.Event { @@ -70,11 +70,11 @@ export class WorkflowStudioEditor { this.getPanel()?.reveal() } - public async refreshPanel(context: vscode.ExtensionContext) { + public async refreshPanel() { if (!this.isPanelDisposed) { this.webviewPanel.dispose() const document = await vscode.workspace.openTextDocument(this.documentUri) - this.setupWebviewPanel(document, context) + this.setupWebviewPanel(document) } } @@ -87,10 +87,9 @@ export class WorkflowStudioEditor { * panel, setting up the webview content, and handling the communication between the webview * and the extension context. * @param textDocument The text document to be displayed in the webview panel. - * @param context The extension context. * @private */ - private setupWebviewPanel(textDocument: vscode.TextDocument, context: vscode.ExtensionContext) { + private setupWebviewPanel(textDocument: vscode.TextDocument) { const documentUri = textDocument.uri const contextObject: WebviewContext = { @@ -131,7 +130,7 @@ export class WorkflowStudioEditor { // Initialise webview panel for Workflow Studio and set up initial content this.webviewPanel.webview.options = { enableScripts: true, - localResourceRoots: [context.extensionUri], + localResourceRoots: [globals.context.extensionUri], } // Set the initial html for the webpage @@ -183,9 +182,9 @@ export class WorkflowStudioEditor { this.isPanelDisposed = true resolve() this.onVisualizationDisposeEmitter.fire() - this.disposables.forEach((disposable) => { + for (const disposable of this.disposables) { disposable.dispose() - }) + } this.onVisualizationDisposeEmitter.dispose() }) ) diff --git a/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditorProvider.ts b/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditorProvider.ts index 1defe10ea48..54d641df1dd 100644 --- a/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditorProvider.ts +++ b/packages/core/src/stepFunctions/workflowStudio/workflowStudioEditorProvider.ts @@ -35,8 +35,8 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv * @remarks This should only be called once per extension. * @param context The extension context */ - public static register(context: vscode.ExtensionContext): vscode.Disposable { - const provider = new WorkflowStudioEditorProvider(context) + public static register(): vscode.Disposable { + const provider = new WorkflowStudioEditorProvider() return vscode.window.registerCustomEditorProvider(WorkflowStudioEditorProvider.viewType, provider, { webviewOptions: { enableFindWidget: true, @@ -45,13 +45,11 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv }) } - protected extensionContext: vscode.ExtensionContext protected webviewHtml: string protected readonly managedVisualizations = new Map() protected readonly logger = getLogger() - constructor(context: vscode.ExtensionContext) { - this.extensionContext = context + constructor() { this.webviewHtml = '' } @@ -65,7 +63,7 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv this.webviewHtml = await response.text() for (const visualization of this.managedVisualizations.values()) { - await visualization.refreshPanel(this.extensionContext) + await visualization.refreshPanel() } } @@ -98,7 +96,7 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv htmlFileSplit = html.split('') const script = await fs.readFileText( - vscode.Uri.joinPath(this.extensionContext.extensionUri, 'resources', 'js', 'vsCodeExtensionInterface.js') + vscode.Uri.joinPath(globals.context.extensionUri, 'resources', 'js', 'vsCodeExtensionInterface.js') ) return `${htmlFileSplit[0]} ${htmlFileSplit[1]}` @@ -151,7 +149,6 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv const newVisualization = new WorkflowStudioEditor( document, webviewPanel, - this.extensionContext, fileId, this.getWebviewContent ) @@ -171,6 +168,6 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv const visualizationDisposable = visualization.onVisualizationDisposeEvent(() => { this.managedVisualizations.delete(key) }) - this.extensionContext.subscriptions.push(visualizationDisposable) + globals.context.subscriptions.push(visualizationDisposable) } } diff --git a/packages/core/src/test/stepFunctions/workflowStudio/workflowStudioApiHandler.test.ts b/packages/core/src/test/stepFunctions/workflowStudio/workflowStudioApiHandler.test.ts new file mode 100644 index 00000000000..7eb8912545c --- /dev/null +++ b/packages/core/src/test/stepFunctions/workflowStudio/workflowStudioApiHandler.test.ts @@ -0,0 +1,101 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import assert from 'assert' +import sinon from 'sinon' +import { WorkflowStudioApiHandler } from '../../../stepFunctions/workflowStudio/workflowStudioApiHandler' +import { MockDocument } from '../../fake/fakeDocument' +import { ApiAction, Command, MessageType, WebviewContext } from '../../../stepFunctions/workflowStudio/types' +import * as vscode from 'vscode' +import { assertTelemetry } from '../../testUtil' +import { DefaultStepFunctionsClient } from '../../../shared/clients/stepFunctionsClient' +import { DefaultIamClient } from '../../../shared/clients/iamClient' + +describe('WorkflowStudioApiHandler', function () { + let postMessageStub: sinon.SinonStub + let apiHandler: WorkflowStudioApiHandler + let testState: sinon.SinonStub + + async function assertTestApiResponse(expectedResponse: any) { + await apiHandler.performApiCall({ + apiName: ApiAction.SFNTestState, + params: { + definition: '', + roleArn: '', + }, + requestId: 'test-request-id', + command: Command.API_CALL, + messageType: MessageType.REQUEST, + }) + + assertTelemetry('ui_click', { + elementId: 'stepfunctions_testState', + }) + assert(postMessageStub.firstCall.calledWithExactly(expectedResponse)) + } + + beforeEach(() => { + const panel = vscode.window.createWebviewPanel('WorkflowStudioMock', 'WorkflowStudioMockTitle', { + viewColumn: vscode.ViewColumn.Active, + preserveFocus: true, + }) + + postMessageStub = sinon.stub(panel.webview, 'postMessage') + + const context: WebviewContext = { + defaultTemplateName: '', + defaultTemplatePath: '', + disposables: [], + panel, + textDocument: new MockDocument('', 'foo', async () => true), + workSpacePath: '', + fileStates: {}, + loaderNotification: undefined, + fileId: '', + } + + const sfnClient = new DefaultStepFunctionsClient('us-east-1') + apiHandler = new WorkflowStudioApiHandler('us-east-1', context, { + sfn: sfnClient, + iam: new DefaultIamClient('us-east-1'), + }) + + testState = sinon.stub(sfnClient, 'testState') + }) + + it('should handle request and response for success', async function () { + testState.resolves({ + output: 'Test state output', + }) + + await assertTestApiResponse({ + messageType: MessageType.RESPONSE, + command: Command.API_CALL, + apiName: ApiAction.SFNTestState, + response: { + output: 'Test state output', + }, + requestId: 'test-request-id', + isSuccess: true, + }) + }) + + it('should handle request and response for error', async function () { + testState.rejects(new Error('Error testing state')) + + await assertTestApiResponse({ + messageType: MessageType.RESPONSE, + command: Command.API_CALL, + apiName: ApiAction.SFNTestState, + error: { + message: 'Error testing state', + name: 'Error', + stack: sinon.match.string, + }, + isSuccess: false, + requestId: 'test-request-id', + }) + }) +})