From 3983b4d1d2d0cbd38b9941911e3f4f66997047a7 Mon Sep 17 00:00:00 2001 From: Arun Bhati Date: Fri, 21 Mar 2025 16:31:08 -0700 Subject: [PATCH 1/3] feat(session): add SageMaker Unified Studio portal redirect MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ** 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 --- .../sagemaker-extension/src/constant.ts | 42 ++++- .../sagemaker-extension/src/extension.ts | 58 +++--- .../sagemaker-extension-smus-support.patch | 178 ++++++++++++++++++ patches/series | 1 + 4 files changed, 256 insertions(+), 23 deletions(-) create mode 100644 patches/sagemaker-extension-smus-support.patch diff --git a/patched-vscode/extensions/sagemaker-extension/src/constant.ts b/patched-vscode/extensions/sagemaker-extension/src/constant.ts index 3417eefb0..684b75ca8 100644 --- a/patched-vscode/extensions/sagemaker-extension/src/constant.ts +++ b/patched-vscode/extensions/sagemaker-extension/src/constant.ts @@ -27,6 +27,10 @@ export const FIVE_MINUTES_INTERVAL_MILLIS = 5 * 60 * 1000; export const SAGEMAKER_METADATA_PATH = '/opt/ml/metadata/resource-metadata.json'; +// Service name identifier for SageMaker Unified Studio +export const SMUS_SERVICE_NAME = 'SageMakerUnifiedStudio'; +export const SERVICE_NAME_ENV_VAR = 'SERVICE_NAME'; + export class SagemakerCookie { authMode: string expiryTime: number @@ -56,6 +60,11 @@ export class SagemakerResourceMetadata { ResourceArn?: string ResourceName?: string AppImageVersion?: string + AdditionalMetadata?: { + DataZoneDomainId?: string + DataZoneProjectId?: string + DataZoneDomainRegion?: string + } }; export function isSSOMode(cookie: SagemakerCookie) { return (cookie.authMode === AUTH_MODE.SSO) @@ -69,4 +78,35 @@ export function getExpiryTime(cookie: SagemakerCookie): number { } else { return -1; } -} \ No newline at end of file +} + +/** + * Constructs the SMUS portal URL using domain, region, and project information + * Returns null if not in SMUS environment or if required fields are missing + */ +export const getSmusVscodePortalUrl = (metadata: SagemakerMetadata): string | null => { + if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) { + return null; + } + + if (!metadata || !metadata.AdditionalMetadata) { + // fail silently not to block users + console.error('[SMUS] Metadata is undefined or null'); + return null; + } + + const { DataZoneDomainId, DataZoneDomainRegion, DataZoneProjectId } = metadata.AdditionalMetadata; + + if (!DataZoneDomainId || !DataZoneDomainRegion || !DataZoneProjectId) { + // fail silently not to block users + // TODO: add monitoring to detect such cases + console.error('[SMUS] Required fields missing in metadata:', { + DataZoneDomainId: !!DataZoneDomainId, + DataZoneDomainRegion: !!DataZoneDomainRegion, + DataZoneProjectId: !!DataZoneProjectId + }); + return null; + } + + return `https://${DataZoneDomainId}.sagemaker.${DataZoneDomainRegion}.on.aws/projects/${DataZoneProjectId}/overview`; +} diff --git a/patched-vscode/extensions/sagemaker-extension/src/extension.ts b/patched-vscode/extensions/sagemaker-extension/src/extension.ts index a64f8e9cd..3f6255504 100644 --- a/patched-vscode/extensions/sagemaker-extension/src/extension.ts +++ b/patched-vscode/extensions/sagemaker-extension/src/extension.ts @@ -11,7 +11,8 @@ import { WARNING_BUTTON_SAVE_AND_RENEW_SESSION, SagemakerCookie, SagemakerResourceMetadata, - getExpiryTime + getExpiryTime, + getSmusVscodePortalUrl } from "./constant"; import * as console from "console"; @@ -19,6 +20,9 @@ import * as console from "console"; const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies'; const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate'; +// Global redirect URL for SMUS environment +let smusRedirectUrl: string | null = null; + function showWarningDialog() { vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => { @@ -59,11 +63,12 @@ function showWarningDialog() { } function signInError(sagemakerCookie: SagemakerCookie) { + const redirectUrl = getRedirectUrl(sagemakerCookie); // The session has expired SessionWarning.signInWarning(sagemakerCookie) .then((selection) => { if (selection === SIGN_IN_BUTTON) { - vscode.env.openExternal(vscode.Uri.parse(sagemakerCookie.redirectURL)); + vscode.env.openExternal(vscode.Uri.parse(redirectUrl)); } }); } @@ -94,32 +99,31 @@ function saveWorkspace() { }); } function renewSession(sagemakerCookie: SagemakerCookie) { + const redirectUrl = getRedirectUrl(sagemakerCookie); // TODO: Log and trigger a Signin - vscode.env.openExternal(vscode.Uri.parse(sagemakerCookie.redirectURL)); + vscode.env.openExternal(vscode.Uri.parse(redirectUrl)); // Trigger the function to show the warning again after 5 minutes again to validate. setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS); } function updateStatusItemWithMetadata(context: vscode.ExtensionContext) { - fs.readFile(SAGEMAKER_METADATA_PATH, 'utf-8', (err, data) => { - if (err) { - // fail silently not to block users - } else { - try { - const jsonData = JSON.parse(data) as SagemakerResourceMetadata; - const spaceName = jsonData.SpaceName; - - if (spaceName != null) { - let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); - spaceNameStatusBarItem.text = `Space: ${spaceName}`; - spaceNameStatusBarItem.show(); - context.subscriptions.push(spaceNameStatusBarItem); - } - } catch (jsonError) { - // fail silently not to block users - } + try { + const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8'); + const jsonData = JSON.parse(data) as SagemakerResourceMetadata; + + if (jsonData?.SpaceName) { + let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); + spaceNameStatusBarItem.text = `Space: ${jsonData.SpaceName}`; + spaceNameStatusBarItem.show(); + context.subscriptions.push(spaceNameStatusBarItem); } - }); + + // Initialize SMUS redirect URL + smusRedirectUrl = getSmusVscodePortalUrl(jsonData); + } catch (error) { + // fail silently not to block users + console.error('Error reading metadata file:', error); + } } // Render warning message regarding auto upgrade disabled @@ -158,15 +162,25 @@ export function activate(context: vscode.ExtensionContext) { // TODO: log activation of extension console.log('Activating Sagemaker Extension...'); + // Initialize metadata first (which will set smusRedirectUrl if in SMUS environment) + updateStatusItemWithMetadata(context); + // execute the get cookie command and save the data to cookies vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => { const sagemakerCookie: SagemakerCookie = r as SagemakerCookie initialize(sagemakerCookie); - updateStatusItemWithMetadata(context); }); // render warning message regarding auto upgrade disabled renderExtensionAutoUpgradeDisabledNotification(); } + +/** + * Returns the appropriate redirect URL based on the environment + * Uses SMUS URL if available, falls back to original redirect URL + */ +function getRedirectUrl(sagemakerCookie: SagemakerCookie): string { + return smusRedirectUrl || sagemakerCookie.redirectURL; +} diff --git a/patches/sagemaker-extension-smus-support.patch b/patches/sagemaker-extension-smus-support.patch new file mode 100644 index 000000000..8a7a32318 --- /dev/null +++ b/patches/sagemaker-extension-smus-support.patch @@ -0,0 +1,178 @@ +Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant.ts +=================================================================== +--- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/constant.ts ++++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant.ts +@@ -27,6 +27,10 @@ export const FIVE_MINUTES_INTERVAL_MILLI + + export const SAGEMAKER_METADATA_PATH = '/opt/ml/metadata/resource-metadata.json'; + ++// Service name identifier for SageMaker Unified Studio ++export const SMUS_SERVICE_NAME = 'SageMakerUnifiedStudio'; ++export const SERVICE_NAME_ENV_VAR = 'SERVICE_NAME'; ++ + export class SagemakerCookie { + authMode: string + expiryTime: number +@@ -56,6 +60,11 @@ export class SagemakerResourceMetadata { + ResourceArn?: string + ResourceName?: string + AppImageVersion?: string ++ AdditionalMetadata?: { ++ DataZoneDomainId?: string ++ DataZoneProjectId?: string ++ DataZoneDomainRegion?: string ++ } + }; + export function isSSOMode(cookie: SagemakerCookie) { + return (cookie.authMode === AUTH_MODE.SSO) +@@ -69,4 +78,35 @@ export function getExpiryTime(cookie: Sa + } else { + return -1; + } +-} +\ No newline at end of file ++} ++ ++/** ++ * Constructs the SMUS portal URL using domain, region, and project information ++ * Returns null if not in SMUS environment or if required fields are missing ++ */ ++export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata): string | null => { ++ if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) { ++ return null; ++ } ++ ++ if (!metadata || !metadata.AdditionalMetadata) { ++ // fail silently not to block users ++ console.error('[SMUS] Metadata is undefined or null'); ++ return null; ++ } ++ ++ const { DataZoneDomainId, DataZoneDomainRegion, DataZoneProjectId } = metadata.AdditionalMetadata; ++ ++ if (!DataZoneDomainId || !DataZoneDomainRegion || !DataZoneProjectId) { ++ // fail silently not to block users ++ // TODO: add monitoring to detect such cases ++ console.error('[SMUS] Required fields missing in metadata:', { ++ DataZoneDomainId: !!DataZoneDomainId, ++ DataZoneDomainRegion: !!DataZoneDomainRegion, ++ DataZoneProjectId: !!DataZoneProjectId ++ }); ++ return null; ++ } ++ ++ return `https://${DataZoneDomainId}.sagemaker.${DataZoneDomainRegion}.on.aws/projects/${DataZoneProjectId}/overview`; ++} +Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts +=================================================================== +--- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/extension.ts ++++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts +@@ -11,13 +11,17 @@ import { + WARNING_BUTTON_SAVE_AND_RENEW_SESSION, + SagemakerCookie, + SagemakerResourceMetadata, +- getExpiryTime ++ getExpiryTime, ++ getSmusVscodePortalUrl + } from "./constant"; + import * as console from "console"; + + + const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies'; + ++// Global redirect URL for SMUS environment ++let smusRedirectUrl: string | null = null; ++ + function showWarningDialog() { + vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => { + +@@ -58,11 +62,12 @@ function showWarningDialog() { + } + + function signInError(sagemakerCookie: SagemakerCookie) { ++ const redirectUrl = getRedirectUrl(sagemakerCookie); + // The session has expired + SessionWarning.signInWarning(sagemakerCookie) + .then((selection) => { + if (selection === SIGN_IN_BUTTON) { +- vscode.env.openExternal(vscode.Uri.parse(sagemakerCookie.redirectURL)); ++ vscode.env.openExternal(vscode.Uri.parse(redirectUrl)); + } + }); + } +@@ -93,32 +98,31 @@ function saveWorkspace() { + }); + } + function renewSession(sagemakerCookie: SagemakerCookie) { ++ const redirectUrl = getRedirectUrl(sagemakerCookie); + // TODO: Log and trigger a Signin +- vscode.env.openExternal(vscode.Uri.parse(sagemakerCookie.redirectURL)); ++ vscode.env.openExternal(vscode.Uri.parse(redirectUrl)); + // Trigger the function to show the warning again after 5 minutes again to validate. + setTimeout(showWarningDialog, FIVE_MINUTES_INTERVAL_MILLIS); + } + + function updateStatusItemWithMetadata(context: vscode.ExtensionContext) { +- fs.readFile(SAGEMAKER_METADATA_PATH, 'utf-8', (err, data) => { +- if (err) { +- // fail silently not to block users +- } else { +- try { +- const jsonData = JSON.parse(data) as SagemakerResourceMetadata; +- const spaceName = jsonData.SpaceName; +- +- if (spaceName != null) { +- let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); +- spaceNameStatusBarItem.text = `Space: ${spaceName}`; +- spaceNameStatusBarItem.show(); +- context.subscriptions.push(spaceNameStatusBarItem); +- } +- } catch (jsonError) { +- // fail silently not to block users +- } ++ try { ++ const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8'); ++ const jsonData = JSON.parse(data) as SagemakerResourceMetadata; ++ ++ if (jsonData?.SpaceName) { ++ let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); ++ spaceNameStatusBarItem.text = `Space: ${jsonData.SpaceName}`; ++ spaceNameStatusBarItem.show(); ++ context.subscriptions.push(spaceNameStatusBarItem); + } +- }); ++ ++ // Initialize SMUS redirect URL ++ smusRedirectUrl = getSmusVscodePortalUrl(jsonData); ++ } catch (error) { ++ // fail silently not to block users ++ console.error('Error reading metadata file:', error); ++ } + } + + export function activate(context: vscode.ExtensionContext) { +@@ -126,12 +130,22 @@ export function activate(context: vscode + // TODO: log activation of extension + console.log('Activating Sagemaker Extension...'); + ++ // Initialize metadata first (which will set smusRedirectUrl if in SMUS environment) ++ updateStatusItemWithMetadata(context); ++ + // execute the get cookie command and save the data to cookies + vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => { + + const sagemakerCookie: SagemakerCookie = r as SagemakerCookie + + initialize(sagemakerCookie); +- updateStatusItemWithMetadata(context); + }); + } ++ ++/** ++ * Returns the appropriate redirect URL based on the environment ++ * Uses SMUS URL if available, falls back to original redirect URL ++ */ ++function getRedirectUrl(sagemakerCookie: SagemakerCookie): string { ++ return smusRedirectUrl || sagemakerCookie.redirectURL; ++} +\ No newline at end of file diff --git a/patches/series b/patches/series index 7c3c12166..d70f9b432 100644 --- a/patches/series +++ b/patches/series @@ -11,3 +11,4 @@ sagemaker-idle-extension.patch terminal-crash-mitigation.patch sagemaker-open-notebook-extension.patch security.diff +sagemaker-extension-smus-support.patch From d864d995bc60b696d6dc45e1a7c9018d02ec62c0 Mon Sep 17 00:00:00 2001 From: Arun Bhati Date: Thu, 27 Mar 2025 13:10:07 -0700 Subject: [PATCH 2/3] Fix patch merge failures from feature/smus remote branch --- .../sagemaker-extension/src/constant.ts | 2 +- patches/sagemaker-extension-smus-support.patch | 17 ++++++++++------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/patched-vscode/extensions/sagemaker-extension/src/constant.ts b/patched-vscode/extensions/sagemaker-extension/src/constant.ts index 684b75ca8..7aabd073d 100644 --- a/patched-vscode/extensions/sagemaker-extension/src/constant.ts +++ b/patched-vscode/extensions/sagemaker-extension/src/constant.ts @@ -84,7 +84,7 @@ export function getExpiryTime(cookie: SagemakerCookie): number { * Constructs the SMUS portal URL using domain, region, and project information * Returns null if not in SMUS environment or if required fields are missing */ -export const getSmusVscodePortalUrl = (metadata: SagemakerMetadata): string | null => { +export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata): string | null => { if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) { return null; } diff --git a/patches/sagemaker-extension-smus-support.patch b/patches/sagemaker-extension-smus-support.patch index 8a7a32318..55ac357f4 100644 --- a/patches/sagemaker-extension-smus-support.patch +++ b/patches/sagemaker-extension-smus-support.patch @@ -67,7 +67,7 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension =================================================================== --- sagemaker-code-editor.orig/vscode/extensions/sagemaker-extension/src/extension.ts +++ sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension.ts -@@ -11,13 +11,17 @@ import { +@@ -11,7 +11,8 @@ import { WARNING_BUTTON_SAVE_AND_RENEW_SESSION, SagemakerCookie, SagemakerResourceMetadata, @@ -77,8 +77,9 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension } from "./constant"; import * as console from "console"; - +@@ -19,6 +20,9 @@ import * as console from "console"; const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies'; + const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate'; +// Global redirect URL for SMUS environment +let smusRedirectUrl: string | null = null; @@ -86,7 +87,7 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension function showWarningDialog() { vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => { -@@ -58,11 +62,12 @@ function showWarningDialog() { +@@ -59,11 +63,12 @@ function showWarningDialog() { } function signInError(sagemakerCookie: SagemakerCookie) { @@ -100,7 +101,7 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension } }); } -@@ -93,32 +98,31 @@ function saveWorkspace() { +@@ -94,32 +99,31 @@ function saveWorkspace() { }); } function renewSession(sagemakerCookie: SagemakerCookie) { @@ -150,8 +151,8 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension + } } - export function activate(context: vscode.ExtensionContext) { -@@ -126,12 +130,22 @@ export function activate(context: vscode + // Render warning message regarding auto upgrade disabled +@@ -158,15 +162,25 @@ export function activate(context: vscode // TODO: log activation of extension console.log('Activating Sagemaker Extension...'); @@ -166,6 +167,9 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension initialize(sagemakerCookie); - updateStatusItemWithMetadata(context); }); + + // render warning message regarding auto upgrade disabled + renderExtensionAutoUpgradeDisabledNotification(); } + +/** @@ -175,4 +179,3 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension +function getRedirectUrl(sagemakerCookie: SagemakerCookie): string { + return smusRedirectUrl || sagemakerCookie.redirectURL; +} -\ No newline at end of file From 987c4551eaca41a985fc2943e49ed21bbed8fffb Mon Sep 17 00:00:00 2001 From: Arun Bhati Date: Thu, 27 Mar 2025 13:50:57 -0700 Subject: [PATCH 3/3] Move resource metadata parsing logic in helper function --- .../sagemaker-extension/src/constant.ts | 2 +- .../sagemaker-extension/src/extension.ts | 42 +++++++------ .../sagemaker-extension-smus-support.patch | 59 ++++++++++--------- 3 files changed, 55 insertions(+), 48 deletions(-) diff --git a/patched-vscode/extensions/sagemaker-extension/src/constant.ts b/patched-vscode/extensions/sagemaker-extension/src/constant.ts index 7aabd073d..ea1d482e4 100644 --- a/patched-vscode/extensions/sagemaker-extension/src/constant.ts +++ b/patched-vscode/extensions/sagemaker-extension/src/constant.ts @@ -84,7 +84,7 @@ export function getExpiryTime(cookie: SagemakerCookie): number { * Constructs the SMUS portal URL using domain, region, and project information * Returns null if not in SMUS environment or if required fields are missing */ -export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata): string | null => { +export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata | null): string | null => { if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) { return null; } diff --git a/patched-vscode/extensions/sagemaker-extension/src/extension.ts b/patched-vscode/extensions/sagemaker-extension/src/extension.ts index 3f6255504..e40febdd7 100644 --- a/patched-vscode/extensions/sagemaker-extension/src/extension.ts +++ b/patched-vscode/extensions/sagemaker-extension/src/extension.ts @@ -23,6 +23,21 @@ const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate // Global redirect URL for SMUS environment let smusRedirectUrl: string | null = null; +function fetchMetadata(): SagemakerResourceMetadata | null { + try { + const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8'); + return JSON.parse(data) as SagemakerResourceMetadata; + } catch (error) { + // fail silently not to block users + console.error('Error reading metadata file:', error); + return null; + } +} + +function initializeSmusRedirectUrl() { + smusRedirectUrl = getSmusVscodePortalUrl(fetchMetadata()); +} + function showWarningDialog() { vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => { @@ -107,22 +122,12 @@ function renewSession(sagemakerCookie: SagemakerCookie) { } function updateStatusItemWithMetadata(context: vscode.ExtensionContext) { - try { - const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8'); - const jsonData = JSON.parse(data) as SagemakerResourceMetadata; - - if (jsonData?.SpaceName) { - let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); - spaceNameStatusBarItem.text = `Space: ${jsonData.SpaceName}`; - spaceNameStatusBarItem.show(); - context.subscriptions.push(spaceNameStatusBarItem); - } - - // Initialize SMUS redirect URL - smusRedirectUrl = getSmusVscodePortalUrl(jsonData); - } catch (error) { - // fail silently not to block users - console.error('Error reading metadata file:', error); + const metadata = fetchMetadata(); + if (metadata?.SpaceName) { + let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); + spaceNameStatusBarItem.text = `Space: ${metadata.SpaceName}`; + spaceNameStatusBarItem.show(); + context.subscriptions.push(spaceNameStatusBarItem); } } @@ -162,8 +167,8 @@ export function activate(context: vscode.ExtensionContext) { // TODO: log activation of extension console.log('Activating Sagemaker Extension...'); - // Initialize metadata first (which will set smusRedirectUrl if in SMUS environment) - updateStatusItemWithMetadata(context); + // First set smusRedirectUrl if we are in SMUS environment + initializeSmusRedirectUrl(); // execute the get cookie command and save the data to cookies vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => { @@ -171,6 +176,7 @@ export function activate(context: vscode.ExtensionContext) { const sagemakerCookie: SagemakerCookie = r as SagemakerCookie initialize(sagemakerCookie); + updateStatusItemWithMetadata(context); }); // render warning message regarding auto upgrade disabled diff --git a/patches/sagemaker-extension-smus-support.patch b/patches/sagemaker-extension-smus-support.patch index 55ac357f4..303039bbc 100644 --- a/patches/sagemaker-extension-smus-support.patch +++ b/patches/sagemaker-extension-smus-support.patch @@ -37,7 +37,7 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/constant. + * Constructs the SMUS portal URL using domain, region, and project information + * Returns null if not in SMUS environment or if required fields are missing + */ -+export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata): string | null => { ++export const getSmusVscodePortalUrl = (metadata: SagemakerResourceMetadata | null): string | null => { + if (process.env[SERVICE_NAME_ENV_VAR] !== SMUS_SERVICE_NAME) { + return null; + } @@ -77,17 +77,32 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension } from "./constant"; import * as console from "console"; -@@ -19,6 +20,9 @@ import * as console from "console"; +@@ -19,6 +20,24 @@ import * as console from "console"; const PARSE_SAGEMAKER_COOKIE_COMMAND = 'sagemaker.parseCookies'; const ENABLE_AUTO_UPDATE_COMMAND = 'workbench.extensions.action.enableAutoUpdate'; +// Global redirect URL for SMUS environment +let smusRedirectUrl: string | null = null; ++ ++function fetchMetadata(): SagemakerResourceMetadata | null { ++ try { ++ const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8'); ++ return JSON.parse(data) as SagemakerResourceMetadata; ++ } catch (error) { ++ // fail silently not to block users ++ console.error('Error reading metadata file:', error); ++ return null; ++ } ++} ++ ++function initializeSmusRedirectUrl() { ++ smusRedirectUrl = getSmusVscodePortalUrl(fetchMetadata()); ++} + function showWarningDialog() { vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(response => { -@@ -59,11 +63,12 @@ function showWarningDialog() { +@@ -59,11 +78,12 @@ function showWarningDialog() { } function signInError(sagemakerCookie: SagemakerCookie) { @@ -101,7 +116,7 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension } }); } -@@ -94,32 +99,31 @@ function saveWorkspace() { +@@ -94,32 +114,21 @@ function saveWorkspace() { }); } function renewSession(sagemakerCookie: SagemakerCookie) { @@ -131,43 +146,29 @@ Index: sagemaker-code-editor/vscode/extensions/sagemaker-extension/src/extension - } catch (jsonError) { - // fail silently not to block users - } -+ try { -+ const data = fs.readFileSync(SAGEMAKER_METADATA_PATH, 'utf-8'); -+ const jsonData = JSON.parse(data) as SagemakerResourceMetadata; -+ -+ if (jsonData?.SpaceName) { -+ let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); -+ spaceNameStatusBarItem.text = `Space: ${jsonData.SpaceName}`; -+ spaceNameStatusBarItem.show(); -+ context.subscriptions.push(spaceNameStatusBarItem); - } +- } - }); -+ -+ // Initialize SMUS redirect URL -+ smusRedirectUrl = getSmusVscodePortalUrl(jsonData); -+ } catch (error) { -+ // fail silently not to block users -+ console.error('Error reading metadata file:', error); ++ const metadata = fetchMetadata(); ++ if (metadata?.SpaceName) { ++ let spaceNameStatusBarItem = vscode.window.createStatusBarItem(vscode.StatusBarAlignment.Left, 100); ++ spaceNameStatusBarItem.text = `Space: ${metadata.SpaceName}`; ++ spaceNameStatusBarItem.show(); ++ context.subscriptions.push(spaceNameStatusBarItem); + } } // Render warning message regarding auto upgrade disabled -@@ -158,15 +162,25 @@ export function activate(context: vscode +@@ -158,6 +167,9 @@ export function activate(context: vscode // TODO: log activation of extension console.log('Activating Sagemaker Extension...'); -+ // Initialize metadata first (which will set smusRedirectUrl if in SMUS environment) -+ updateStatusItemWithMetadata(context); ++ // First set smusRedirectUrl if we are in SMUS environment ++ initializeSmusRedirectUrl(); + // execute the get cookie command and save the data to cookies vscode.commands.executeCommand(PARSE_SAGEMAKER_COOKIE_COMMAND).then(r => { - const sagemakerCookie: SagemakerCookie = r as SagemakerCookie - - initialize(sagemakerCookie); -- updateStatusItemWithMetadata(context); - }); - +@@ -170,3 +182,11 @@ export function activate(context: vscode // render warning message regarding auto upgrade disabled renderExtensionAutoUpgradeDisabledNotification(); }