feat(sagemaker): disable Q signout on SMUS #6850

parameja1 · web-flow · commit 9f805be4c87b · 2025-04-01T13:26:36.000-07:00
## Problem In Sagemaker Unified Studio(SMUS), user shouldnt be allowed to sign out of Q Chat extension as per business requirement. In addition to that, Q Chat Introduction message should display Q Free Tier is only supported with the upcoming SMUS release ## Solution 1. `isSageMaker` is used to identify if it is SageMaker AI (CodeEditor) which comes from [here](https://github.com/aws/sagemaker-code-editor/blob/main/patched-vscode/product.json#L2-L3) as shared by Sagemaker team. Since SageMaker Unified Studio also uses the same repository/branch for CodeEditor we required a new environment variable `SERVICE_NAME` to differentiate SMUS case. 2. Read `SERVICE_NAME` and set a new vscode env context variable `aws.isSagemakerStudio` to control visibility of Sign out label 3. Sign out option to users is controlled at two places: - VS Code command palette - Quick pick UI action - Both the sign out options are disabled based on the SMUS flag 4. Similarly Q Chat Introduction message in WebView is modified for SMUS case alone
diff --git a/packages/amazonq/.changes/next-release/Feature-390d2096-ea1e-43b9-a889-2494d6e9e6d3.json b/packages/amazonq/.changes/next-release/Feature-390d2096-ea1e-43b9-a889-2494d6e9e6d3.json
@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "SageMaker Unified Studio: Disable Sign out"
+}
diff --git a/packages/amazonq/.changes/next-release/Feature-3c42f632-46b8-4e59-b5fa-75d638f3f08d.json b/packages/amazonq/.changes/next-release/Feature-3c42f632-46b8-4e59-b5fa-75d638f3f08d.json
@@ -0,0 +1,4 @@
+{
+	"type": "Feature",
+	"description": "SageMaker Unified Studio: Update Q Chat Introduction message"
+}
diff --git a/packages/amazonq/package.json b/packages/amazonq/package.json
@@ -359,7 +359,7 @@
                 },
                 {
                     "command": "aws.amazonq.signout",
-                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected",
+                    "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio",
                     "group": "2_amazonQ@4"
                 },
                 {
@@ -635,7 +635,7 @@
                 "title": "%AWS.command.codewhisperer.signout%",
                 "category": "%AWS.amazonq.title%",
                 "icon": "$(debug-disconnect)",
-                "enablement": "aws.codewhisperer.connected"
+                "enablement": "aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio"
             },
             {
                 "command": "aws.amazonq.learnMore",
diff --git a/packages/amazonq/src/extension.ts b/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')
 
diff --git a/packages/core/src/amazonq/webview/generators/webViewContent.ts b/packages/core/src/amazonq/webview/generators/webViewContent.ts
@@ -83,6 +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 isSMUS = isSageMaker('SMUS')
 
         return `
         <script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>
@@ -91,7 +92,7 @@ export class WebViewContentGenerator {
             const init = () => {
                 createMynahUI(acquireVsCodeApi(), ${
                     (await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'
-                },${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString});
+                },${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString},${isSMUS});
             }
         </script>
         `
diff --git a/packages/core/src/amazonq/webview/ui/main.ts b/packages/core/src/amazonq/webview/ui/main.ts
@@ -47,7 +47,8 @@ export const createMynahUI = (
     featureConfigsSerialized: [string, FeatureContext][],
     welcomeCount: number,
     disclaimerAcknowledged: boolean,
-    disabledCommands?: string[]
+    disabledCommands?: string[],
+    isSMUS?: boolean
 ) => {
     let disclaimerCardActive = !disclaimerAcknowledged
     // eslint-disable-next-line prefer-const
@@ -624,7 +625,7 @@ export const createMynahUI = (
                 tabsStorage.updateTabTypeFromUnknown(tabID, 'cwc')
                 mynahUI?.updateTabDefaults({
                     store: {
-                        ...tabDataGenerator.getTabData('cwc', true),
+                        ...tabDataGenerator.getTabData('cwc', true, undefined, isSMUS),
                         tabHeaderDetails: void 0,
                         compactMode: false,
                         tabBackground: false,
@@ -931,15 +932,15 @@ export const createMynahUI = (
                 store: {
                     ...(showWelcomePage()
                         ? welcomeScreenTabData(tabDataGenerator).store
-                        : tabDataGenerator.getTabData('cwc', true)),
+                        : tabDataGenerator.getTabData('cwc', true, undefined, isSMUS)),
                     ...(disclaimerCardActive ? { promptInputStickyCard: disclaimerCard } : {}),
                 },
             },
         },
         defaults: {
             store: showWelcomePage()
                 ? welcomeScreenTabData(tabDataGenerator).store
-                : tabDataGenerator.getTabData('cwc', true),
+                : tabDataGenerator.getTabData('cwc', true, undefined, isSMUS),
         },
         config: {
             maxTabs: 10,
diff --git a/packages/core/src/amazonq/webview/ui/tabs/constants.ts b/packages/core/src/amazonq/webview/ui/tabs/constants.ts
@@ -6,6 +6,16 @@ import { TabType } from '../storages/tabsStorage'
 import { QuickActionCommandGroup } from '@aws/mynah-ui'
 import { userGuideURL } from '../texts/constants'
 
+const qChatIntroMessage = `Hi, I'm Amazon Q. I can answer your software development questions.
+  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 type TabTypeData = {
     title: string
     placeholder: string
@@ -26,9 +36,7 @@ export const workspaceCommand: QuickActionCommandGroup = {
 export const commonTabData: TabTypeData = {
     title: 'Chat',
     placeholder: 'Ask a question. Use @ to add context, / for quick actions',
-    welcome: `Hi, I'm Amazon Q. I can answer your software development questions.
-  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.`,
+    welcome: qChatIntroMessage,
     contextCommands: [workspaceCommand],
 }
 
diff --git a/packages/core/src/amazonq/webview/ui/tabs/generator.ts b/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 { TabTypeDataMap } from './constants'
+import { qChatIntroMessageForSMUS, TabTypeDataMap } from './constants'
 import { agentWalkthroughDataModel } from '../walkthrough/agent'
 import { FeatureContext } from '../../../../shared/featureConfig'
 
@@ -39,7 +39,12 @@ export class TabDataGenerator {
         this.highlightCommand = props.commandHighlight
     }
 
-    public getTabData(tabType: TabType, needWelcomeMessages: boolean, taskName?: string): MynahUIDataModel {
+    public getTabData(
+        tabType: TabType,
+        needWelcomeMessages: boolean,
+        taskName?: string,
+        isSMUS?: boolean
+    ): MynahUIDataModel {
         if (tabType === 'agentWalkthrough') {
             return agentWalkthroughDataModel
         }
@@ -59,7 +64,7 @@ export class TabDataGenerator {
                 ? [
                       {
                           type: ChatItemType.ANSWER,
-                          body: TabTypeDataMap[tabType].welcome,
+                          body: isSMUS ? qChatIntroMessageForSMUS : TabTypeDataMap[tabType].welcome,
                       },
                       {
                           type: ChatItemType.ANSWER,
diff --git a/packages/core/src/codewhisperer/ui/statusBarMenu.ts b/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
diff --git a/packages/core/src/extension.ts b/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()
diff --git a/packages/core/src/shared/extensionUtilities.ts b/packages/core/src/shared/extensionUtilities.ts
@@ -22,13 +22,16 @@ 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'
 
 function _isAmazonQ() {
@@ -65,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) {
@@ -145,6 +150,14 @@ function createCloud9Properties(company: string): IdeProperties {
     }
 }
 
+function isSageMakerUnifiedStudio(): boolean {
+    if (serviceName === notInitialized) {
+        serviceName = process.env.SERVICE_NAME ?? ''
+        isSMUS = serviceName === sageMakerUnifiedStudio
+    }
+    return isSMUS
+}
+
 /**
  * Decides if the current system is (the specified flavor of) Cloud9.
  */
@@ -157,8 +170,20 @@ 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
+/**
+ *
+ * @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
+    }
 }
 
 export function isCn(): boolean {
diff --git a/packages/core/src/shared/vscode/constants.ts b/packages/core/src/shared/vscode/constants.ts
@@ -0,0 +1,9 @@
+/*!
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+export const vscodeAppname = 'Visual Studio Code'
+export const cloud9Appname = 'AWS Cloud9'
+export const cloud9CnAppname = 'Amazon Cloud9'
+export const sageMakerAppname = 'SageMaker Code Editor'
+export const sageMakerUnifiedStudio = 'SageMakerUnifiedStudio'
diff --git a/packages/core/src/shared/vscode/setContext.ts b/packages/core/src/shared/vscode/setContext.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'
diff --git a/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts b/packages/core/src/test/codewhisperer/commands/basicCommands.test.ts
@@ -68,6 +68,7 @@ import { CodeWhispererSettings } from '../../../codewhisperer/util/codewhisperer
 import { confirm } from '../../../shared'
 import * as commentUtils from '../../../shared/utilities/commentUtils'
 import * as startCodeFixGeneration from '../../../codewhisperer/commands/startCodeFixGeneration'
+import * as extUtils from '../../../shared/extensionUtilities'
 
 describe('CodeWhisperer-basicCommands', function () {
     let targetCommand: Command<any> & vscode.Disposable
@@ -520,6 +521,42 @@ describe('CodeWhisperer-basicCommands', function () {
             })
             await listCodeWhispererCommands.execute()
         })
+
+        it('includes sign out when connected and not in SageMaker', async function () {
+            sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
+            sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(false)
+            sinon.stub(extUtils, 'isSageMaker').value(false)
+            await CodeScansState.instance.setScansEnabled(false)
+
+            getTestWindow().onDidShowQuickPick((e) => {
+                e.assertContainsItems(
+                    createAutoSuggestions(true),
+                    createOpenReferenceLog(),
+                    createGettingStarted(),
+                    createAutoScans(false),
+                    switchToAmazonQNode(),
+                    ...genericItems(),
+                    createSettingsNode(),
+                    createSignout()
+                )
+                e.dispose()
+            })
+
+            await listCodeWhispererCommands.execute()
+        })
+
+        it('shows expected items when connection is expired and in SageMaker', async function () {
+            sinon.stub(AuthUtil.instance, 'isConnected').returns(true)
+            sinon.stub(AuthUtil.instance, 'isConnectionExpired').returns(true)
+            sinon.stub(extUtils, 'isSageMaker').value(true)
+
+            getTestWindow().onDidShowQuickPick((e) => {
+                e.assertContainsItems(createReconnect(), createLearnMore(), ...genericItems())
+                e.dispose()
+            })
+
+            await listCodeWhispererCommands.execute()
+        })
     })
 
     describe('applySecurityFix', function () {

-Original file line number
+Diff line change
@@ @@ -0,0 +1,4 @@ @@
 +{
 +	"type": "Feature",
 +	"description": "SageMaker Unified Studio: Disable Sign out"
 +}
Original file line number	Diff line number	Diff line change
`@@ -359,7 +359,7 @@`
`359`	`359`	`},`
`360`	`360`	`{`
`361`	`361`	`"command": "aws.amazonq.signout",`
`362`		`- "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected",`
	`362`	`+ "when": "(view == aws.amazonq.AmazonQChatView) && aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio",`
`363`	`363`	`"group": "2_amazonQ@4"`
`364`	`364`	`},`
`365`	`365`	`{`
`@@ -635,7 +635,7 @@`
`635`	`635`	`"title": "%AWS.command.codewhisperer.signout%",`
`636`	`636`	`"category": "%AWS.amazonq.title%",`
`637`	`637`	`"icon": "$(debug-disconnect)",`
`638`		`- "enablement": "aws.codewhisperer.connected"`
	`638`	`+ "enablement": "aws.codewhisperer.connected && !aws.isSageMakerUnifiedStudio"`
`639`	`639`	`},`
`640`	`640`	`{`
`641`	`641`	`"command": "aws.amazonq.learnMore",`
Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,7 @@ export class WebViewContentGenerator {`
`83`	`83`	`const disclaimerAcknowledged = globals.globalState.tryGet('aws.amazonq.disclaimerAcknowledged', Boolean, false)`
`84`	`84`
`85`	`85`	`const welcomeLoadCount = globals.globalState.tryGet('aws.amazonq.welcomeChatShowCount', Number, 0)`
	`86`	`+ const isSMUS = isSageMaker('SMUS')`
`86`	`87`
`87`	`88`	return `
`88`	`89`	`<script type="text/javascript" src="${javascriptEntrypoint.toString()}" defer onload="init()"></script>`
`@@ -91,7 +92,7 @@ export class WebViewContentGenerator {`
`91`	`92`	`const init = () => {`
`92`	`93`	`createMynahUI(acquireVsCodeApi(), ${`
`93`	`94`	`(await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'`
`94`		`- },${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString});`
	`95`	`+ },${featureConfigsString},${welcomeLoadCount},${disclaimerAcknowledged},${disabledCommandsString},${isSMUS});`
`95`	`96`	`}`
`96`	`97`	`</script>`
`97`	`98`	`