Skip to content

Commit 3983b4d

Browse files
author
Arun Bhati
committed
feat(session): add SageMaker Unified Studio portal redirect
** Description ** • Add support for redirecting to SMUS portal during VSCode session renewal • Detect SMUS environment using SERVICE_NAME variable • Update metadata handling to support SMUS portal URL construction ** Motivation ** • Enable users to renew VSCode sessions through SMUS portal • Improve session renewal experience for SMUS users ** Testing ** • Verified with local VSCode server • Tested for both SMUS and SageMaker AI with manually cookie population and metadata file creation ** Backward compatibility ** • Maintains existing behavior for non-SMUS environments • Silent failure handling prevents disruption of current flows
1 parent 914e854 commit 3983b4d

File tree

4 files changed

+256
-23
lines changed

4 files changed

+256
-23
lines changed

‎patched-vscode/extensions/sagemaker-extension/src/constant.ts‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,10 @@ export const FIVE_MINUTES_INTERVAL_MILLIS = 5 * 60 * 1000;
2727

2828
export const SAGEMAKER_METADATA_PATH = '/opt/ml/metadata/resource-metadata.json';
2929

30+
// Service name identifier for SageMaker Unified Studio
31+
export const SMUS_SERVICE_NAME = 'SageMakerUnifiedStudio';
32+
export const SERVICE_NAME_ENV_VAR = 'SERVICE_NAME';
33+
3034
export class SagemakerCookie {
3135
authMode: string
3236
expiryTime: number
@@ -56,6 +60,11 @@ export class SagemakerResourceMetadata {
5660
ResourceArn?: string
5761
ResourceName?: string
5862
AppImageVersion?: string
63+
AdditionalMetadata?: {
64+
DataZoneDomainId?: string
65+
DataZoneProjectId?: string
66+
DataZoneDomainRegion?: string
67+
}
5968
};
6069
export function isSSOMode(cookie: SagemakerCookie) {
6170
return (cookie.authMode === AUTH_MODE.SSO)
@@ -69,4 +78,35 @@ export function getExpiryTime(cookie: SagemakerCookie): number {
6978
} else {
7079
return -1;
7180
}
72-
}
81+
}
82+
83+
/**
84+
* Constructs the SMUS portal URL using domain, region, and project information
85+
* Returns null if not in SMUS environment or if required fields are missing
86+
*/
87+
export const getSmusVscodePortalUrl = (metadata: SagemakerMetadata): string | null => {
88+
if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) {
89+
return null;
90+
}
91+
92+
if (!metadata || !metadata.AdditionalMetadata) {
93+
// fail silently not to block users
94+
console.error('[SMUS] Metadata is undefined or null');
95+
return null;
96+
}
97+
98+
const { DataZoneDomainId, DataZoneDomainRegion, DataZoneProjectId } = metadata.AdditionalMetadata;
99+
100+
if (!DataZoneDomainId || !DataZoneDomainRegion || !DataZoneProjectId) {
101+
// fail silently not to block users
102+
// TODO: add monitoring to detect such cases
103+
console.error('[SMUS] Required fields missing in metadata:', {
104+
DataZoneDomainId: !!DataZoneDomainId,
105+
DataZoneDomainRegion: !!DataZoneDomainRegion,
106+
DataZoneProjectId: !!DataZoneProjectId
107+
});
108+
return null;
109+
}
110+
111+
return `https://${DataZoneDomainId}.sagemaker.${DataZoneDomainRegion}.on.aws/projects/${DataZoneProjectId}/overview`;
112+
}

‎patched-vscode/extensions/sagemaker-extension/src/extension.ts‎

Lines changed: 36 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,18 @@ import {
1111
WARNING_BUTTON_SAVE_AND_RENEW_SESSION,
1212
SagemakerCookie,
1313
SagemakerResourceMetadata,
14-
getExpiryTime
14+
getExpiryTime,
15+
getSmusVscodePortalUrl
1516
} from "./constant";
1617
import * as console from "console";
1718

1819

1920
const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies';
2021
const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate';
2122

23+
// Global redirect URL for SMUS environment
24+
let smusRedirectUrl: string | null = null;
25+
2226
function showWarningDialog() {
2327
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => {
2428

@@ -59,11 +63,12 @@ function showWarningDialog() {
5963
}
6064

6165
function signInError(sagemakerCookie: SagemakerCookie) {
66+
const redirectUrl = getRedirectUrl(sagemakerCookie);
6267
// The session has expired
6368
SessionWarning.signInWarning(sagemakerCookie)
6469
.then((selection) => {
6570
if (selection === SIGN_IN_BUTTON) {
66-
vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
71+
vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
6772
}
6873
});
6974
}
@@ -94,32 +99,31 @@ function saveWorkspace() {
9499
});
95100
}
96101
function renewSession(sagemakerCookie: SagemakerCookie) {
102+
const redirectUrl = getRedirectUrl(sagemakerCookie);
97103
// TODO: Log and trigger a Signin
98-
vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
104+
vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
99105
// Trigger the function to show the warning again after 5 minutes again to validate.
100106
setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS);
101107
}
102108

103109
function updateStatusItemWithMetadata(context: vscode.ExtensionContext) {
104-
fs.readFile(SAGEMAKER_METADATA_PATH, 'utf-8', (err, data) => {
105-
if (err) {
106-
// fail silently not to block users
107-
} else {
108-
try {
109-
const jsonData = JSON.parse(data) as SagemakerResourceMetadata;
110-
const spaceName = jsonData.SpaceName;
111-
112-
if (spaceName != null) {
113-
let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
114-
spaceNameStatusBarItem.text = `Space: ${spaceName}`;
115-
spaceNameStatusBarItem.show();
116-
context.subscriptions.push(spaceNameStatusBarItem);
117-
}
118-
} catch (jsonError) {
119-
// fail silently not to block users
120-
}
110+
try {
111+
const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8');
112+
const jsonData = JSON.parse(data) as SagemakerResourceMetadata;
113+
114+
if (jsonData?.SpaceName) {
115+
let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
116+
spaceNameStatusBarItem.text = `Space: ${jsonData.SpaceName}`;
117+
spaceNameStatusBarItem.show();
118+
context.subscriptions.push(spaceNameStatusBarItem);
121119
}
122-
});
120+
121+
// Initialize SMUS redirect URL
122+
smusRedirectUrl = getSmusVscodePortalUrl(jsonData);
123+
} catch (error) {
124+
// fail silently not to block users
125+
console.error('Error reading metadata file:', error);
126+
}
123127
}
124128

125129
// Render warning message regarding auto upgrade disabled
@@ -158,15 +162,25 @@ export function activate(context: vscode.ExtensionContext) {
158162
// TODO: log activation of extension
159163
console.log('Activating Sagemaker Extension...');
160164

165+
// Initialize metadata first (which will set smusRedirectUrl if in SMUS environment)
166+
updateStatusItemWithMetadata(context);
167+
161168
// execute the get cookie command and save the data to cookies
162169
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => {
163170

164171
const sagemakerCookie: SagemakerCookie = r as SagemakerCookie
165172

166173
initialize(sagemakerCookie);
167-
updateStatusItemWithMetadata(context);
168174
});
169175

170176
// render warning message regarding auto upgrade disabled
171177
renderExtensionAutoUpgradeDisabledNotification();
172178
}
179+
180+
/**
181+
* Returns the appropriate redirect URL based on the environment
182+
* Uses SMUS URL if available, falls back to original redirect URL
183+
*/
184+
function getRedirectUrl(sagemakerCookie: SagemakerCookie): string {
185+
return smusRedirectUrl || sagemakerCookie.redirectURL;
186+
}

‎patches/sagemaker-extension-smus-support.patch‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant.ts
2+
===================================================================
3+
--- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/constant.ts
4+
+++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant.ts
5+
@@ -27,6 +27,10 @@ export const FIVE_MINUTES_INTERVAL_MILLI
6+
7+
export const SAGEMAKER_METADATA_PATH = '/opt/ml/metadata/resource-metadata.json';
8+
9+
+// Service name identifier for SageMaker Unified Studio
10+
+export const SMUS_SERVICE_NAME = 'SageMakerUnifiedStudio';
11+
+export const SERVICE_NAME_ENV_VAR = 'SERVICE_NAME';
12+
+
13+
export class SagemakerCookie {
14+
authMode: string
15+
expiryTime: number
16+
@@ -56,6 +60,11 @@ export class SagemakerResourceMetadata {
17+
ResourceArn?: string
18+
ResourceName?: string
19+
AppImageVersion?: string
20+
+ AdditionalMetadata?: {
21+
+ DataZoneDomainId?: string
22+
+ DataZoneProjectId?: string
23+
+ DataZoneDomainRegion?: string
24+
+ }
25+
};
26+
export function isSSOMode(cookie: SagemakerCookie) {
27+
return (cookie.authMode === AUTH_MODE.SSO)
28+
@@ -69,4 +78,35 @@ export function getExpiryTime(cookie: Sa
29+
} else {
30+
return -1;
31+
}
32+
-}
33+
\ No newline at end of file
34+
+}
35+
+
36+
+/**
37+
+ * Constructs the SMUS portal URL using domain, region, and project information
38+
+ * Returns null if not in SMUS environment or if required fields are missing
39+
+ */
40+
+export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata): string | null => {
41+
+ if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) {
42+
+ return null;
43+
+ }
44+
+
45+
+ if (!metadata || !metadata.AdditionalMetadata) {
46+
+ // fail silently not to block users
47+
+ console.error('[SMUS] Metadata is undefined or null');
48+
+ return null;
49+
+ }
50+
+
51+
+ const { DataZoneDomainId, DataZoneDomainRegion, DataZoneProjectId } = metadata.AdditionalMetadata;
52+
+
53+
+ if (!DataZoneDomainId || !DataZoneDomainRegion || !DataZoneProjectId) {
54+
+ // fail silently not to block users
55+
+ // TODO: add monitoring to detect such cases
56+
+ console.error('[SMUS] Required fields missing in metadata:', {
57+
+ DataZoneDomainId: !!DataZoneDomainId,
58+
+ DataZoneDomainRegion: !!DataZoneDomainRegion,
59+
+ DataZoneProjectId: !!DataZoneProjectId
60+
+ });
61+
+ return null;
62+
+ }
63+
+
64+
+ return `https://${DataZoneDomainId}.sagemaker.${DataZoneDomainRegion}.on.aws/projects/${DataZoneProjectId}/overview`;
65+
+}
66+
Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts
67+
===================================================================
68+
--- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/extension.ts
69+
+++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts
70+
@@ -11,13 +11,17 @@ import {
71+
WARNING_BUTTON_SAVE_AND_RENEW_SESSION,
72+
SagemakerCookie,
73+
SagemakerResourceMetadata,
74+
- getExpiryTime
75+
+ getExpiryTime,
76+
+ getSmusVscodePortalUrl
77+
} from "./constant";
78+
import * as console from "console";
79+
80+
81+
const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies';
82+
83+
+// Global redirect URL for SMUS environment
84+
+let smusRedirectUrl: string | null = null;
85+
+
86+
function showWarningDialog() {
87+
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => {
88+
89+
@@ -58,11 +62,12 @@ function showWarningDialog() {
90+
}
91+
92+
function signInError(sagemakerCookie: SagemakerCookie) {
93+
+ const redirectUrl = getRedirectUrl(sagemakerCookie);
94+
// The session has expired
95+
SessionWarning.signInWarning(sagemakerCookie)
96+
.then((selection) => {
97+
if (selection === SIGN_IN_BUTTON) {
98+
- vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
99+
+ vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
100+
}
101+
});
102+
}
103+
@@ -93,32 +98,31 @@ function saveWorkspace() {
104+
});
105+
}
106+
function renewSession(sagemakerCookie: SagemakerCookie) {
107+
+ const redirectUrl = getRedirectUrl(sagemakerCookie);
108+
// TODO: Log and trigger a Signin
109+
- vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
110+
+ vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
111+
// Trigger the function to show the warning again after 5 minutes again to validate.
112+
setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS);
113+
}
114+
115+
function updateStatusItemWithMetadata(context: vscode.ExtensionContext) {
116+
- fs.readFile(SAGEMAKER_METADATA_PATH, 'utf-8', (err, data) => {
117+
- if (err) {
118+
- // fail silently not to block users
119+
- } else {
120+
- try {
121+
- const jsonData = JSON.parse(data) as SagemakerResourceMetadata;
122+
- const spaceName = jsonData.SpaceName;
123+
-
124+
- if (spaceName != null) {
125+
- let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
126+
- spaceNameStatusBarItem.text = `Space: ${spaceName}`;
127+
- spaceNameStatusBarItem.show();
128+
- context.subscriptions.push(spaceNameStatusBarItem);
129+
- }
130+
- } catch (jsonError) {
131+
- // fail silently not to block users
132+
- }
133+
+ try {
134+
+ const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8');
135+
+ const jsonData = JSON.parse(data) as SagemakerResourceMetadata;
136+
+
137+
+ if (jsonData?.SpaceName) {
138+
+ let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100);
139+
+ spaceNameStatusBarItem.text = `Space: ${jsonData.SpaceName}`;
140+
+ spaceNameStatusBarItem.show();
141+
+ context.subscriptions.push(spaceNameStatusBarItem);
142+
}
143+
- });
144+
+
145+
+ // Initialize SMUS redirect URL
146+
+ smusRedirectUrl = getSmusVscodePortalUrl(jsonData);
147+
+ } catch (error) {
148+
+ // fail silently not to block users
149+
+ console.error('Error reading metadata file:', error);
150+
+ }
151+
}
152+
153+
export function activate(context: vscode.ExtensionContext) {
154+
@@ -126,12 +130,22 @@ export function activate(context: vscode
155+
// TODO: log activation of extension
156+
console.log('Activating Sagemaker Extension...');
157+
158+
+ // Initialize metadata first (which will set smusRedirectUrl if in SMUS environment)
159+
+ updateStatusItemWithMetadata(context);
160+
+
161+
// execute the get cookie command and save the data to cookies
162+
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => {
163+
164+
const sagemakerCookie: SagemakerCookie = r as SagemakerCookie
165+
166+
initialize(sagemakerCookie);
167+
- updateStatusItemWithMetadata(context);
168+
});
169+
}
170+
+
171+
+/**
172+
+ * Returns the appropriate redirect URL based on the environment
173+
+ * Uses SMUS URL if available, falls back to original redirect URL
174+
+ */
175+
+function getRedirectUrl(sagemakerCookie: SagemakerCookie): string {
176+
+ return smusRedirectUrl || sagemakerCookie.redirectURL;
177+
+}
178+
\ No newline at end of file

‎patches/series‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,3 +11,4 @@ sagemaker-idle-extension.patch
1111
terminal-crash-mitigation.patch
1212
sagemaker-open-notebook-extension.patch
1313
security.diff
14+
sagemaker-extension-smus-support.patch

0 commit comments

Comments
 (0)