refactor(stepfunctions): migrate to sdkv3 (#7560)

zelzhou · web-flow · commit e07833d92eb0 · 2025-06-27T14:07:03.000-04:00
## Problem step functions still uses sdk v2 ## Solution - Refactor `DefaultStepFunctionsClient` to `StepFunctionsClient` using ClientWrapper. - Refactor references to `DefaultStepFunctionsClient` with `StepFunctionsClient` and reference sdk v3 types ## Verification - Verified CreateStateMachine, UpdateStateMachine works and creates/updates StateMachine - Verified AWS Explorer lists all StateMachines using ListStateMachine - Verified StartExecution works - Verified TestState, ListIAMRoles works in Workflow Studio, which also fixes the bug that variables cannot be used in TestState in Workflow Studio - fixes #6819 ![TestStateWithVariables](https://github.com/user-attachments/assets/ab981622-b773-4983-a9ce-a70c8fcfc711) --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
diff --git a/package-lock.json b/package-lock.json
diff --git a/packages/core/package.json b/packages/core/package.json
@@ -517,6 +517,7 @@
         "@aws-sdk/client-ssm": "<3.731.0",
         "@aws-sdk/client-sso": "<3.731.0",
         "@aws-sdk/client-sso-oidc": "<3.731.0",
+        "@aws-sdk/client-sfn": "<3.731.0",
         "@aws-sdk/credential-provider-env": "<3.731.0",
         "@aws-sdk/credential-provider-process": "<3.731.0",
         "@aws-sdk/credential-provider-sso": "<3.731.0",
diff --git a/packages/core/src/shared/clients/stepFunctions.ts b/packages/core/src/shared/clients/stepFunctions.ts
@@ -0,0 +1,68 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import {
+    CreateStateMachineCommand,
+    CreateStateMachineCommandInput,
+    CreateStateMachineCommandOutput,
+    DescribeStateMachineCommand,
+    DescribeStateMachineCommandInput,
+    DescribeStateMachineCommandOutput,
+    ListStateMachinesCommand,
+    ListStateMachinesCommandInput,
+    ListStateMachinesCommandOutput,
+    SFNClient,
+    StartExecutionCommand,
+    StartExecutionCommandInput,
+    StartExecutionCommandOutput,
+    StateMachineListItem,
+    TestStateCommand,
+    TestStateCommandInput,
+    TestStateCommandOutput,
+    UpdateStateMachineCommand,
+    UpdateStateMachineCommandInput,
+    UpdateStateMachineCommandOutput,
+} from '@aws-sdk/client-sfn'
+import { ClientWrapper } from './clientWrapper'
+
+export class StepFunctionsClient extends ClientWrapper<SFNClient> {
+    public constructor(regionCode: string) {
+        super(regionCode, SFNClient)
+    }
+
+    public async *listStateMachines(
+        request: ListStateMachinesCommandInput = {}
+    ): AsyncIterableIterator<StateMachineListItem> {
+        do {
+            const response: ListStateMachinesCommandOutput = await this.makeRequest(ListStateMachinesCommand, request)
+            if (response.stateMachines) {
+                yield* response.stateMachines
+            }
+            request.nextToken = response.nextToken
+        } while (request.nextToken)
+    }
+
+    public async getStateMachineDetails(
+        request: DescribeStateMachineCommandInput
+    ): Promise<DescribeStateMachineCommandOutput> {
+        return this.makeRequest(DescribeStateMachineCommand, request)
+    }
+
+    public async executeStateMachine(request: StartExecutionCommandInput): Promise<StartExecutionCommandOutput> {
+        return this.makeRequest(StartExecutionCommand, request)
+    }
+
+    public async createStateMachine(request: CreateStateMachineCommandInput): Promise<CreateStateMachineCommandOutput> {
+        return this.makeRequest(CreateStateMachineCommand, request)
+    }
+
+    public async updateStateMachine(request: UpdateStateMachineCommandInput): Promise<UpdateStateMachineCommandOutput> {
+        return this.makeRequest(UpdateStateMachineCommand, request)
+    }
+
+    public async testState(request: TestStateCommandInput): Promise<TestStateCommandOutput> {
+        return this.makeRequest(TestStateCommand, request)
+    }
+}
diff --git a/packages/core/src/shared/clients/stepFunctionsClient.ts b/packages/core/src/shared/clients/stepFunctionsClient.ts
diff --git a/packages/core/src/stepFunctions/commands/downloadStateMachineDefinition.ts b/packages/core/src/stepFunctions/commands/downloadStateMachineDefinition.ts
@@ -7,10 +7,10 @@ import * as nls from 'vscode-nls'
 import * as os from 'os'
 const localize = nls.loadMessageBundle()
 
-import { StepFunctions } from 'aws-sdk'
+import * as StepFunctions from '@aws-sdk/client-sfn'
 import * as path from 'path'
 import * as vscode from 'vscode'
-import { DefaultStepFunctionsClient, StepFunctionsClient } from '../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../shared/clients/stepFunctions'
 
 import { getLogger, Logger } from '../../shared/logger/logger'
 import { Result } from '../../shared/telemetry/telemetry'
@@ -28,10 +28,11 @@ export async function downloadStateMachineDefinition(params: {
     let downloadResult: Result = 'Succeeded'
     const stateMachineName = params.stateMachineNode.details.name
     try {
-        const client: StepFunctionsClient = new DefaultStepFunctionsClient(params.stateMachineNode.regionCode)
-        const stateMachineDetails: StepFunctions.DescribeStateMachineOutput = await client.getStateMachineDetails(
-            params.stateMachineNode.details.stateMachineArn
-        )
+        const client: StepFunctionsClient = new StepFunctionsClient(params.stateMachineNode.regionCode)
+        const stateMachineDetails: StepFunctions.DescribeStateMachineCommandOutput =
+            await client.getStateMachineDetails({
+                stateMachineArn: params.stateMachineNode.details.stateMachineArn,
+            })
 
         if (params.isPreviewAndRender) {
             const doc = await vscode.workspace.openTextDocument({
@@ -53,7 +54,7 @@ export async function downloadStateMachineDefinition(params: {
 
             if (fileInfo) {
                 const filePath = fileInfo.fsPath
-                await fs.writeFile(filePath, stateMachineDetails.definition, 'utf8')
+                await fs.writeFile(filePath, stateMachineDetails.definition || '', 'utf8')
                 const openPath = vscode.Uri.file(filePath)
                 const doc = await vscode.workspace.openTextDocument(openPath)
                 await vscode.window.showTextDocument(doc)
diff --git a/packages/core/src/stepFunctions/commands/publishStateMachine.ts b/packages/core/src/stepFunctions/commands/publishStateMachine.ts
@@ -7,7 +7,7 @@ import { load } from 'js-yaml'
 import * as vscode from 'vscode'
 import * as nls from 'vscode-nls'
 import { AwsContext } from '../../shared/awsContext'
-import { DefaultStepFunctionsClient, StepFunctionsClient } from '../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../shared/clients/stepFunctions'
 import { getLogger, Logger } from '../../shared/logger/logger'
 import { showViewLogsMessage } from '../../shared/utilities/messages'
 import { VALID_SFN_PUBLISH_FORMATS, YAML_FORMATS } from '../constants/aslFormats'
@@ -64,7 +64,7 @@ export async function publishStateMachine(params: publishStateMachineParams) {
         if (!response) {
             return
         }
-        const client = new DefaultStepFunctionsClient(response.region)
+        const client = new StepFunctionsClient(response.region)
 
         if (response?.createResponse) {
             await createStateMachine(response.createResponse, text, params.outputChannel, response.region, client)
@@ -109,7 +109,7 @@ async function createStateMachine(
                 wizardResponse.name
             )
         )
-        outputChannel.appendLine(result.stateMachineArn)
+        outputChannel.appendLine(result.stateMachineArn || '')
         logger.info(`Created "${result.stateMachineArn}"`)
     } catch (err) {
         const msg = localize(
diff --git a/packages/core/src/stepFunctions/explorer/stepFunctionsNodes.ts b/packages/core/src/stepFunctions/explorer/stepFunctionsNodes.ts
@@ -7,9 +7,9 @@ import * as os from 'os'
 import * as nls from 'vscode-nls'
 const localize = nls.loadMessageBundle()
 
-import { StepFunctions } from 'aws-sdk'
+import * as StepFunctions from '@aws-sdk/client-sfn'
 import * as vscode from 'vscode'
-import { DefaultStepFunctionsClient } from '../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../shared/clients/stepFunctions'
 
 import { AWSResourceNode } from '../../shared/treeview/nodes/awsResourceNode'
 import { AWSTreeNodeBase } from '../../shared/treeview/nodes/awsTreeNodeBase'
@@ -40,7 +40,7 @@ export class StepFunctionsNode extends AWSTreeNodeBase {
 
     public constructor(
         public override readonly regionCode: string,
-        private readonly client = new DefaultStepFunctionsClient(regionCode)
+        private readonly client = new StepFunctionsClient(regionCode)
     ) {
         super('Step Functions', vscode.TreeItemCollapsibleState.Collapsed)
         this.stateMachineNodes = new Map<string, StateMachineNode>()
@@ -101,7 +101,7 @@ export class StateMachineNode extends AWSTreeNodeBase implements AWSResourceNode
     }
 
     public get arn(): string {
-        return this.details.stateMachineArn
+        return this.details.stateMachineArn || ''
     }
 
     public get name(): string {
diff --git a/packages/core/src/stepFunctions/utils.ts b/packages/core/src/stepFunctions/utils.ts
@@ -5,10 +5,10 @@
 import * as nls from 'vscode-nls'
 const localize = nls.loadMessageBundle()
 
-import { StepFunctions } from 'aws-sdk'
+import * as StepFunctions from '@aws-sdk/client-sfn'
 import * as yaml from 'js-yaml'
 import * as vscode from 'vscode'
-import { StepFunctionsClient } from '../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../shared/clients/stepFunctions'
 import {
     DiagnosticSeverity,
     DocumentLanguageSettings,
diff --git a/packages/core/src/stepFunctions/vue/executeStateMachine/executeStateMachine.ts b/packages/core/src/stepFunctions/vue/executeStateMachine/executeStateMachine.ts
@@ -6,7 +6,7 @@
 import * as nls from 'vscode-nls'
 const localize = nls.loadMessageBundle()
 
-import { DefaultStepFunctionsClient } from '../../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../../shared/clients/stepFunctions'
 
 import { getLogger } from '../../../shared/logger/logger'
 import { Result } from '../../../shared/telemetry/telemetry'
@@ -56,11 +56,14 @@ export class ExecuteStateMachineWebview extends VueWebview {
         this.channel.appendLine('')
 
         try {
-            const client = new DefaultStepFunctionsClient(this.stateMachine.region)
-            const startExecResponse = await client.executeStateMachine(this.stateMachine.arn, input)
+            const client = new StepFunctionsClient(this.stateMachine.region)
+            const startExecResponse = await client.executeStateMachine({
+                stateMachineArn: this.stateMachine.arn,
+                input,
+            })
             this.logger.info('started execution for Step Functions State Machine')
             this.channel.appendLine(localize('AWS.stepFunctions.executeStateMachine.info.started', 'Execution started'))
-            this.channel.appendLine(startExecResponse.executionArn)
+            this.channel.appendLine(startExecResponse.executionArn || '')
         } catch (e) {
             executeResult = 'Failed'
             const error = e as Error
@@ -82,8 +85,8 @@ const Panel = VueWebview.compilePanel(ExecuteStateMachineWebview)
 
 export async function executeStateMachine(context: ExtContext, node: StateMachineNode): Promise<void> {
     const wv = new Panel(context.extensionContext, context.outputChannel, {
-        arn: node.details.stateMachineArn,
-        name: node.details.name,
+        arn: node.details.stateMachineArn || '',
+        name: node.details.name || '',
         region: node.regionCode,
     })
 
diff --git a/packages/core/src/stepFunctions/wizards/publishStateMachineWizard.ts b/packages/core/src/stepFunctions/wizards/publishStateMachineWizard.ts
@@ -22,7 +22,7 @@ import { Wizard, WIZARD_BACK } from '../../shared/wizards/wizard'
 import { isStepFunctionsRole } from '../utils'
 import { createRolePrompter } from '../../shared/ui/common/roles'
 import { IamClient } from '../../shared/clients/iam'
-import { DefaultStepFunctionsClient } from '../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../shared/clients/stepFunctions'
 
 export enum PublishStateMachineAction {
     QuickCreate,
@@ -109,14 +109,14 @@ function createStepFunctionsRolePrompter(region: string) {
 }
 
 async function* listStateMachines(region: string) {
-    const client = new DefaultStepFunctionsClient(region)
+    const client = new StepFunctionsClient(region)
 
     for await (const machine of client.listStateMachines()) {
         yield [
             {
-                label: machine.name,
-                data: machine.stateMachineArn,
-                description: machine.stateMachineArn,
+                label: machine.name || '',
+                data: machine.stateMachineArn || '',
+                description: machine.stateMachineArn || '',
             },
         ]
     }
diff --git a/packages/core/src/stepFunctions/workflowStudio/types.ts b/packages/core/src/stepFunctions/workflowStudio/types.ts
@@ -2,7 +2,8 @@
  * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
-import { IAM, StepFunctions } from 'aws-sdk'
+import { IAM } from 'aws-sdk'
+import * as StepFunctions from '@aws-sdk/client-sfn'
 import * as vscode from 'vscode'
 
 export enum WorkflowMode {
diff --git a/packages/core/src/stepFunctions/workflowStudio/workflowStudioApiHandler.ts b/packages/core/src/stepFunctions/workflowStudio/workflowStudioApiHandler.ts
@@ -3,9 +3,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { StepFunctions } from 'aws-sdk'
+import * as StepFunctions from '@aws-sdk/client-sfn'
 import { IamClient, IamRole } from '../../shared/clients/iam'
-import { DefaultStepFunctionsClient } from '../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../shared/clients/stepFunctions'
 import { ApiAction, ApiCallRequestMessage, Command, MessageType, WebviewContext } from './types'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { ListRolesRequest } from '@aws-sdk/client-iam'
@@ -15,7 +15,7 @@ export class WorkflowStudioApiHandler {
         region: string,
         private readonly context: WebviewContext,
         private readonly clients = {
-            sfn: new DefaultStepFunctionsClient(region),
+            sfn: new StepFunctionsClient(region),
             iam: new IamClient(region),
         }
     ) {}
diff --git a/packages/core/src/test/stepFunctions/explorer/stepFunctionNodes.test.ts b/packages/core/src/test/stepFunctions/explorer/stepFunctionNodes.test.ts
@@ -15,14 +15,14 @@ import {
 } from '../../utilities/explorerNodeAssertions'
 import { asyncGenerator } from '../../../shared/utilities/collectionUtils'
 import globals from '../../../shared/extensionGlobals'
-import { DefaultStepFunctionsClient } from '../../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../../shared/clients/stepFunctions'
 import { stub } from '../../utilities/stubber'
 
 const regionCode = 'someregioncode'
 
 describe('StepFunctionsNode', function () {
     function createStatesClient(...stateMachineNames: string[]) {
-        const client = stub(DefaultStepFunctionsClient, { regionCode })
+        const client = stub(StepFunctionsClient, { regionCode })
         client.listStateMachines.returns(
             asyncGenerator(
                 stateMachineNames.map((name) => {
diff --git a/packages/core/src/test/stepFunctions/workflowStudio/workflowStudioApiHandler.test.ts b/packages/core/src/test/stepFunctions/workflowStudio/workflowStudioApiHandler.test.ts
@@ -16,7 +16,7 @@ import {
 } from '../../../stepFunctions/workflowStudio/types'
 import * as vscode from 'vscode'
 import { assertTelemetry } from '../../testUtil'
-import { DefaultStepFunctionsClient } from '../../../shared/clients/stepFunctionsClient'
+import { StepFunctionsClient } from '../../../shared/clients/stepFunctions'
 import { IamClient } from '../../../shared/clients/iam'
 
 describe('WorkflowStudioApiHandler', function () {
@@ -64,7 +64,7 @@ describe('WorkflowStudioApiHandler', function () {
             fileId: '',
         }
 
-        const sfnClient = new DefaultStepFunctionsClient('us-east-1')
+        const sfnClient = new StepFunctionsClient('us-east-1')
         apiHandler = new WorkflowStudioApiHandler('us-east-1', context, {
             sfn: sfnClient,
             iam: new IamClient('us-east-1'),
diff --git a/packages/toolkit/.changes/next-release/Bug Fix-3da4aea6-cfe4-40dc-a3d6-8dbd35e120e6.json b/packages/toolkit/.changes/next-release/Bug Fix-3da4aea6-cfe4-40dc-a3d6-8dbd35e120e6.json

Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ import { Wizard, WIZARD_BACK } from '../../shared/wizards/wizard'`
`22`	`22`	`import { isStepFunctionsRole } from '../utils'`
`23`	`23`	`import { createRolePrompter } from '../../shared/ui/common/roles'`
`24`	`24`	`import { IamClient } from '../../shared/clients/iam'`
`25`		`-import { DefaultStepFunctionsClient } from '../../shared/clients/stepFunctionsClient'`
	`25`	`+import { StepFunctionsClient } from '../../shared/clients/stepFunctions'`
`26`	`26`
`27`	`27`	`export enum PublishStateMachineAction {`
`28`	`28`	`QuickCreate,`
`@@ -109,14 +109,14 @@ function createStepFunctionsRolePrompter(region: string) {`
`109`	`109`	`}`
`110`	`110`
`111`	`111`	`async function* listStateMachines(region: string) {`
`112`		`- const client = new DefaultStepFunctionsClient(region)`
	`112`	`+ const client = new StepFunctionsClient(region)`
`113`	`113`
`114`	`114`	`for await (const machine of client.listStateMachines()) {`
`115`	`115`	`yield [`
`116`	`116`	`{`
`117`		`- label: machine.name,`
`118`		`- data: machine.stateMachineArn,`
`119`		`- description: machine.stateMachineArn,`
	`117`	`+ label: machine.name \|\| '',`
	`118`	`+ data: machine.stateMachineArn \|\| '',`
	`119`	`+ description: machine.stateMachineArn \|\| '',`
`120`	`120`	`},`
`121`	`121`	`]`
`122`	`122`	`}`