feat(sagemakerunifiedstudio): Support for SageMaker Unified Studio

.github/workflows/lintcommit.js

@@ -57,6 +57,7 @@ const scopes = new Set([
     'telemetry',
     'toolkit',
     'ui',
+    'sagemakerunifiedstudio',
 ])
 void scopes
 

.gitignore

@@ -32,6 +32,8 @@ src.gen/*
 **/src/codewhisperer/client/codewhispererclient.d.ts
 **/src/codewhisperer/client/codewhispereruserclient.d.ts
 **/src/auth/sso/oidcclientpkce.d.ts
+**/src/sagemakerunifiedstudio/shared/client/gluecatalogapi.d.ts
+**/src/sagemakerunifiedstudio/shared/client/sqlworkbench.d.ts
 
 # Generated by tests
 **/src/testFixtures/**/bin

packages/core/package.json

@@ -421,26 +421,61 @@
                     "fontCharacter": "\\f1e0"
                 }
             },
-            "aws-schemas-registry": {
+            "aws-sagemakerunifiedstudio-catalog": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e1"
                 }
             },
-            "aws-schemas-schema": {
+            "aws-sagemakerunifiedstudio-spaces": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e2"
                 }
             },
-            "aws-stepfunctions-preview": {
+            "aws-sagemakerunifiedstudio-spaces-dark": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e3"
                 }
+            },
+            "aws-sagemakerunifiedstudio-symbol-int": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e4"
+                }
+            },
+            "aws-sagemakerunifiedstudio-table": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e5"
+                }
+            },
+            "aws-schemas-registry": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e6"
+                }
+            },
+            "aws-schemas-schema": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e7"
+                }
+            },
+            "aws-stepfunctions-preview": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e8"
+                }
             }
         }
     },
@@ -539,12 +574,15 @@
         "@aws-sdk/client-cloudwatch-logs": "<3.731.0",
         "@aws-sdk/client-codecatalyst": "<3.731.0",
         "@aws-sdk/client-cognito-identity": "<3.731.0",
+        "@aws-sdk/client-datazone": "^3.848.0",
         "@aws-sdk/client-docdb": "<3.731.0",
         "@aws-sdk/client-docdb-elastic": "<3.731.0",
         "@aws-sdk/client-ec2": "<3.731.0",
+        "@aws-sdk/client-glue": "^3.852.0",
         "@aws-sdk/client-iam": "<3.731.0",
         "@aws-sdk/client-lambda": "<3.731.0",
         "@aws-sdk/client-s3": "<3.731.0",
+        "@aws-sdk/client-s3-control": "^3.830.0",
         "@aws-sdk/client-sagemaker": "<3.696.0",
         "@aws-sdk/client-ssm": "<3.731.0",
         "@aws-sdk/client-sso": "<3.731.0",

packages/core/package.nls.json

@@ -231,6 +231,9 @@
     "AWS.command.s3.createFolder": "Create Folder...",
     "AWS.command.s3.uploadFile": "Upload Files...",
     "AWS.command.s3.uploadFileToParent": "Upload to Parent...",
+    "AWS.command.smus.switchProject": "Switch Project",
+    "AWS.command.smus.refreshProject": "Refresh Project",
+    "AWS.command.smus.signOut": "Sign Out",
     "AWS.command.sagemaker.filterSpaces": "Filter Sagemaker Spaces",
     "AWS.command.stepFunctions.createStateMachineFromTemplate": "Create a new Step Functions state machine",
     "AWS.command.stepFunctions.publishStateMachine": "Publish state machine to Step Functions",
@@ -297,6 +300,7 @@
     "AWS.appcomposer.explorerTitle": "Infrastructure Composer",
     "AWS.cdk.explorerTitle": "CDK",
     "AWS.codecatalyst.explorerTitle": "CodeCatalyst",
+    "AWS.sagemakerunifiedstudio.explorerTitle": "SageMaker Unified Studio",
     "AWS.cwl.limit.desc": "Maximum amount of log entries pulled per request from CloudWatch Logs. For LiveTail, when the limit is reached, the oldest events will be removed to accomodate new events. (max 10000)",
     "AWS.samcli.deploy.bucket.recentlyUsed": "Buckets recently used for SAM deployments",
     "AWS.submenu.amazonqEditorContextSubmenu.title": "Amazon Q",

packages/core/scripts/build/generateServiceClient.ts

@@ -241,6 +241,14 @@ void (async () => {
             serviceJsonPath: 'src/codewhisperer/client/user-service-2.json',
             serviceName: 'CodeWhispererUserClient',
         },
+        {
+            serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/gluecatalogapi.json',
+            serviceName: 'GlueCatalogApi',
+        },
+        {
+            serviceJsonPath: 'src/sagemakerunifiedstudio/shared/client/sqlworkbench.json',
+            serviceName: 'SQLWorkbench',
+        },
     ]
     await generateServiceClients(serviceClientDefinitions)
 })()

packages/core/src/auth/auth.ts

@@ -219,6 +219,30 @@ export class Auth implements AuthService, ConnectionManager {
         }
     }
 
+    /**
+     * Gets the SSO access token for a connection
+     * @param connection The SSO connection to get the token for
+     * @returns Promise resolving to the access token string
+     */
+    @withTelemetryContext({ name: 'getSsoAccessToken', class: authClassName })
+    public async getSsoAccessToken(connection: Pick<SsoConnection, 'id'>): Promise<string> {
+        const profile = this.store.getProfileOrThrow(connection.id)
+
+        if (profile.type !== 'sso') {
+            throw new Error(`Connection ${connection.id} is not an SSO connection`)
+        }
+
+        const provider = this.getSsoTokenProvider(connection.id, profile)
+        // Calling existing getToken private method - It will handle setting the connection state etc.
+        const token = await this._getToken(connection.id, provider)
+
+        if (!token?.accessToken) {
+            throw new Error(`No access token available for connection ${connection.id}`)
+        }
+
+        return token.accessToken
+    }
+
     public async useConnection({ id }: Pick<SsoConnection, 'id'>): Promise<SsoConnection>
     public async useConnection({ id }: Pick<IamConnection, 'id'>): Promise<IamConnection>
     @withTelemetryContext({ name: 'useConnection', class: authClassName })
@@ -923,10 +947,22 @@ export class Auth implements AuthService, ConnectionManager {
         if (previousState === 'valid') {
             // Non-token expiration errors can happen. We must log it here, otherwise they are lost.
             getLogger().warn(`auth: valid connection became invalid. Last error: %s`, this.#validationErrors.get(id))
-
             const timeout = new Timeout(60000)
             this.#invalidCredentialsTimeouts.set(id, timeout)
 
+            // Check if this is a SMUS profile - if so, skip the generic prompt
+            // as SMUS has its own reauthentication flow
+            const isSmusConnection = profile.type === 'sso' && 'domainUrl' in profile && 'domainId' in profile
+            if (isSmusConnection) {
+                getLogger().debug(`auth: Skipping generic reauthentication prompt for SMUS connection ${id}`)
+                // For SMUS connections, just throw the InvalidConnection error
+                // The SMUS auth provider will handle showing the appropriate prompt
+                throw new ToolkitError('Connection is invalid or expired. Try logging in again.', {
+                    code: errorCode.invalidConnection,
+                    cause: this.#validationErrors.get(id),
+                })
+            }
+
             const connLabel = profile.metadata.label ?? (profile.type === 'sso' ? this.getSsoProfileLabel(profile) : id)
             const message = localize(
                 'aws.auth.invalidConnection',

packages/core/src/auth/secondaryAuth.ts

@@ -18,7 +18,7 @@ import { withTelemetryContext } from '../shared/telemetry/util'
 import { isNetworkError } from '../shared/errors'
 import globals from '../shared/extensionGlobals'
 
-export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId'
+export type ToolId = 'codecatalyst' | 'codewhisperer' | 'testId' | 'smus'
 
 let currentConn: Auth['activeConnection']
 const auths = new Map<string, SecondaryAuth>()

packages/core/src/awsService/sagemaker/commands.ts

@@ -20,6 +20,7 @@ import { ToolkitError } from '../../shared/errors'
 import { showConfirmationMessage } from '../../shared/utilities/messages'
 import { RemoteSessionError } from '../../shared/remoteSession'
 import { ConnectFromRemoteWorkspaceMessage, InstanceTypeError } from './constants'
+import { SagemakerUnifiedStudioSpaceNode } from '../../sagemakerunifiedstudio/explorer/nodes/sageMakerUnifiedStudioSpaceNode'
 
 const localize = nls.loadMessageBundle()
 
@@ -101,6 +102,8 @@ export async function deeplinkConnect(
             connectionIdentifier,
             ctx.extensionContext,
             'sm_dl',
+            false /* isSMUS */,
+            undefined /* node */,
             session,
             wsUrl,
             token,
@@ -125,7 +128,11 @@ export async function deeplinkConnect(
     }
 }
 
-export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionContext) {
+export async function stopSpace(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    sageMakerClient?: SagemakerClient
+) {
     const spaceName = node.spaceApp.SpaceName!
     const confirmed = await showConfirmationMessage({
         prompt: `You are about to stop this space. Any active resource will also be stopped. Are you sure you want to stop the space?`,
@@ -137,8 +144,8 @@ export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionC
     if (!confirmed) {
         return
     }
-
-    const client = new SagemakerClient(node.regionCode)
+    //  In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client.
+    const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode)
     try {
         await client.deleteApp({
             DomainId: node.spaceApp.DomainId!,
@@ -151,36 +158,50 @@ export async function stopSpace(node: SagemakerSpaceNode, ctx: vscode.ExtensionC
         if (error.name === 'AccessDeniedException') {
             throw new ToolkitError('You do not have permission to stop spaces. Please contact your administrator', {
                 cause: error,
+                code: error.name,
             })
         } else {
-            throw err
+            throw new ToolkitError(`Failed to stop space ${spaceName}: ${(error as Error).message}`, {
+                cause: error,
+                code: error.name,
+            })
         }
     }
     await tryRefreshNode(node)
 }
 
-export async function openRemoteConnect(node: SagemakerSpaceNode, ctx: vscode.ExtensionContext) {
+export async function openRemoteConnect(
+    node: SagemakerSpaceNode | SagemakerUnifiedStudioSpaceNode,
+    ctx: vscode.ExtensionContext,
+    sageMakerClient?: SagemakerClient
+) {
     if (isRemoteWorkspace()) {
         void vscode.window.showErrorMessage(ConnectFromRemoteWorkspaceMessage)
         return
     }
 
     if (node.getStatus() === 'Stopped') {
-        const client = new SagemakerClient(node.regionCode)
+        //  In case of SMUS, we pass in a SM Client and for SM AI, it creates a new SM Client.
+        const client = sageMakerClient ? sageMakerClient : new SagemakerClient(node.regionCode)
 
         try {
             await client.startSpace(node.spaceApp.SpaceName!, node.spaceApp.DomainId!)
             await tryRefreshNode(node)
             const appType = node.spaceApp.SpaceSettingsSummary?.AppType
             if (!appType) {
-                throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.')
+                throw new ToolkitError('AppType is undefined for the selected space. Cannot start remote connection.', {
+                    code: 'undefinedAppType',
+                })
             }
             await client.waitForAppInService(node.spaceApp.DomainId!, node.spaceApp.SpaceName!, appType)
             await tryRemoteConnection(node, ctx)
         } catch (err: any) {
             // Ignore InstanceTypeError since it means the user decided not to use an instanceType with more memory
             if (err.code !== InstanceTypeError) {
-                throw err
+                throw new ToolkitError(`Remote connection failed: ${(err as Error).message}`, {
+                    cause: err as Error,
+                    code: err.code,
+                })
             }
         }
     } else if (node.getStatus() === 'Running') {