Skip to content

Commit 097a7ba

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 c2a7812 commit 097a7ba

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,13 +11,17 @@ 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

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

@@ -58,11 +62,12 @@ function showWarningDialog() {
5862
}
5963

6064
function signInError(sagemakerCookie: SagemakerCookie) {
65+
const redirectUrl = getRedirectUrl(sagemakerCookie);
6166
// The session has expired
6267
SessionWarning.signInWarning(sagemakerCookie)
6368
.then((selection) => {
6469
if (selection === SIGN_IN_BUTTON) {
65-
vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
70+
vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
6671
}
6772
});
6873
}
@@ -93,45 +98,54 @@ function saveWorkspace() {
9398
});
9499
}
95100
function renewSession(sagemakerCookie: SagemakerCookie) {
101+
const redirectUrl = getRedirectUrl(sagemakerCookie);
96102
// TODO: Log and trigger a Signin
97-
vscode.env.openExternal(vscode.Uri.parse(<string>sagemakerCookie.redirectURL));
103+
vscode.env.openExternal(vscode.Uri.parse(redirectUrl));
98104
// Trigger the function to show the warning again after 5 minutes again to validate.
99105
setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS);
100106
}
101107

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

124128
export function activate(context: vscode.ExtensionContext) {
125129

126130
// TODO: log activation of extension
127131
console.log('Activating Sagemaker Extension...');
128132

133+
// Initialize metadata first (which will set smusRedirectUrl if in SMUS environment)
134+
updateStatusItemWithMetadata(context);
135+
129136
// execute the get cookie command and save the data to cookies
130137
vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => {
131138

132139
const sagemakerCookie: SagemakerCookie = r as SagemakerCookie
133140

134141
initialize(sagemakerCookie);
135-
updateStatusItemWithMetadata(context);
136142
});
137143
}
144+
145+
/**
146+
* Returns the appropriate redirect URL based on the environment
147+
* Uses SMUS URL if available, falls back to original redirect URL
148+
*/
149+
function getRedirectUrl(sagemakerCookie: SagemakerCookie): string {
150+
return smusRedirectUrl || sagemakerCookie.redirectURL;
151+
}

‎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
@@ -9,3 +9,4 @@ license.diff
99
sagemaker-idle-extension.patch
1010
terminal-crash-mitigation.patch
1111
sagemaker-open-notebook-extension.patch
12+
sagemaker-extension-smus-support.patch

0 commit comments

Comments
 (0)