Skip to content

Commit ae6b2d1

Browse files
aws-toolkit-automationaws-toolkit-automation
authored
Merge master into feature/ui-e2e-tests
2 parents 2cfe0c6 + 4edc815 commit ae6b2d1

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

61 files changed

+13569
-5523
lines changed

‎package-lock.json

Lines changed: 9111 additions & 5505 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/core/package.json

Lines changed: 20 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -407,26 +407,40 @@
407407
"fontCharacter": "\\f1de"
408408
}
409409
},
410-
"aws-schemas-registry": {
410+
"aws-sagemaker-code-editor": {
411411
"description": "AWS Contributed Icon",
412412
"default": {
413413
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
414414
"fontCharacter": "\\f1df"
415415
}
416416
},
417-
"aws-schemas-schema": {
417+
"aws-sagemaker-jupyter-lab": {
418418
"description": "AWS Contributed Icon",
419419
"default": {
420420
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
421421
"fontCharacter": "\\f1e0"
422422
}
423423
},
424-
"aws-stepfunctions-preview": {
424+
"aws-schemas-registry": {
425425
"description": "AWS Contributed Icon",
426426
"default": {
427427
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
428428
"fontCharacter": "\\f1e1"
429429
}
430+
},
431+
"aws-schemas-schema": {
432+
"description": "AWS Contributed Icon",
433+
"default": {
434+
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
435+
"fontCharacter": "\\f1e2"
436+
}
437+
},
438+
"aws-stepfunctions-preview": {
439+
"description": "AWS Contributed Icon",
440+
"default": {
441+
"fontPath": "./resources/fonts/aws-toolkit-icons.woff",
442+
"fontCharacter": "\\f1e3"
443+
}
430444
}
431445
}
432446
},
@@ -516,6 +530,8 @@
516530
"dependencies": {
517531
"@amzn/amazon-q-developer-streaming-client": "file:../../src.gen/@amzn/amazon-q-developer-streaming-client",
518532
"@amzn/codewhisperer-streaming": "file:../../src.gen/@amzn/codewhisperer-streaming",
533+
"@amzn/sagemaker-client": "file:../../src.gen/@amzn/sagemaker-client/1.0.0.tgz",
534+
"@aws-sdk/credential-providers": "<3.731.0",
519535
"@aws-sdk/client-api-gateway": "<3.731.0",
520536
"@aws-sdk/client-apprunner": "<3.731.0",
521537
"@aws-sdk/client-cloudcontrol": "<3.731.0",
@@ -529,6 +545,7 @@
529545
"@aws-sdk/client-iam": "<3.731.0",
530546
"@aws-sdk/client-lambda": "<3.731.0",
531547
"@aws-sdk/client-s3": "<3.731.0",
548+
"@aws-sdk/client-sagemaker": "<3.696.0",
532549
"@aws-sdk/client-ssm": "<3.731.0",
533550
"@aws-sdk/client-sso": "<3.731.0",
534551
"@aws-sdk/client-sso-oidc": "<3.731.0",

‎packages/core/package.nls.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -229,6 +229,7 @@
229229
"AWS.command.s3.createFolder": "Create Folder...",
230230
"AWS.command.s3.uploadFile": "Upload Files...",
231231
"AWS.command.s3.uploadFileToParent": "Upload to Parent...",
232+
"AWS.command.sagemaker.filterSpaces": "Filter Sagemaker Spaces",
232233
"AWS.command.stepFunctions.createStateMachineFromTemplate": "Create a new Step Functions state machine",
233234
"AWS.command.stepFunctions.publishStateMachine": "Publish state machine to Step Functions",
234235
"AWS.command.stepFunctions.openWithWorkflowStudio": "Open with Workflow Studio",

‎packages/core/resources/icons/aws/sagemaker/code-editor.svg

Lines changed: 4 additions & 0 deletions
Loading

‎packages/core/resources/icons/aws/sagemaker/jupyter-lab.svg

Lines changed: 4 additions & 0 deletions
Loading

‎packages/core/resources/icons/fonts/aws-toolkit-icons.woff

58 KB
Binary file not shown.

‎packages/core/resources/sagemaker_connect

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
#!/bin/bash
2+
3+
set -x
4+
5+
_get_ssm_session_info() {
6+
local credentials_type="$1"
7+
local aws_resource_arn="$2"
8+
local local_endpoint_port="$3"
9+
10+
local url_to_get_session_info="http://localhost:${local_endpoint_port}/get_session?connection_identifier=${aws_resource_arn}&credentials_type=${credentials_type}"
11+
12+
# Use curl with --write-out to capture HTTP status
13+
response=$(curl -s -w "%{http_code}" -o /tmp/ssm_session_response.json "$url_to_get_session_info")
14+
http_status="${response: -3}"
15+
session_json=$(cat /tmp/ssm_session_response.json)
16+
17+
# Clean up temporary file
18+
rm -f /tmp/ssm_session_response.json
19+
20+
if [[ "$http_status" -ne 200 ]]; then
21+
echo "Error: Failed to get SSM session info. HTTP status: $http_status"
22+
echo "Response: $session_json"
23+
exit 1
24+
fi
25+
26+
if [ -z "$session_json" ]; then
27+
echo "Error: SSM connection info is empty."
28+
exit 1
29+
fi
30+
31+
export SSM_SESSION_JSON="$session_json"
32+
}
33+
34+
_get_ssm_session_info_async() {
35+
local credentials_type="$1"
36+
local aws_resource_arn="$2"
37+
local local_endpoint_port="$3"
38+
39+
local request_id=$(date +%s%3N)
40+
local url_base="http://localhost:${local_endpoint_port}/get_session_async"
41+
local url_to_get_session_info="${url_base}?connection_identifier=${aws_resource_arn}&credentials_type=${credentials_type}&request_id=${request_id}"
42+
43+
local max_retries=60
44+
local retry_interval=5
45+
local attempt=1
46+
47+
while (( attempt <= max_retries )); do
48+
response=$(curl -s -w "%{http_code}" -o /tmp/ssm_session_response.json "$url_to_get_session_info")
49+
http_status="${response: -3}"
50+
session_json=$(cat /tmp/ssm_session_response.json)
51+
52+
if [[ "$http_status" -eq 200 ]]; then
53+
export SSM_SESSION_JSON="$session_json"
54+
return 0
55+
elif [[ "$http_status" -eq 202 || "$http_status" -eq 204 ]]; then
56+
echo "Info: Session not ready (HTTP $http_status). Retrying in $retry_interval seconds... [$attempt/$max_retries]"
57+
sleep $retry_interval
58+
((attempt++))
59+
else
60+
echo "Error: Failed to get SSM session info. HTTP status: $http_status"
61+
echo "Response: $session_json"
62+
exit 1
63+
fi
64+
done
65+
66+
echo "Error: Timed out after $max_retries attempts waiting for session to be ready."
67+
exit 1
68+
}
69+
70+
71+
# Validate input
72+
if [ "$#" -ne 1 ]; then
73+
echo "Usage: $0 <hostname>"
74+
exit 1
75+
fi
76+
77+
HOSTNAME="$1"
78+
79+
# Parse creds_type and AWS resource ARN from HOSTNAME
80+
if [[ "$HOSTNAME" =~ ^sm_([^_]+)_(arn_._aws.*)$ ]]; then
81+
CREDS_TYPE="${BASH_REMATCH[1]}"
82+
AWS_RESOURCE_ARN="${BASH_REMATCH[2]}"
83+
else
84+
echo "Hostname: $HOSTNAME"
85+
echo "Invalid hostname format. Expected format: sm_<creds-type>_<AWSResourceARN>"
86+
exit 1
87+
fi
88+
89+
# Workaround: Replace "__" with "/" in ARN
90+
AWS_RESOURCE_ARN=$(echo "${AWS_RESOURCE_ARN}" | sed 's|__|/|g; s|_._|:|g; s|jupyterlab/default|JupyterLab/default|g')
91+
REGION=$(echo "$AWS_RESOURCE_ARN" | cut -d: -f4)
92+
93+
# Validate credentials type
94+
if [[ "$CREDS_TYPE" != "lc" && "$CREDS_TYPE" != "dl" ]]; then
95+
echo "Invalid creds_type. Must be 'lc' or 'dl'."
96+
exit 1
97+
fi
98+
99+
# Validate required env var and file
100+
if [ -z "$SAGEMAKER_LOCAL_SERVER_FILE_PATH" ]; then
101+
echo "[Error] SAGEMAKER_LOCAL_SERVER_FILE_PATH is not set"
102+
exit 1
103+
fi
104+
105+
if [ ! -f "$SAGEMAKER_LOCAL_SERVER_FILE_PATH" ]; then
106+
echo "[Error] File not found at SAGEMAKER_LOCAL_SERVER_FILE_PATH: $SAGEMAKER_LOCAL_SERVER_FILE_PATH"
107+
exit 1
108+
fi
109+
110+
# Extract port from file
111+
LOCAL_ENDPOINT_PORT=$(jq -r '.port' "$SAGEMAKER_LOCAL_SERVER_FILE_PATH")
112+
if [ -z "$LOCAL_ENDPOINT_PORT" ] || [ "$LOCAL_ENDPOINT_PORT" == "null" ]; then
113+
echo "[Error] 'port' field is missing or invalid in $SAGEMAKER_LOCAL_SERVER_FILE_PATH"
114+
exit 1
115+
fi
116+
117+
# Determine region from ARN
118+
if [ "$CREDS_TYPE" == "lc" ]; then
119+
credentials_type="local"
120+
_get_ssm_session_info "$credentials_type" "$AWS_RESOURCE_ARN" "$LOCAL_ENDPOINT_PORT"
121+
elif [ "$CREDS_TYPE" == "dl" ]; then
122+
credentials_type="deeplink"
123+
_get_ssm_session_info_async "$credentials_type" "$AWS_RESOURCE_ARN" "$LOCAL_ENDPOINT_PORT"
124+
else
125+
echo "[Error] Invalid creds_type. Must be 'lc' or 'dl'."
126+
exit 1
127+
fi
128+
129+
echo "Extracting AWS_SSM_CLI: $AWS_SSM_CLI"
130+
131+
# Execute the session
132+
AWS_SSM_CLI="${AWS_SSM_CLI:=session-manager-plugin}"
133+
exec "${AWS_SSM_CLI}" "$SSM_SESSION_JSON" "$REGION" StartSession

‎packages/core/resources/sagemaker_connect.ps1

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
param (
2+
[Parameter(Mandatory = $true)][string]$Hostname
3+
)
4+
5+
Write-Host "`n--- Script Start ---"
6+
Write-Host "Start Time: $(Get-Date -Format o)"
7+
Write-Host "Hostname argument received: $Hostname"
8+
9+
Set-PSDebug -Trace 1
10+
11+
function Get-SSMSessionInfo {
12+
param (
13+
[string]$CredentialsType,
14+
[string]$AwsResourceArn,
15+
[int]$LocalEndpointPort
16+
)
17+
18+
Write-Host "Calling Get-SSMSessionInfo with credsType=${CredentialsType}, arn=${AwsResourceArn}, port=${LocalEndpointPort}"
19+
20+
$url = "http://127.0.0.1:$LocalEndpointPort/get_session?connection_identifier=$AwsResourceArn&credentials_type=$CredentialsType"
21+
Write-Host "Request URL: $url"
22+
23+
try {
24+
$response = Invoke-WebRequest -Uri $url -UseBasicParsing -ErrorAction Stop
25+
Write-Host "Received response with status: $($response.StatusCode)"
26+
27+
if ($response.StatusCode -ne 200) {
28+
Write-Error "Failed to get SSM session info. HTTP status: $($response.StatusCode)"
29+
Write-Error "Response: $($response.Content)"
30+
exit 1
31+
}
32+
33+
if (-not $response.Content) {
34+
Write-Error "SSM connection info is empty."
35+
exit 1
36+
}
37+
38+
$script:SSM_SESSION_JSON = $response.Content
39+
Write-Host "Session JSON successfully retrieved"
40+
} catch {
41+
Write-Error "Exception in Get-SSMSessionInfo: $_"
42+
exit 1
43+
}
44+
}
45+
46+
function Get-SSMSessionInfoAsync {
47+
param (
48+
[string]$CredentialsType,
49+
[string]$AwsResourceArn,
50+
[int]$LocalEndpointPort
51+
)
52+
53+
$requestId = [string][DateTimeOffset]::Now.ToUnixTimeMilliseconds()
54+
$url = "http://localhost:$LocalEndpointPort/get_session_async?connection_identifier=$AwsResourceArn&credentials_type=$CredentialsType&request_id=$requestId"
55+
Write-Host "Calling Get-SSMSessionInfoAsync with URL: $url"
56+
57+
$maxRetries = 60
58+
$retryInterval = 5
59+
60+
for ($attempt = 1; $attempt -le $maxRetries; $attempt++) {
61+
try {
62+
$response = Invoke-WebRequest -Uri $url -UseBasicParsing -ErrorAction Stop
63+
$statusCode = $response.StatusCode
64+
Write-Host "Attempt ${attempt}: HTTP ${statusCode}"
65+
66+
if ($statusCode -eq 200) {
67+
$script:SSM_SESSION_JSON = $response.Content
68+
Write-Host "Session JSON successfully retrieved"
69+
return
70+
} elseif ($statusCode -eq 202 -or $statusCode -eq 204) {
71+
Write-Host "Session not ready. Retrying in ${retryInterval} seconds... [${attempt}/${maxRetries}]"
72+
Start-Sleep -Seconds $retryInterval
73+
} else {
74+
Write-Error "Failed to get SSM session info. HTTP status: ${statusCode}"
75+
Write-Error "Response: $($response.Content)"
76+
exit 1
77+
}
78+
} catch {
79+
Write-Error "Exception in Get-SSMSessionInfoAsync: $_"
80+
exit 1
81+
}
82+
}
83+
84+
Write-Error "Timed out after ${maxRetries} attempts waiting for session to be ready."
85+
exit 1
86+
}
87+
88+
# Parse creds_type and AWS resource ARN from HOSTNAME
89+
Write-Host "`nParsing hostname..."
90+
if ($Hostname -match "^sm_([^_]+)_(arn_._aws.*)$") {
91+
$CREDS_TYPE = $matches[1]
92+
$AWS_RESOURCE_ARN = $matches[2] -replace '_._', ':' -replace '__', '/'
93+
} else {
94+
Write-Error "Invalid hostname format. Expected format: sm_<creds-type>_<AWSResourceARN>"
95+
exit 1
96+
}
97+
98+
$REGION = ($AWS_RESOURCE_ARN -split ':')[3]
99+
Write-Host "Parsed values:"
100+
Write-Host " CREDS_TYPE: ${CREDS_TYPE}"
101+
Write-Host " AWS_RESOURCE_ARN: ${AWS_RESOURCE_ARN}"
102+
Write-Host " REGION: ${REGION}"
103+
104+
# Validate credentials type
105+
if ($CREDS_TYPE -ne "lc" -and $CREDS_TYPE -ne "dl") {
106+
Write-Error "Invalid creds_type. Must be 'lc' or 'dl'."
107+
exit 1
108+
}
109+
110+
# Read port from local info JSON
111+
Write-Host "`nReading SAGEMAKER_LOCAL_SERVER_FILE_PATH: $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH"
112+
try {
113+
$jsonContent = Get-Content $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH -Raw | ConvertFrom-Json
114+
$LOCAL_ENDPOINT_PORT = $jsonContent.port
115+
Write-Host "Extracted port: $LOCAL_ENDPOINT_PORT"
116+
} catch {
117+
Write-Error "Failed to read or parse JSON file at $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH"
118+
exit 1
119+
}
120+
121+
if (-not $LOCAL_ENDPOINT_PORT -or $LOCAL_ENDPOINT_PORT -eq "null") {
122+
Write-Error "'port' field is missing or invalid in $env:SAGEMAKER_LOCAL_SERVER_FILE_PATH"
123+
exit 1
124+
}
125+
126+
# Retrieve SSM session
127+
Write-Host "`nStarting session retrieval..."
128+
if ($CREDS_TYPE -eq "lc") {
129+
Get-SSMSessionInfo -CredentialsType "local" -AwsResourceArn $AWS_RESOURCE_ARN -LocalEndpointPort $LOCAL_ENDPOINT_PORT
130+
} elseif ($CREDS_TYPE -eq "dl") {
131+
Get-SSMSessionInfoAsync -CredentialsType "deeplink" -AwsResourceArn $AWS_RESOURCE_ARN -LocalEndpointPort $LOCAL_ENDPOINT_PORT
132+
}
133+
134+
# Execute the session
135+
Write-Host "`nLaunching session-manager-plugin..."
136+
$sessionPlugin = if ($env:AWS_SSM_CLI) { $env:AWS_SSM_CLI } else { "session-manager-plugin" }
137+
138+
$jsonObj = $script:SSM_SESSION_JSON | ConvertFrom-Json
139+
$streamUrl = $jsonObj.StreamUrl
140+
$tokenValue = $jsonObj.TokenValue
141+
$sessionId = $jsonObj.SessionId
142+
143+
Write-Host "Session Values:"
144+
Write-Host " Stream URL: ${streamUrl}"
145+
Write-Host " Token Value: ${tokenValue}"
146+
Write-Host " Session ID: ${sessionId}"
147+
148+
& $sessionPlugin "{\`"streamUrl\`":\`"${streamUrl}\`",\`"tokenValue\`":\`"${tokenValue}\`",\`"sessionId\`":\`"${sessionId}\`"}" "$REGION" "StartSession"

‎packages/core/src/auth/auth.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,10 @@ export class Auth implements AuthService, ConnectionManager {
182182
return Object.values(this._declaredConnections)
183183
}
184184

185+
public getCurrentProfileId() {
186+
return this.store.getCurrentProfileId()
187+
}
188+
185189
@withTelemetryContext({ name: 'restorePreviousSession', class: authClassName })
186190
public async restorePreviousSession(): Promise<Connection | undefined> {
187191
const id = this.store.getCurrentProfileId()

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ export interface CachedCredentials {
1818
* Simple cache for credentials
1919
*/
2020
export class CredentialsStore {
21-
private readonly credentialsCache: { [key: string]: CachedCredentials }
21+
public readonly credentialsCache: { [key: string]: CachedCredentials }
2222

2323
public constructor() {
2424
this.credentialsCache = {}

0 commit comments

Comments
 (0)