feat(stepfunctions): Miltiple Workflow Studio integration improvements and fixes

Author: Vlad Nikolaenko

packages/core/package.nls.json

@@ -26,7 +26,9 @@
     "AWS.stepFunctions.asl.maxItemsComputed.desc": "The maximum number of outline symbols and folding regions computed (limited for performance reasons).",
     "AWS.stepFunctions.workflowStudio.actions.progressMessage": "Opening asl file in Workflow Studio",
     "AWS.stepFunctions.workflowStudio.actions.saveSuccessMessage": "{0} has been saved",
-    "AWS.stepFunctions.workflowStudio.actions.invalidJson": "The Workflow Studio editor was not opened because the JSON in the file is invalid. To access Workflow Studio, please fix the JSON and manually reopen the integration.",
+    "AWS.stepFunctions.workflowStudio.actions.InvalidJSONContent": "The Workflow Studio editor was not opened because the JSON in the file is invalid. To access Workflow Studio, please fix the JSON and manually reopen the integration.",
+    "AWS.stepFunctions.workflowStudio.actions.InvalidYAMLContent": "The Workflow Studio editor was not opened because the YAML in the file is invalid. To access Workflow Studio, please fix the YAML and manually reopen the integration.",
+    "AWS.stepFunctions.workflowStudio.actions.webviewFetchFailed": "Failed to load Workflow Studio editor. Please check your network connection and try again.",
     "AWS.configuration.description.awssam.debug.api": "API Gateway configuration",
     "AWS.configuration.description.awssam.debug.api.clientCertId": "The API Gateway client certificate ID",
     "AWS.configuration.description.awssam.debug.api.headers": "Additional HTTP headers",

packages/core/src/stepFunctions/utils.ts

@@ -6,6 +6,7 @@ import * as nls from 'vscode-nls'
 const localize = nls.loadMessageBundle()
 
 import { IAM, StepFunctions } from 'aws-sdk'
+import * as yaml from 'js-yaml'
 import * as vscode from 'vscode'
 import { StepFunctionsClient } from '../shared/clients/stepFunctionsClient'
 import { fileExists } from '../shared/filesystemUtilities'
@@ -232,7 +233,7 @@ export async function isDocumentValid(text: string, textDocument?: vscode.TextDo
 }
 
 /**
- * Checks if the JSON content in a text document is invalid.
+ * Checks if the JSON content in an ASL text document is invalid.
  * Returns `true` for invalid JSON; `false` for valid JSON, empty content, or non-JSON files.
  *
  * @param textDocument - The text document to check.
@@ -241,9 +242,25 @@ export async function isDocumentValid(text: string, textDocument?: vscode.TextDo
 export const isInvalidJsonFile = (textDocument: vscode.TextDocument): boolean => {
     const documentContent = textDocument.getText().trim()
     // An empty file or whitespace-only text is considered valid JSON for our use case
-    return textDocument.fileName.toLowerCase().endsWith('.json') && documentContent
-        ? isInvalidJson(documentContent)
-        : false
+    return textDocument.languageId === 'asl' && documentContent ? isInvalidJson(documentContent) : false
+}
+
+/**
+ * Checks if the YAML content in an ASL text document is invalid.
+ * Returns `true` for invalid YAML; `false` for valid YAML, empty content, or non-YAML files.
+ *
+ * @param textDocument - The text document to check.
+ * @returns `true` if invalid; `false` otherwise.
+ */
+export const isInvalidYamlFile = (textDocument: vscode.TextDocument): boolean => {
+    try {
+        if (textDocument.languageId === 'asl-yaml') {
+            yaml.load(textDocument.getText())
+        }
+        return false
+    } catch {
+        return true
+    }
 }
 
 const isInvalidJson = (content: string): boolean => {

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

@@ -12,9 +12,10 @@ import { telemetry } from '../../shared/telemetry/telemetry'
 import globals from '../../shared/extensionGlobals'
 import { getRandomString, getStringHash } from '../../shared/utilities/textUtilities'
 import { ToolkitError } from '../../shared/errors'
+import { getTabSizeSetting } from '../../shared/utilities/editorUtilities'
 import { WorkflowStudioEditor } from './workflowStudioEditor'
 import { i18n } from '../../shared/i18n-helper'
-import { isInvalidJsonFile } from '../utils'
+import { isInvalidJsonFile, isInvalidYamlFile } from '../utils'
 
 const isLocalDev = false
 const localhost = 'http://127.0.0.1:3002'
@@ -72,9 +73,6 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv
      * @private
      */
     private getWebviewContent = async () => {
-        if (!this.webviewHtml) {
-            await this.fetchWebviewHtml()
-        }
         let htmlFileSplit = this.webviewHtml.split('<head>')
 
         // Set asset source to CDN
@@ -86,13 +84,15 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv
         const localeTag = `<meta name='locale' content='${locale}'>`
         const theme = vscode.window.activeColorTheme.kind
         const isDarkMode = theme === vscode.ColorThemeKind.Dark || theme === vscode.ColorThemeKind.HighContrast
+        const tabSizeTag = `<meta name='tab-size' content='${getTabSizeSetting()}'>`
         const darkModeTag = `<meta name='dark-mode' content='${isDarkMode}'>`
-        let html = `${htmlFileSplit[0]} <head> ${baseTag} ${localeTag} ${darkModeTag} ${htmlFileSplit[1]}`
+        let html = `${htmlFileSplit[0]} <head> ${baseTag} ${localeTag} ${darkModeTag} ${tabSizeTag} ${htmlFileSplit[1]}`
 
         const nonce = getRandomString()
+        const localDevURL = isLocalDev ? localhost : ''
         htmlFileSplit = html.split("script-src 'self'")
 
-        html = `${htmlFileSplit[0]} script-src 'self' 'nonce-${nonce}' ${isLocalDev && localhost} ${htmlFileSplit[1]}`
+        html = `${htmlFileSplit[0]} script-src 'self' 'nonce-${nonce}' ${localDevURL} ${htmlFileSplit[1]}`
         htmlFileSplit = html.split('<body>')
 
         const script = await fs.readFileText(
@@ -115,35 +115,56 @@ export class WorkflowStudioEditorProvider implements vscode.CustomTextEditorProv
         _token: vscode.CancellationToken
     ): Promise<void> {
         await telemetry.stepfunctions_openWorkflowStudio.run(async () => {
-            // For invalid JSON, open default editor and show warning message
-            if (isInvalidJsonFile(document)) {
+            const reopenWithDefaultEditor = async () => {
                 await vscode.commands.executeCommand('vscode.openWith', document.uri, 'default')
                 webviewPanel.dispose()
-                void vscode.window.showWarningMessage(i18n('AWS.stepFunctions.workflowStudio.actions.invalidJson'))
+            }
+
+            const isInvalidJson = isInvalidJsonFile(document)
+            const isInvalidYaml = isInvalidYamlFile(document)
 
+            if (isInvalidJson || isInvalidYaml) {
+                const language = isInvalidJson ? 'JSON' : 'YAML'
+                const errorKey = isInvalidJson ? 'InvalidJSONContent' : 'InvalidYAMLContent'
+
+                await reopenWithDefaultEditor()
+                void vscode.window.showWarningMessage(i18n(`AWS.stepFunctions.workflowStudio.actions.${errorKey}`))
                 throw ToolkitError.chain(
-                    'Invalid JSON file',
-                    'The Workflow Studio editor was not opened because the JSON in the file is invalid',
-                    {
-                        code: 'InvalidJSONContent',
-                    }
+                    `Invalid ${language} file`,
+                    `The Workflow Studio editor was not opened because the ${language} in the file is invalid`,
+                    { code: errorKey }
                 )
             }
 
             if (!this.webviewHtml) {
-                await this.fetchWebviewHtml()
+                try {
+                    await this.fetchWebviewHtml()
+                } catch (e) {
+                    await reopenWithDefaultEditor()
+
+                    void vscode.window.showWarningMessage(
+                        i18n('AWS.stepFunctions.workflowStudio.actions.webviewFetchFailed')
+                    )
+                    throw ToolkitError.chain(
+                        'Failed to fetch editor content',
+                        'Could not retrieve content for the Workflow Studio editor',
+                        {
+                            code: 'webviewFetchFailed',
+                        }
+                    )
+                }
             }
 
             if (clientId === '') {
                 clientId = getClientId(globals.globalState)
             }
-            // Attempt to retrieve existing visualization if it exists.
-            const existingVisualization = this.managedVisualizations.get(document.uri.fsPath)
 
+            const existingVisualization = this.managedVisualizations.get(document.uri.fsPath)
             if (existingVisualization) {
-                existingVisualization.showPanel()
+                // Prevent opening multiple custom editors for a single file
+                await reopenWithDefaultEditor()
             } else {
-                // Existing visualization does not exist, construct new visualization
+                // Construct new visualization
                 try {
                     const fileId = getStringHash(`${document.uri.fsPath}${clientId}`)
                     const newVisualization = new WorkflowStudioEditor(

packages/toolkit/package.json

@@ -1412,6 +1412,11 @@
                     "command": "aws.openInApplicationComposer",
                     "when": "isFileSystemResource && !(resourceFilename =~ /^.*\\.tc\\.json$/) && resourceFilename =~ /^.*\\.(json|yml|yaml|template)$/",
                     "group": "z_aws@1"
+                },
+                {
+                    "command": "aws.stepfunctions.openWithWorkflowStudio",
+                    "when": "isFileSystemResource && resourceFilename =~ /^.*\\.asl\\.(json|yml|yaml)$/",
+                    "group": "z_aws@1"
                 }
             ],
             "view/item/context": [
@@ -3917,6 +3922,16 @@
                     }
                 }
             },
+            {
+                "command": "aws.stepfunctions.openWithWorkflowStudio",
+                "title": "%AWS.command.stepFunctions.openWithWorkflowStudio%",
+                "category": "%AWS.title%",
+                "cloud9": {
+                    "cn": {
+                        "category": "%AWS.title.cn%"
+                    }
+                }
+            },
             {
                 "command": "aws.createNewThreatComposer",
                 "title": "%AWS.command.threatComposer.createNew%",
@@ -4654,7 +4669,8 @@
         ],
         "configurationDefaults": {
             "workbench.editorAssociations": {
-                "{git,gitlens,conflictResolution,vscode-local-history}:/**/*.tc.json": "default"
+                "{git,gitlens,conflictResolution,vscode-local-history}:/**/*.tc.json": "default",
+                "{git,gitlens,conflictResolution,vscode-local-history}:/**/*.{asl.json,asl.yaml,asl.yml}": "default"
             }
         }
     },