fix(amazonq): fix to move isSagemaker to env

Author: ashishrp-aws

packages/core/src/shared/extensionUtilities.ts

@@ -11,7 +11,7 @@ import { getLogger } from './logger/logger'
 import { VSCODE_EXTENSION_ID, extensionAlphaVersion } from './extensions'
 import { Ec2MetadataClient } from './clients/ec2MetadataClient'
 import { DefaultEc2MetadataClient } from './clients/ec2MetadataClient'
-import { extensionVersion, getCodeCatalystDevEnvId } from './vscode/env'
+import { extensionVersion, getCodeCatalystDevEnvId, hasSageMakerEnvVars } from './vscode/env'
 import globals from './extensionGlobals'
 import { once } from './utilities/functionUtils'
 import {
@@ -170,26 +170,14 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo
     return (flavor === 'classic' && !codecat) || (flavor === 'codecatalyst' && codecat)
 }
 
-/**
- * Checks if the current environment has SageMaker-specific environment variables
- * @returns true if SageMaker environment variables are detected
- */
-function hasSageMakerEnvVars(): boolean {
-    return (
-        process.env.SAGEMAKER_APP_TYPE !== undefined ||
-        process.env.SAGEMAKER_INTERNAL_IMAGE_URI !== undefined ||
-        process.env.STUDIO_LOGGING_DIR?.includes('/var/log/studio') === true
-    )
-}
-
 /**
  *
  * @param appName to identify the proper SM instance
  * @returns true if the current system is SageMaker(SMAI or SMUS)
  */
 export function isSageMaker(appName: 'SMAI' | 'SMUS' = 'SMAI'): boolean {
     // Check for SageMaker-specific environment variables first
-    if (hasSageMakerEnvVars() || process.env.SERVICE_NAME === sageMakerUnifiedStudio) {
+    if (hasSageMakerEnvVars()) {
         getLogger().debug('SageMaker environment detected via environment variables')
         return true
     }

packages/core/src/shared/vscode/env.ts

@@ -124,6 +124,25 @@ export function isRemoteWorkspace(): boolean {
     return vscode.env.remoteName === 'ssh-remote'
 }
 
+/**
+ * Checks if the current environment has SageMaker-specific environment variables
+ * @returns true if SageMaker environment variables are detected
+ */
+export function hasSageMakerEnvVars(): boolean {
+    // Check both old and new environment variable names
+    // SageMaker is renaming their environment variables in their Docker images
+    return (
+        // Original environment variables
+        process.env.SAGEMAKER_APP_TYPE !== undefined ||
+        process.env.SAGEMAKER_INTERNAL_IMAGE_URI !== undefined ||
+        process.env.STUDIO_LOGGING_DIR?.includes('/var/log/studio') === true ||
+        // New environment variables (update these with the actual new names)
+        process.env.SM_APP_TYPE !== undefined ||
+        process.env.SM_INTERNAL_IMAGE_URI !== undefined ||
+        process.env.SERVICE_NAME === 'SageMakerUnifiedStudio'
+    )
+}
+
 /**
  * Checks if the current environment is running on Amazon Linux 2.
  *
@@ -135,11 +154,7 @@ export function isRemoteWorkspace(): boolean {
 export function isAmazonLinux2() {
     // First check if we're in a SageMaker environment, which should not be treated as AL2
     // even if the underlying host is AL2
-    if (
-        process.env.SAGEMAKER_APP_TYPE ||
-        process.env.SERVICE_NAME === 'SageMakerUnifiedStudio' ||
-        process.env.SAGEMAKER_INTERNAL_IMAGE_URI
-    ) {
+    if (hasSageMakerEnvVars()) {
         return false
     }