Refactor to use IsSageMaker to differentiate SMUS, remove global state

change and modified unit tests

Author: parameja1

packages/amazonq/package.json

@@ -351,7 +351,7 @@
                 },
                 {
                     "command": "aws.amazonq.signout",
-                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.amazonq.isSagemakerUnifiedStudio",
+                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio",
                     "group": "2_amazonQ@4"
                 },
                 {
@@ -627,7 +627,7 @@
                 "title": "%AWS.command.codewhisperer.signout%",
                 "category": "%AWS.amazonq.title%",
                 "icon": "$(debug-disconnect)",
-                "enablement": "aws.codewhisperer.connected && !aws.amazonq.isSagemakerUnifiedStudio"
+                "enablement": "aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio"
             },
             {
                 "command": "aws.amazonq.learnMore",

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, getLogger, setContext } from 'aws-core-vscode/shared'
+import { Commands, placeholder, funcUtil, setContext, isSageMaker } 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,7 @@ export async function activate(context: ExtensionContext) {
     const appInitContext = amazonq.DefaultAmazonQAppInitContext.instance
 
     registerApps(appInitContext, context)
-    await setupAmazonQState()
+    await setIsSMUSContext()
 
     const provider = new amazonq.AmazonQChatViewProvider(
         context,
@@ -65,12 +65,10 @@ export async function activate(context: ExtensionContext) {
     void setupLsp()
     void setupAuthNotification()
 
-    async function setupAmazonQState() {
-        const serviceName = process.env.SERVICE_NAME ?? 'AmazonQ' // Fallback service name to a generic name AmazonQ
-        const amazonQState = amazonq.AmazonQState.initialize(serviceName)
-        getLogger().info(`In Extension activation - q state is initialized ${amazonQState.serviceName}`)
-        // Create a context key for SageMaker Studio state and set in context
-        await setContext('aws.amazonq.isSagemakerUnifiedStudio', amazonQState.isSageMakerUnifiedStudio())
+    async function setIsSMUSContext() {
+        // set context key for SageMaker Unified Studio
+        const isSMUS = isSageMaker(true)
+        await setContext('aws.isSageMakerUnifiedStudio', isSMUS)
     }
 }
 

packages/core/src/amazonq/index.ts

@@ -43,7 +43,6 @@ export { ChatItemType } from './commons/model'
 export { ExtensionMessage } from '../amazonq/webview/ui/commands'
 export { CodeReference } from '../codewhispererChat/view/connector/connector'
 export { extractAuthFollowUp } from './util/authUtils'
-export { AmazonQState } from './util/amazonQState'
 export { Messenger } from './commons/connector/baseMessenger'
 export * from './lsp/config'
 export * as WorkspaceLspInstaller from './lsp/workspaceInstaller'

packages/core/src/amazonq/util/amazonQState.ts

@@ -10,19 +10,20 @@ const qChatIntroMessage = `Hi, I'm Amazon Q. I can answer your software developm
   Ask me to explain, debug, or optimize your code.
   You can enter \`/\` to see a list of quick actions. Use \`@\` to add saved prompts, files, folders, or your entire workspace as context.`
 
+export const qChatIntroMessageForSMUS = `Hi, I'm Amazon Q. I can answer your software development questions.\n\
+  Ask me to explain, debug, or optimize your code.\n\
+  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
     placeholder: string
     welcome: string
     contextCommands?: QuickActionCommandGroup[]
 }
 
-export const qChatIntroMessageForSMUS = `Hi, I'm Amazon Q. I can answer your software development questions.\n\
-Ask me to explain, debug, or optimize your code.\n\
-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 workspaceCommand: QuickActionCommandGroup = {
     groupName: 'Mention code',
     commands: [

packages/core/src/amazonq/webview/ui/tabs/constants.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, TabTypeDataMap } from './constants'
+import { qChatIntroMessageForSMUS, sageMakerUnifiedStudio, TabTypeDataMap } from './constants'
 import { agentWalkthroughDataModel } from '../walkthrough/agent'
 import { FeatureContext } from '../../../../shared/featureConfig'
 
@@ -53,7 +53,7 @@ export class TabDataGenerator {
             return {}
         }
 
-        const isSMUS = serviceName ? serviceName === 'SageMakerUnifiedStudio' : false
+        const isSMUS = serviceName ? serviceName === sageMakerUnifiedStudio : false
 
         const tabData: MynahUIDataModel = {
             tabTitle: taskName ?? TabTypeDataMap[tabType].title,

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

@@ -1108,10 +1108,12 @@ 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) {
+export function hasVendedIamCredentials(isC9?: boolean, isSM?: boolean, isSMUS?: boolean) {
     isC9 ??= isCloud9()
     isSM ??= isSageMaker()
-    return isSM || isC9
+    isSMUS ??= isSageMaker(true)
+    return isSM || isC9 || isSMUS
 }

packages/core/src/auth/auth.ts

@@ -30,7 +30,6 @@ 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 { AmazonQState } from '../../amazonq/util/amazonQState'
 
 function getAmazonQCodeWhispererNodes() {
     const autoTriggerEnabled = CodeSuggestionsState.instance.isSuggestionsEnabled()
@@ -93,11 +92,7 @@ export function getQuickPickItems(): DataQuickPickItem<string>[] {
         // Add settings and signout
         createSeparator(),
         createSettingsNode(),
-        ...(AuthUtil.instance.isConnected() &&
-        !hasVendedIamCredentials() &&
-        !AmazonQState.instance.isSageMakerUnifiedStudio()
-            ? [createSignout()]
-            : []),
+        ...(AuthUtil.instance.isConnected() && !hasVendedIamCredentials() ? [createSignout()] : []),
     ]
 
     return children

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

@@ -30,6 +30,7 @@ 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
@@ -145,6 +146,14 @@ function createCloud9Properties(company: string): IdeProperties {
     }
 }
 
+function isSageMakerUnifiedStudio(): boolean {
+    const service = process.env.SERVICE_NAME
+    if (service) {
+        return service === sageMakerUnifiedStudio
+    }
+    return false
+}
+
 /**
  * Decides if the current system is (the specified flavor of) Cloud9.
  */
@@ -157,8 +166,12 @@ export function isCloud9(flavor: 'classic' | 'codecatalyst' | 'any' = 'any'): bo
     return (flavor === 'classic' && !codecat) || (flavor === 'codecatalyst' && codecat)
 }
 
-export function isSageMaker(): boolean {
-    return vscode.env.appName === sageMakerAppname
+export function isSageMaker(isUnifiedStudio?: boolean): boolean {
+    const isSM = vscode.env.appName === sageMakerAppname
+    if (isUnifiedStudio) {
+        return isSM && isSageMakerUnifiedStudio()
+    }
+    return isSM
 }
 
 export function isCn(): boolean {

packages/core/src/shared/extensionUtilities.ts

@@ -13,6 +13,7 @@ import * as vscode from 'vscode'
 export type contextKey =
     | 'aws.isDevMode'
     | 'aws.isSageMaker'
+    | 'aws.isSageMakerUnifiedStudio'
     | 'aws.isWebExtHost'
     | 'aws.isInternalUser'
     | 'aws.amazonq.showLoginView'
@@ -37,7 +38,6 @@ export type contextKey =
     | 'gumby.wasQCodeTransformationUsed'
     | 'amazonq.inline.codelensShortcutEnabled'
     | 'aws.toolkit.lambda.walkthroughSelected'
-    | 'aws.amazonq.isSagemakerUnifiedStudio'
 
 const contextMap: Partial<Record<contextKey, any>> = {}