diff --git a/patched-vscode/extensions/sagemaker-extension/src/constant.ts b/patched-vscode/extensions/sagemaker-extension/src/constant.ts index 3417eefb0..ea1d482e4 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: SagemakerResourceMetadata | null): 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..e40febdd7 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,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 +78,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 +114,21 @@ 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 - } - } - }); + 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,6 +167,9 @@ export function activate(context: vscode.ExtensionContext) { // TODO: log activation of extension console.log('Activating Sagemaker Extension...'); + // 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 => { @@ -170,3 +182,11 @@ export function activate(context: vscode.ExtensionContext) { // 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..303039bbc --- /dev/null +++ b/patches/sagemaker-extension-smus-support.patch @@ -0,0 +1,182 @@ +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 | null): 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,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,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 +78,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 +114,21 @@ 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 +- } +- } +- }); ++ 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,6 +167,9 @@ export function activate(context: vscode + // TODO: log activation of extension + console.log('Activating Sagemaker Extension...'); + ++ // 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 => { + +@@ -170,3 +182,11 @@ export function activate(context: vscode + // 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/series b/patches/series index fc0b16b2f..d8b108499 100644 --- a/patches/series +++ b/patches/series @@ -13,3 +13,4 @@ sagemaker-open-notebook-extension.patch security.diff sagemaker-ui-dark-theme.patch sagemaker-ui-post-startup.patch +sagemaker-extension-smus-support.patch