feat(sagemaker): Merging Feature/sagemaker connect to master

packages/core/package.json

@@ -407,26 +407,40 @@
                     "fontCharacter": "\\f1de"
                 }
             },
-            "aws-schemas-registry": {
+            "aws-sagemaker-code-editor": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1df"
                 }
             },
-            "aws-schemas-schema": {
+            "aws-sagemaker-jupyter-lab": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e0"
                 }
             },
-            "aws-stepfunctions-preview": {
+            "aws-schemas-registry": {
                 "description": "AWS Contributed Icon",
                 "default": {
                     "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
                     "fontCharacter": "\\f1e1"
                 }
+            },
+            "aws-schemas-schema": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e2"
+                }
+            },
+            "aws-stepfunctions-preview": {
+                "description": "AWS Contributed Icon",
+                "default": {
+                    "fontPath": "./resources/fonts/aws-toolkit-icons.woff",
+                    "fontCharacter": "\\f1e3"
+                }
             }
         }
     },
@@ -516,6 +530,8 @@
     "dependencies": {
         "@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",
         "@amzn/codewhisperer-streaming": "file:../../src.gen/@amzn/codewhisperer-streaming",
+        "@amzn/sagemaker-client": "file:../../src.gen/@amzn/sagemaker-client/1.0.0.tgz",
+        "@aws-sdk/credential-providers": "<3.731.0",
         "@aws-sdk/client-api-gateway": "<3.731.0",
         "@aws-sdk/client-apprunner": "<3.731.0",
         "@aws-sdk/client-cloudcontrol": "<3.731.0",
@@ -529,6 +545,7 @@
         "@aws-sdk/client-iam": "<3.731.0",
         "@aws-sdk/client-lambda": "<3.731.0",
         "@aws-sdk/client-s3": "<3.731.0",
+        "@aws-sdk/client-sagemaker": "<3.696.0",
         "@aws-sdk/client-ssm": "<3.731.0",
         "@aws-sdk/client-sso": "<3.731.0",
         "@aws-sdk/client-sso-oidc": "<3.731.0",

packages/core/package.nls.json

@@ -229,6 +229,7 @@
     "AWS.command.s3.createFolder": "Create Folder...",
     "AWS.command.s3.uploadFile": "Upload Files...",
     "AWS.command.s3.uploadFileToParent": "Upload to Parent...",
+    "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",
     "AWS.command.stepFunctions.openWithWorkflowStudio": "Open with Workflow Studio",

packages/core/resources/sagemaker_connect

@@ -0,0 +1,133 @@
+#!/bin/bash
+
+set -x
+
+_get_ssm_session_info() {
+    local credentials_type="$1"
+    local aws_resource_arn="$2"
+    local local_endpoint_port="$3"
+
+    local url_to_get_session_info="http://localhost:${local_endpoint_port}/get_session?connection_identifier=${aws_resource_arn}&credentials_type=${credentials_type}"
+
+    # Use curl with --write-out to capture HTTP status
+    response=$(curl -s -w "%{http_code}" -o /tmp/ssm_session_response.json "$url_to_get_session_info")
+    http_status="${response: -3}"
+    session_json=$(cat /tmp/ssm_session_response.json)
+
+    # Clean up temporary file
+    rm -f /tmp/ssm_session_response.json
+
+    if [[ "$http_status" -ne 200 ]]; then
+        echo "Error: Failed to get SSM session info. HTTP status: $http_status"
+        echo "Response: $session_json"
+        exit 1
+    fi
+
+    if [ -z "$session_json" ]; then
+        echo "Error: SSM connection info is empty."
+        exit 1
+    fi
+
+    export SSM_SESSION_JSON="$session_json"
+}
+
+_get_ssm_session_info_async() {
+    local credentials_type="$1"
+    local aws_resource_arn="$2"
+    local local_endpoint_port="$3"
+
+    local request_id=$(date +%s%3N)
+    local url_base="http://localhost:${local_endpoint_port}/get_session_async"
+    local url_to_get_session_info="${url_base}?connection_identifier=${aws_resource_arn}&credentials_type=${credentials_type}&request_id=${request_id}"
+
+    local max_retries=60
+    local retry_interval=5
+    local attempt=1
+
+    while (( attempt <= max_retries )); do
+        response=$(curl -s -w "%{http_code}" -o /tmp/ssm_session_response.json "$url_to_get_session_info")
+        http_status="${response: -3}"
+        session_json=$(cat /tmp/ssm_session_response.json)
+
+        if [[ "$http_status" -eq 200 ]]; then
+            export SSM_SESSION_JSON="$session_json"
+            return 0
+        elif [[ "$http_status" -eq 202 || "$http_status" -eq 204 ]]; then
+            echo "Info: Session not ready (HTTP $http_status). Retrying in $retry_interval seconds... [$attempt/$max_retries]"
+            sleep $retry_interval
+            ((attempt++))
+        else
+            echo "Error: Failed to get SSM session info. HTTP status: $http_status"
+            echo "Response: $session_json"
+            exit 1
+        fi
+    done
+
+    echo "Error: Timed out after $max_retries attempts waiting for session to be ready."
+    exit 1
+}
+
+
+# Validate input
+if [ "$#" -ne 1 ]; then
+    echo "Usage: $0 <hostname>"
+    exit 1
+fi
+
+HOSTNAME="$1"
+
+# Parse creds_type and AWS resource ARN from HOSTNAME
+if [[ "$HOSTNAME" =~ ^sm_([^_]+)_(arn_._aws.*)$ ]]; then
+    CREDS_TYPE="${BASH_REMATCH[1]}"
+    AWS_RESOURCE_ARN="${BASH_REMATCH[2]}"
+else
+    echo "Hostname: $HOSTNAME"
+    echo "Invalid hostname format. Expected format: sm_<creds-type>_<AWSResourceARN>"
+    exit 1
+fi
+
+# Workaround: Replace "__" with "/" in ARN
+AWS_RESOURCE_ARN=$(echo "${AWS_RESOURCE_ARN}" | sed 's|__|/|g; s|_._|:|g; s|jupyterlab/default|JupyterLab/default|g')
+REGION=$(echo "$AWS_RESOURCE_ARN" | cut -d: -f4)
+
+# Validate credentials type
+if [[ "$CREDS_TYPE" != "lc" && "$CREDS_TYPE" != "dl" ]]; then
+    echo "Invalid creds_type. Must be 'lc' or 'dl'."
+    exit 1
+fi
+
+# Validate required env var and file
+if [ -z "$SAGEMAKER_LOCAL_SERVER_FILE_PATH" ]; then
+    echo "[Error] SAGEMAKER_LOCAL_SERVER_FILE_PATH is not set"
+    exit 1
+fi
+
+if [ ! -f "$SAGEMAKER_LOCAL_SERVER_FILE_PATH" ]; then
+    echo "[Error] File not found at SAGEMAKER_LOCAL_SERVER_FILE_PATH: $SAGEMAKER_LOCAL_SERVER_FILE_PATH"
+    exit 1
+fi
+
+# Extract port from file
+LOCAL_ENDPOINT_PORT=$(jq -r '.port' "$SAGEMAKER_LOCAL_SERVER_FILE_PATH")
+if [ -z "$LOCAL_ENDPOINT_PORT" ] || [ "$LOCAL_ENDPOINT_PORT" == "null" ]; then
+    echo "[Error] 'port' field is missing or invalid in $SAGEMAKER_LOCAL_SERVER_FILE_PATH"
+    exit 1
+fi
+
+# Determine region from ARN
+if [ "$CREDS_TYPE" == "lc" ]; then
+    credentials_type="local"
+    _get_ssm_session_info "$credentials_type" "$AWS_RESOURCE_ARN" "$LOCAL_ENDPOINT_PORT"
+elif [ "$CREDS_TYPE" == "dl" ]; then
+    credentials_type="deeplink"
+    _get_ssm_session_info_async "$credentials_type" "$AWS_RESOURCE_ARN" "$LOCAL_ENDPOINT_PORT"
+else
+    echo "[Error] Invalid creds_type. Must be 'lc' or 'dl'."
+    exit 1
+fi
+
+echo "Extracting AWS_SSM_CLI: $AWS_SSM_CLI"
+
+# Execute the session
+AWS_SSM_CLI="${AWS_SSM_CLI:=session-manager-plugin}" 
+exec "${AWS_SSM_CLI}" "$SSM_SESSION_JSON" "$REGION" StartSession

packages/core/resources/sagemaker_connect.ps1

@@ -0,0 +1,148 @@
+param (
+    [Parameter(Mandatory = $true)][string]$Hostname
+)
+
+Write-Host "`n--- Script Start ---"
+Write-Host "Start Time: $(Get-Date -Format o)"
+Write-Host "Hostname argument received: $Hostname"
+
+Set-PSDebug -Trace 1
+
+function Get-SSMSessionInfo {
+    param (
+        [string]$CredentialsType,
+        [string]$AwsResourceArn,
+        [int]$LocalEndpointPort
+    )
+
+    Write-Host "Calling Get-SSMSessionInfo with credsType=${CredentialsType}, arn=${AwsResourceArn}, port=${LocalEndpointPort}"
+
+    $url = "http://127.0.0.1:$LocalEndpointPort/get_session?connection_identifier=$AwsResourceArn&credentials_type=$CredentialsType"
+    Write-Host "Request URL: $url"
+
+    try {
+        $response = Invoke-WebRequest -Uri $url -UseBasicParsing -ErrorAction Stop
+        Write-Host "Received response with status: $($response.StatusCode)"
+
+        if ($response.StatusCode -ne 200) {
+            Write-Error "Failed to get SSM session info. HTTP status: $($response.StatusCode)"
+            Write-Error "Response: $($response.Content)"
+            exit 1
+        }
+
+        if (-not $response.Content) {
+            Write-Error "SSM connection info is empty."
+            exit 1
+        }
+
+        $script:SSM_SESSION_JSON = $response.Content
+        Write-Host "Session JSON successfully retrieved"
+    } catch {
+        Write-Error "Exception in Get-SSMSessionInfo: $_"
+        exit 1
+    }
+}
+
+function Get-SSMSessionInfoAsync {
+    param (
+        [string]$CredentialsType,
+        [string]$AwsResourceArn,
+        [int]$LocalEndpointPort
+    )
+
+    $requestId = [string][DateTimeOffset]::Now.ToUnixTimeMilliseconds()
+    $url = "http://localhost:$LocalEndpointPort/get_session_async?connection_identifier=$AwsResourceArn&credentials_type=$CredentialsType&request_id=$requestId"
+    Write-Host "Calling Get-SSMSessionInfoAsync with URL: $url"
+
+    $maxRetries = 60
+    $retryInterval = 5
+
+    for ($attempt = 1; $attempt -le $maxRetries; $attempt++) {
+        try {
+            $response = Invoke-WebRequest -Uri $url -UseBasicParsing -ErrorAction Stop
+            $statusCode = $response.StatusCode
+            Write-Host "Attempt ${attempt}: HTTP ${statusCode}"
+
+            if ($statusCode -eq 200) {
+                $script:SSM_SESSION_JSON = $response.Content
+                Write-Host "Session JSON successfully retrieved"
+                return
+            } elseif ($statusCode -eq 202 -or $statusCode -eq 204) {
+                Write-Host "Session not ready. Retrying in ${retryInterval} seconds... [${attempt}/${maxRetries}]"
+                Start-Sleep -Seconds $retryInterval
+            } else {
+                Write-Error "Failed to get SSM session info. HTTP status: ${statusCode}"
+                Write-Error "Response: $($response.Content)"
+                exit 1
+            }
+        } catch {
+            Write-Error "Exception in Get-SSMSessionInfoAsync: $_"
+            exit 1
+        }
+    }
+
+    Write-Error "Timed out after ${maxRetries} attempts waiting for session to be ready."
+    exit 1
+}
+
+# Parse creds_type and AWS resource ARN from HOSTNAME
+Write-Host "`nParsing hostname..."
+if ($Hostname -match "^sm_([^_]+)_(arn_._aws.*)$") {
+    $CREDS_TYPE = $matches[1]
+    $AWS_RESOURCE_ARN = $matches[2] -replace '_._', ':' -replace '__', '/'
+} else {
+    Write-Error "Invalid hostname format. Expected format: sm_<creds-type>_<AWSResourceARN>"
+    exit 1
+}
+
+$REGION = ($AWS_RESOURCE_ARN -split ':')[3]
+Write-Host "Parsed values:"
+Write-Host "  CREDS_TYPE: ${CREDS_TYPE}"
+Write-Host "  AWS_RESOURCE_ARN: ${AWS_RESOURCE_ARN}"
+Write-Host "  REGION: ${REGION}"
+
+# Validate credentials type
+if ($CREDS_TYPE -ne "lc" -and $CREDS_TYPE -ne "dl") {
+    Write-Error "Invalid creds_type. Must be 'lc' or 'dl'."
+    exit 1
+}
+
+# Read port from local info JSON
+Write-Host "`nReading SAGEMAKER_LOCAL_SERVER_FILE_PATH: $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH"
+try {
+    $jsonContent = Get-Content $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH -Raw | ConvertFrom-Json
+    $LOCAL_ENDPOINT_PORT = $jsonContent.port
+    Write-Host "Extracted port: $LOCAL_ENDPOINT_PORT"
+} catch {
+    Write-Error "Failed to read or parse JSON file at $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH"
+    exit 1
+}
+
+if (-not $LOCAL_ENDPOINT_PORT -or $LOCAL_ENDPOINT_PORT -eq "null") {
+    Write-Error "'port' field is missing or invalid in $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH"
+    exit 1
+}
+
+# Retrieve SSM session
+Write-Host "`nStarting session retrieval..."
+if ($CREDS_TYPE -eq "lc") {
+    Get-SSMSessionInfo -CredentialsType "local" -AwsResourceArn $AWS_RESOURCE_ARN -LocalEndpointPort $LOCAL_ENDPOINT_PORT
+} elseif ($CREDS_TYPE -eq "dl") {
+    Get-SSMSessionInfoAsync -CredentialsType "deeplink" -AwsResourceArn $AWS_RESOURCE_ARN -LocalEndpointPort $LOCAL_ENDPOINT_PORT
+}
+
+# Execute the session
+Write-Host "`nLaunching session-manager-plugin..."
+$sessionPlugin = if ($env:AWS_SSM_CLI) { $env:AWS_SSM_CLI } else { "session-manager-plugin" }
+
+$jsonObj = $script:SSM_SESSION_JSON | ConvertFrom-Json
+$streamUrl = $jsonObj.StreamUrl
+$tokenValue = $jsonObj.TokenValue
+$sessionId = $jsonObj.SessionId
+
+Write-Host "Session Values:"
+Write-Host "  Stream URL: ${streamUrl}"
+Write-Host "  Token Value: ${tokenValue}"
+Write-Host "  Session ID: ${sessionId}"
+
+& $sessionPlugin "{\`"streamUrl\`":\`"${streamUrl}\`",\`"tokenValue\`":\`"${tokenValue}\`",\`"sessionId\`":\`"${sessionId}\`"}" "$REGION" "StartSession"

packages/core/src/auth/auth.ts

@@ -182,6 +182,10 @@ export class Auth implements AuthService, ConnectionManager {
         return Object.values(this._declaredConnections)
     }
 
+    public getCurrentProfileId() {
+        return this.store.getCurrentProfileId()
+    }
+
     @withTelemetryContext({ name: 'restorePreviousSession', class: authClassName })
     public async restorePreviousSession(): Promise<Connection | undefined> {
         const id = this.store.getCurrentProfileId()

packages/core/src/auth/credentials/store.ts

@@ -18,7 +18,7 @@ export interface CachedCredentials {
  * Simple cache for credentials
  */
 export class CredentialsStore {
-    private readonly credentialsCache: { [key: string]: CachedCredentials }
+    public readonly credentialsCache: { [key: string]: CachedCredentials }
 
     public constructor() {
         this.credentialsCache = {}