fix(stepfunctions): Prevent opening Workflow Studio for invalid JSON/YAML; add save and deploy message handler

Author: Vlad Nikolaenko

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

@@ -45,6 +45,9 @@ export async function handleMessage(message: Message, context: WebviewContext) {
             case Command.SAVE_FILE:
                 void saveFileMessageHandler(message as SaveFileRequestMessage, context)
                 break
+            case Command.SAVE_FILE_AND_DEPLOY:
+                void saveFileAndDeployMessageHandler(message as SaveFileRequestMessage, context)
+                break
             case Command.AUTO_SYNC_FILE:
                 void autoSyncFileMessageHandler(message as SyncFileRequestMessage, context)
                 break
@@ -132,13 +135,13 @@ async function loadStageMessageHandler(context: WebviewContext) {
  * Handler for closing WFS custom editor. When called, disposes webview panel and opens default VSCode editor
  * @param context The context object containing the necessary information for the webview.
  */
-export function closeCustomEditorMessageHandler(context: WebviewContext) {
-    telemetry.stepfunctions_closeWorkflowStudio.run((span) => {
+export async function closeCustomEditorMessageHandler(context: WebviewContext) {
+    await telemetry.stepfunctions_closeWorkflowStudio.run(async (span) => {
         span.record({
             id: context.fileId,
         })
         context.panel.dispose()
-        void vscode.commands.executeCommand('vscode.openWith', context.textDocument.uri, 'default')
+        await vscode.commands.executeCommand('vscode.openWith', context.textDocument.uri, 'default')
     })
 }
 
@@ -174,6 +177,18 @@ async function saveFileMessageHandler(request: SaveFileRequestMessage, context:
     })
 }
 
+/**
+ * Handler for saving a file and starting the state machine deployment flow, while also switching to default editor.
+ * Triggered when the user triggers 'Save and Deploy' action in WFS
+ * @param request The request message containing the file contents.
+ * @param context The webview context containing the necessary information for saving the file.
+ */
+async function saveFileAndDeployMessageHandler(request: SaveFileRequestMessage, context: WebviewContext) {
+    await saveFileMessageHandler(request, context)
+    await closeCustomEditorMessageHandler(context)
+    void vscode.commands.executeCommand('aws.stepfunctions.publishStateMachine')
+}
+
 /**
  * Handler for auto syncing a file from the webview which updates the workspace but does not save the file.
  * Triggered on every code change from WFS, including invalid JSON.

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

@@ -35,6 +35,7 @@ export enum MessageType {
 export enum Command {
     INIT = 'INIT',
     SAVE_FILE = 'SAVE_FILE',
+    SAVE_FILE_AND_DEPLOY = 'SAVE_FILE_AND_DEPLOY',
     AUTO_SYNC_FILE = 'AUTO_SYNC_FILE',
     FILE_CHANGED = 'FILE_CHANGED',
     LOAD_STAGE = 'LOAD_STAGE',

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

@@ -30,11 +30,49 @@ let clientId = ''
 export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProvider {
     public static readonly viewType = 'workflowStudio.asl'
 
+    /**
+     * Opens a file in the Workflow Studio editor, after validating the document.
+     * @param uri The URI of the document to open in the Workflow Studio editor.
+     * @param params Optional parameters to customize the WebView panel.
+     */
     public static async openWithWorkflowStudio(
         uri: vscode.Uri,
         params?: Parameters<typeof vscode.window.createWebviewPanel>[2]
     ) {
-        await vscode.commands.executeCommand('vscode.openWith', uri, WorkflowStudioEditorProvider.viewType, params)
+        const document = await vscode.workspace.openTextDocument(uri)
+        await telemetry.stepfunctions_openWorkflowStudio.run(async () => {
+            await this.validateAndThrowIfInvalid(document)
+            await vscode.commands.executeCommand('vscode.openWith', uri, WorkflowStudioEditorProvider.viewType, params)
+        })
+    }
+
+    /**
+     * Validates the document to ensure it is either valid JSON or YAML.
+     * If the document is invalid, a warning is shown (if specified) and an error is thrown to stop further execution.
+     * @param document The document to validate.
+     * @param alertOnError If true, shows a warning message when the document is invalid. Defaults to `false`.
+     */
+    private static async validateAndThrowIfInvalid(
+        document: vscode.TextDocument,
+        alertOnError?: boolean
+    ): Promise<void> {
+        const isInvalidJson = isInvalidJsonFile(document)
+        const isInvalidYaml = isInvalidYamlFile(document)
+
+        if (isInvalidJson || isInvalidYaml) {
+            const language = isInvalidJson ? 'JSON' : 'YAML'
+            const errorKey = isInvalidJson ? 'InvalidJSONContent' : 'InvalidYAMLContent'
+
+            if (alertOnError) {
+                void vscode.window.showErrorMessage(i18n(`AWS.stepFunctions.workflowStudio.actions.${errorKey}`))
+            }
+
+            throw ToolkitError.chain(
+                `Invalid ${language} file`,
+                `The Workflow Studio editor was not opened because the ${language} in the file is invalid`,
+                { code: errorKey }
+            )
+        }
     }
 
     /**
@@ -114,20 +152,12 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv
                 webviewPanel.dispose()
             }
 
-            const isInvalidJson = isInvalidJsonFile(document)
-            const isInvalidYaml = isInvalidYamlFile(document)
-
-            if (isInvalidJson || isInvalidYaml) {
-                const language = isInvalidJson ? 'JSON' : 'YAML'
-                const errorKey = isInvalidJson ? 'InvalidJSONContent' : 'InvalidYAMLContent'
-
+            try {
+                await WorkflowStudioEditorProvider.validateAndThrowIfInvalid(document, true)
+            } catch (e) {
+                // If the document is invalid, reopen it with the default editor and re-throw the error
                 await reopenWithDefaultEditor()
-                void vscode.window.showWarningMessage(i18n(`AWS.stepFunctions.workflowStudio.actions.${errorKey}`))
-                throw ToolkitError.chain(
-                    `Invalid ${language} file`,
-                    `The Workflow Studio editor was not opened because the ${language} in the file is invalid`,
-                    { code: errorKey }
-                )
+                throw e
             }
 
             if (!this.webviewHtml) {