Update context var setup in extension activation for both node and web

instead of chat activation, utils refactored

Author: parameja1

packages/amazonq/src/app/chat/activation.ts

@@ -7,7 +7,7 @@ import * as vscode from 'vscode'
 import { ExtensionContext, window } from 'vscode'
 import { telemetry } from 'aws-core-vscode/telemetry'
 import { AuthUtil, CodeWhispererSettings } from 'aws-core-vscode/codewhisperer'
-import { Commands, placeholder, funcUtil, setContext, isSageMaker } from 'aws-core-vscode/shared'
+import { Commands, placeholder, funcUtil } from 'aws-core-vscode/shared'
 import * as amazonq from 'aws-core-vscode/amazonq'
 import { scanChatAppInit } from '../amazonqScan'
 import { init as inlineChatInit } from '../../inlineChat/app'
@@ -16,7 +16,6 @@ export async function activate(context: ExtensionContext) {
     const appInitContext = amazonq.DefaultAmazonQAppInitContext.instance
 
     registerApps(appInitContext, context)
-    await setIsSMUSContext()
 
     const provider = new amazonq.AmazonQChatViewProvider(
         context,
@@ -64,12 +63,6 @@ export async function activate(context: ExtensionContext) {
 
     void setupLsp()
     void setupAuthNotification()
-
-    async function setIsSMUSContext() {
-        // set context key for SageMaker Unified Studio
-        const isSMUS = isSageMaker(true)
-        await setContext('aws.isSageMakerUnifiedStudio', isSMUS)
-    }
 }
 
 function registerApps(appInitContext: amazonq.AmazonQAppInitContext, context: ExtensionContext) {

packages/amazonq/src/extension.ts

@@ -32,6 +32,7 @@ import {
     setupUninstallHandler,
     maybeShowMinVscodeWarning,
     Experiments,
+    isSageMaker,
 } from 'aws-core-vscode/shared'
 import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
 import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -146,6 +147,9 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
     // Hide the Amazon Q tree in toolkit explorer
     await setContext('aws.toolkit.amazonq.dismissed', true)
 
+    // set context var to check if its SageMaker Unified Studio or not
+    await setContext('aws.isSageMakerUnifiedStudio', isSageMaker('SMUS'))
+
     // reload webviews
     await vscode.commands.executeCommand('workbench.action.webview.reloadWebviewAction')
 

packages/core/src/amazonq/webview/generators/webViewContent.ts

@@ -83,7 +83,7 @@ export class WebViewContentGenerator {
         const disclaimerAcknowledged = globals.globalState.tryGet('aws.amazonq.disclaimerAcknowledged', Boolean, false)
 
         const welcomeLoadCount = globals.globalState.tryGet('aws.amazonq.welcomeChatShowCount', Number, 0)
-        const serviceName = process.env.SERVICE_NAME ?? ''
+        const isSMUS = isSageMaker('SMUS')
 
         return `
         <script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>
@@ -92,7 +92,7 @@ export class WebViewContentGenerator {
             const init = () => {
                 createMynahUI(acquireVsCodeApi(), ${
                     (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
-                },${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString}, '${serviceName}');
+                },${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString},${isSMUS});
             }
         </script>
         `

packages/core/src/amazonq/webview/ui/main.ts

@@ -48,7 +48,7 @@ export const createMynahUI = (
     welcomeCount: number,
     disclaimerAcknowledged: boolean,
     disabledCommands?: string[],
-    service?: string
+    isSMUS?: boolean
 ) => {
     let disclaimerCardActive = !disclaimerAcknowledged
     // eslint-disable-next-line prefer-const
@@ -59,7 +59,6 @@ export const createMynahUI = (
     const responseMetadata = new Map<string, string[]>()
 
     let savedContextCommands: MynahUIDataModel['contextCommands'] = []
-    const serviceName = service ?? ''
 
     window.addEventListener('error', (e) => {
         const { error, message } = e
@@ -626,7 +625,7 @@ export const createMynahUI = (
                 tabsStorage.updateTabTypeFromUnknown(tabID, 'cwc')
                 mynahUI?.updateTabDefaults({
                     store: {
-                        ...tabDataGenerator.getTabData('cwc', true, undefined, serviceName),
+                        ...tabDataGenerator.getTabData('cwc', true, undefined, isSMUS),
                         tabHeaderDetails: void 0,
                         compactMode: false,
                         tabBackground: false,
@@ -933,15 +932,15 @@ export const createMynahUI = (
                 store: {
                     ...(showWelcomePage()
                         ? welcomeScreenTabData(tabDataGenerator).store
-                        : tabDataGenerator.getTabData('cwc', true, undefined, serviceName)),
+                        : tabDataGenerator.getTabData('cwc', true, undefined, isSMUS)),
                     ...(disclaimerCardActive ? { promptInputStickyCard: disclaimerCard } : {}),
                 },
             },
         },
         defaults: {
             store: showWelcomePage()
                 ? welcomeScreenTabData(tabDataGenerator).store
-                : tabDataGenerator.getTabData('cwc', true, undefined, serviceName),
+                : tabDataGenerator.getTabData('cwc', true, undefined, isSMUS),
         },
         config: {
             maxTabs: 10,

packages/core/src/amazonq/webview/ui/tabs/constants.ts

@@ -15,7 +15,6 @@ export const qChatIntroMessageForSMUS = `Hi, I'm Amazon Q. I can answer your sof
   You can enter \`/\` to see a list of quick actions. Use \`@\` to add saved prompts, files, folders, or your entire workspace as context.
   You are now using Q free tier.\n\
   `
-export const sageMakerUnifiedStudio = 'SageMakerUnifiedStudio'
 
 export type TabTypeData = {
     title: string

packages/core/src/amazonq/webview/ui/tabs/generator.ts

@@ -7,7 +7,7 @@ import { ChatItemType, MynahUIDataModel, QuickActionCommandGroup } from '@aws/my
 import { TabType } from '../storages/tabsStorage'
 import { FollowUpGenerator } from '../followUps/generator'
 import { QuickActionGenerator } from '../quickActions/generator'
-import { qChatIntroMessageForSMUS, sageMakerUnifiedStudio, TabTypeDataMap } from './constants'
+import { qChatIntroMessageForSMUS, TabTypeDataMap } from './constants'
 import { agentWalkthroughDataModel } from '../walkthrough/agent'
 import { FeatureContext } from '../../../../shared/featureConfig'
 
@@ -43,7 +43,7 @@ export class TabDataGenerator {
         tabType: TabType,
         needWelcomeMessages: boolean,
         taskName?: string,
-        serviceName?: string
+        isSMUS?: boolean
     ): MynahUIDataModel {
         if (tabType === 'agentWalkthrough') {
             return agentWalkthroughDataModel
@@ -53,8 +53,6 @@ export class TabDataGenerator {
             return {}
         }
 
-        const isSMUS = serviceName ? serviceName === sageMakerUnifiedStudio : false
-
         const tabData: MynahUIDataModel = {
             tabTitle: taskName ?? TabTypeDataMap[tabType].title,
             promptInputInfo:

packages/core/src/auth/auth.ts

@@ -1108,12 +1108,10 @@ export class Auth implements AuthService, ConnectionManager {
  *
  * @param isC9 boolean for if Cloud9 is host
  * @param isSM boolean for if SageMaker is host
- * @param isSMUS boolean for if SageMaker Unified Studio is host
  * @returns boolean for if C9 "OR" SM
  */
-export function hasVendedIamCredentials(isC9?: boolean, isSM?: boolean, isSMUS?: boolean) {
+export function hasVendedIamCredentials(isC9?: boolean, isSM?: boolean) {
     isC9 ??= isCloud9()
     isSM ??= isSageMaker()
-    isSMUS ??= isSageMaker(true)
-    return isSM || isC9 || isSMUS
+    return isSM || isC9
 }

packages/core/src/codewhisperer/ui/statusBarMenu.ts

@@ -30,6 +30,7 @@ import { Commands } from '../../shared/vscode/commands2'
 import { createExitButton } from '../../shared/ui/buttons'
 import { telemetry } from '../../shared/telemetry/telemetry'
 import { getLogger } from '../../shared/logger/logger'
+import { isSageMaker } from '../../shared/extensionUtilities'
 
 function getAmazonQCodeWhispererNodes() {
     const autoTriggerEnabled = CodeSuggestionsState.instance.isSuggestionsEnabled()
@@ -92,7 +93,9 @@ export function getQuickPickItems(): DataQuickPickItem<string>[] {
         // Add settings and signout
         createSeparator(),
         createSettingsNode(),
-        ...(AuthUtil.instance.isConnected() && !hasVendedIamCredentials() ? [createSignout()] : []),
+        ...(AuthUtil.instance.isConnected() && !(hasVendedIamCredentials() || isSageMaker('SMUS'))
+            ? [createSignout()]
+            : []),
     ]
 
     return children

packages/core/src/extension.ts

@@ -17,7 +17,7 @@ import globals, { initialize, isWeb } from './shared/extensionGlobals'
 import { join } from 'path'
 import { Commands } from './shared/vscode/commands2'
 import { endpointsFileUrl, githubCreateIssueUrl, githubUrl } from './shared/constants'
-import { getIdeProperties, aboutExtension, getDocUrl } from './shared/extensionUtilities'
+import { getIdeProperties, aboutExtension, getDocUrl, isSageMaker } from './shared/extensionUtilities'
 import { logAndShowError, logAndShowWebviewError } from './shared/utilities/logAndShowUtils'
 import { telemetry } from './shared/telemetry/telemetry'
 import { openUrl } from './shared/utilities/vsCodeUtils'
@@ -54,6 +54,7 @@ import { AWSClientBuilderV3 } from './shared/awsClientBuilderV3'
 import { setupUninstallHandler } from './shared/handleUninstall'
 import { maybeShowMinVscodeWarning } from './shared/extensionStartup'
 import { getLogger } from './shared/logger/logger'
+import { setContext } from './shared/vscode/setContext'
 
 disableAwsSdkWarning()
 
@@ -118,6 +119,9 @@ export async function activateCommon(
     // telemetry
     await activateTelemetry(context, globals.awsContext, Settings.instance, 'AWS Toolkit For VS Code')
 
+    // set context var to identify if its SageMaker Unified Studio or not
+    await setContext('aws.isSageMakerUnifiedStudio', isSageMaker('SMUS'))
+
     // Create this now, but don't call vscode.window.registerUriHandler() until after all
     // Toolkit services have a chance to register their path handlers. #4105
     globals.uriHandler = new UriHandler()

packages/core/src/shared/extensionUtilities.ts

@@ -22,15 +22,17 @@ import {
     samDeployDocUrl,
     samInitDocUrl,
 } from './constants'
+import {
+    cloud9Appname,
+    cloud9CnAppname,
+    sageMakerAppname,
+    sageMakerUnifiedStudio,
+    vscodeAppname,
+} from './vscode/constants'
 
 const localize = nls.loadMessageBundle()
 
-const vscodeAppname = 'Visual Studio Code'
-const cloud9Appname = 'AWS Cloud9'
-const cloud9CnAppname = 'Amazon Cloud9'
-const sageMakerAppname = 'SageMaker Code Editor'
 const notInitialized = 'notInitialized'
-const sageMakerUnifiedStudio = 'SageMakerUnifiedStudio'
 
 function _isAmazonQ() {
     const id = globals.context.extension.id
@@ -66,6 +68,8 @@ export function commandsPrefix(): string {
 }
 
 let computeRegion: string | undefined = notInitialized
+let serviceName: string = notInitialized
+let isSMUS: boolean = false
 
 export function getIdeType(): 'vscode' | 'cloud9' | 'sagemaker' | 'unknown' {
     if (vscode.env.appName === cloud9Appname || vscode.env.appName === cloud9CnAppname) {
@@ -147,11 +151,11 @@ function createCloud9Properties(company: string): IdeProperties {
 }
 
 function isSageMakerUnifiedStudio(): boolean {
-    const service = process.env.SERVICE_NAME
-    if (service) {
-        return service === sageMakerUnifiedStudio
+    if (serviceName === notInitialized) {
+        serviceName = process.env.SERVICE_NAME ?? ''
+        isSMUS = serviceName === sageMakerUnifiedStudio
     }
-    return false
+    return isSMUS
 }
 
 /**
@@ -166,12 +170,20 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo
     return (flavor === 'classic' && !codecat) || (flavor === 'codecatalyst' && codecat)
 }
 
-export function isSageMaker(isUnifiedStudio?: boolean): boolean {
-    const isSM = vscode.env.appName === sageMakerAppname
-    if (isUnifiedStudio) {
-        return isSM && isSageMakerUnifiedStudio()
+/**
+ *
+ * @param appName to identify the proper SM instance
+ * @returns true if the current system is SageMaker(SMAI or SMUS)
+ */
+export function isSageMaker(appName: string = 'SMAI'): boolean {
+    switch (appName) {
+        case 'SMAI':
+            return vscode.env.appName === sageMakerAppname
+        case 'SMUS':
+            return vscode.env.appName === sageMakerAppname && isSageMakerUnifiedStudio()
+        default:
+            return false
     }
-    return isSM
 }
 
 export function isCn(): boolean {