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](#6375) being merged
first since it requires an [aws-sdk version
update](f4e0893).

Author: anthonyting

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)
 

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<StepFunctions.TestStateOutput> {
+        const client = await this.createSdkClient()
+
+        return await client.testState(params).promise()
+    }
+
     private async createSdkClient(): Promise<StepFunctions> {
         return await globals.sdkClientBuilder.createAwsService(StepFunctions, undefined, this.regionCode)
     }

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(

packages/core/src/stepFunctions/workflowStudio/activation.ts

@@ -6,24 +6,24 @@
 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<void> {
-    extensionContext.subscriptions.push(WorkflowStudioEditorProvider.register(extensionContext))
+export async function activate(): Promise<void> {
+    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)
         })
     )
 
     // 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)

packages/core/src/stepFunctions/workflowStudio/handleMessage.ts

@@ -13,13 +13,16 @@ import {
     FileChangedMessage,
     FileChangeEventTrigger,
     SyncFileRequestMessage,
+    ApiCallRequestMessage,
 } from './types'
 import { submitFeedback } from '../../feedback/vue/submitFeedback'
 import { placeholder } from '../../shared/vscode/commands2'
 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))
+}

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<ApiName extends ApiAction> 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<ApiAction.IAMListRoles>
+    | ApiCallRequestMessageBase<ApiAction.SFNTestState>

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<void> {
+        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<StepFunctions.TestStateOutput> {
+        telemetry.ui_click.emit({
+            elementId: 'stepfunctions_testState',
+        })
+        return this.clients.sfn.testState(params)
+    }
+
+    public async listRoles(params: IAM.ListRolesRequest): Promise<IAM.Role[]> {
+        return this.clients.iam.listRoles(params)
+    }
+}

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<string>
     ) {
@@ -54,7 +54,7 @@ export class WorkflowStudioEditor {
             id: this.fileId,
         })
 
-        this.setupWebviewPanel(textDocument, context)
+        this.setupWebviewPanel(textDocument)
     }
 
     public get onVisualizationDisposeEvent(): vscode.Event<void> {
@@ -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()
                         })
                     )

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<string, WorkflowStudioEditor>()
     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('<body>')
 
         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]} <body> <script nonce='${nonce}'>${script}</script> ${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)
     }
 }