feat(stepfunctions): allow closing Workflow Studio and reopening it

Author: Vlad Nikolaenko

packages/core/package.json

@@ -1397,6 +1397,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": [
@@ -3579,6 +3584,16 @@
                         "category": "%AWS.title.cn%"
                     }
                 }
+            },
+            {
+                "command": "aws.stepfunctions.openWithWorkflowStudio",
+                "title": "%AWS.command.stepFunctions.openWithWorkflowStudio%",
+                "category": "%AWS.title%",
+                "cloud9": {
+                    "cn": {
+                        "category": "%AWS.title.cn%"
+                    }
+                }
             }
         ],
         "jsonValidation": [

packages/core/package.nls.json

@@ -185,6 +185,7 @@
     "AWS.command.stepFunctions.createStateMachineFromTemplate": "Create a new Step Functions state machine",
     "AWS.command.stepFunctions.publishStateMachine": "Publish state machine to Step Functions",
     "AWS.command.stepFunctions.previewStateMachine": "Render state machine graph",
+    "AWS.command.stepFunctions.openWithWorkflowStudio": "Open with Workflow Studio",
     "AWS.command.cdk.previewStateMachine": "Render state machine graph from CDK application",
     "AWS.command.copyLogResource": "Copy Log Stream or Group",
     "AWS.command.saveCurrentLogDataContent": "Save Log to File",

packages/core/src/shared/telemetry/vscodeTelemetry.json

@@ -516,6 +516,16 @@
                 }
             ]
         },
+        {
+            "name": "stepfunctions_closeWorkflowStudio",
+            "description": "Called after closing Step Functions Workflow Studio custom editor",
+            "metadata": [
+                {
+                    "type": "id",
+                    "required": true
+                }
+            ]
+        },
         {
             "name": "stepfunctions_saveFile",
             "description": "Called after the user saves local ASL file (inlcuding autosave) from VSCode editor or Workflow Studio",

packages/core/src/stepFunctions/activation.ts

@@ -144,6 +144,13 @@ function initializeCodeLens(context: vscode.ExtensionContext) {
         public async provideCodeLenses(document: vscode.TextDocument): Promise<vscode.CodeLens[]> {
             const topOfDocument = new vscode.Range(0, 0, 0, 0)
 
+            const openCustomEditorCommand: vscode.Command = {
+                command: 'aws.stepfunctions.switchToWorkflowStudio',
+                title: localize('AWS.command.stepFunctions.openWithWorkflowStudio', 'Open with Workflow Studio'),
+                arguments: [document.uri],
+            }
+            const openCustomEditor = new vscode.CodeLens(topOfDocument, openCustomEditorCommand)
+
             const renderCodeLens = previewStateMachineCommand.build().asCodeLens(topOfDocument, {
                 title: localize('AWS.stepFunctions.render', 'Render graph'),
             })
@@ -155,9 +162,9 @@ function initializeCodeLens(context: vscode.ExtensionContext) {
                 }
                 const publishCodeLens = new vscode.CodeLens(topOfDocument, publishCommand)
 
-                return [publishCodeLens, renderCodeLens]
+                return [openCustomEditor, publishCodeLens, renderCodeLens]
             } else {
-                return [renderCodeLens]
+                return [openCustomEditor, renderCodeLens]
             }
         }
     }

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

@@ -5,11 +5,28 @@
 
 import * as vscode from 'vscode'
 import { WorkflowStudioEditorProvider } from './workflowStudioEditorProvider'
+import { Commands } from '../../shared/vscode/commands2'
 
 /**
  * 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))
+
+    // 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(
+        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(
+        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

@@ -41,6 +41,9 @@ export async function handleMessage(message: Message, context: WebviewContext) {
             case Command.AUTO_SAVE_FILE:
                 void autoSaveFileMessageHandler(message as SaveFileRequestMessage, context)
                 break
+            case Command.CLOSE_WFS:
+                void closeCustomEditorMessageHandler(context)
+                break
             case Command.OPEN_FEEDBACK:
                 void submitFeedback(placeholder, 'Workflow Studio')
                 break
@@ -107,6 +110,20 @@ async function loadStageMessageHandler(context: WebviewContext) {
     }, 100)
 }
 
+/**
+ * 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) => {
+        span.record({
+            id: context.fileId,
+        })
+        context.panel.dispose()
+        void vscode.commands.executeCommand('vscode.openWith', context.textDocument.uri, 'default')
+    })
+}
+
 /**
  * Handler for saving a file from the webview which updates the workspace and saves the file.
  * Triggered when the user explicitly applies save action in WFS

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

@@ -38,6 +38,7 @@ export enum Command {
     FILE_CHANGED = 'FILE_CHANGED',
     LOAD_STAGE = 'LOAD_STAGE',
     OPEN_FEEDBACK = 'OPEN_FEEDBACK',
+    CLOSE_WFS = 'CLOSE_WFS',
 }
 
 export type FileWatchInfo = {

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

@@ -6,7 +6,6 @@
 import * as path from 'path'
 import * as vscode from 'vscode'
 import { telemetry } from '../../shared/telemetry/telemetry'
-import { getLogger } from '../../shared/logger'
 import { i18n } from '../../shared/i18n-helper'
 import { broadcastFileChange } from './handleMessage'
 import { FileWatchInfo, WebviewContext } from './types'
@@ -113,7 +112,6 @@ export class WorkflowStudioEditor {
             async (progress, token) => {
                 token.onCancellationRequested(async () => {
                     // Cancel opening in Worflow Studio and open regular code editor instead
-                    getLogger().debug('WorkflowStudio: Canceled opening')
                     contextObject.panel.dispose()
                     await vscode.commands.executeCommand('vscode.openWith', documentUri, 'default')
                     throw new CancellationError('user')