feat(auth): add IAM credentials LSP requests to AuthUtils and auth2

aws-toolkit-vscode.code-workspace

@@ -12,6 +12,15 @@
         {
             "path": "packages/amazonq",
         },
+        {
+            "path": "../language-servers",
+        },
+        {
+            "path": "../language-server-runtimes",
+        },
+        {
+            "path": "../aws-toolkit-common",
+        },
     ],
     "settings": {
         "typescript.tsdk": "node_modules/typescript/lib",

packages/amazonq/.vscode/launch.json

@@ -13,10 +13,10 @@
             "args": ["--extensionDevelopmentPath=${workspaceFolder}"],
             "env": {
                 "SSMDOCUMENT_LANGUAGESERVER_PORT": "6010",
-                "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080"
+                "WEBPACK_DEVELOPER_SERVER": "http://localhost:8080",
                 // Below allows for overrides used during development
-                // "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js",
-                // "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js"
+                "__AMAZONQLSP_PATH": "${workspaceFolder}/../../../language-servers/app/aws-lsp-codewhisperer-runtimes/out/agent-standalone.js",
+                "__AMAZONQLSP_UI": "${workspaceFolder}/../../../language-servers/chat-client/build/amazonq-ui.js"
             },
             "envFile": "${workspaceFolder}/.local.env",
             "outFiles": ["${workspaceFolder}/dist/**/*.js", "${workspaceFolder}/../core/dist/**/*.js"],

packages/amazonq/src/extensionNode.ts

@@ -100,8 +100,11 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
 async function getAuthState(): Promise<Omit<AuthUserState, 'source'>> {
     const state = AuthUtil.instance.getAuthState()
 
-    if (AuthUtil.instance.isConnected() && !(AuthUtil.instance.isSsoSession() || isSageMaker())) {
-        getLogger().error('Current Amazon Q connection is not SSO')
+    if (
+        AuthUtil.instance.isConnected() &&
+        !(AuthUtil.instance.isSsoSession() || AuthUtil.instance.isIamSession() || isSageMaker())
+    ) {
+        getLogger().error('Current Amazon Q connection is not SSO nor IAM')
     }
 
     return {

packages/amazonq/src/inlineChat/provider/inlineChatProvider.ts

@@ -16,6 +16,7 @@ import {
     ChatTriggerType,
     EditorContextExtractor,
     PromptMessage,
+    TriggerEvent,
     TriggerEventsStorage,
     TriggerPayload,
     triggerPayloadToChatRequest,
@@ -30,6 +31,7 @@ import { extractAuthFollowUp } from 'aws-core-vscode/amazonq'
 import { InlineChatParams, InlineChatResult } from '@aws/language-server-runtimes-types'
 import { decryptResponse, encryptRequest } from '../../lsp/encryption'
 import { getCursorState } from '../../lsp/utils'
+import { CwsprChatTriggerInteraction, telemetry } from 'aws-core-vscode/telemetry'
 
 export class InlineChatProvider {
     private readonly editorContextExtractor: EditorContextExtractor
@@ -68,7 +70,34 @@ export class InlineChatProvider {
         }
     }
 
+    private getTriggerInteractionFromTriggerEvent(triggerEvent: TriggerEvent | undefined): CwsprChatTriggerInteraction {
+        switch (triggerEvent?.type) {
+            case 'editor_context_command':
+                return triggerEvent.command?.triggerType === 'keybinding' ? 'hotkeys' : 'contextMenu'
+            case 'follow_up':
+            case 'chat_message':
+            default:
+                return 'click'
+        }
+    }
+
     public async processPromptMessageLSP(message: PromptMessage): Promise<InlineChatResult> {
+        const triggerInteraction = this.getTriggerInteractionFromTriggerEvent(
+            this.triggerEventsStorage.getLastTriggerEventByTabID(message.tabID)
+        )
+        if (!AuthUtil.instance.isSsoSession()) {
+            telemetry.amazonq_messageResponseError.emit({
+                result: 'Failed',
+                cwsprChatConversationType: 'Chat',
+                cwsprChatRequestLength: message.message?.length ?? 0,
+                cwsprChatResponseCode: 401,
+                cwsprChatTriggerInteraction: triggerInteraction,
+                reason: 'AuthenticationError',
+                reasonDesc: 'Inline chat requires SSO authentication, but current session is not',
+            })
+            throw new ToolkitError('Inline chat is only available with SSO authentication')
+        }
+
         // TODO: handle partial responses.
         getLogger().info('Making inline chat request with message %O', message)
         const params = this.getCurrentEditorParams(message.message ?? '')
@@ -83,6 +112,23 @@ export class InlineChatProvider {
 
     // TODO: remove in favor of LSP implementation.
     public async processPromptMessage(message: PromptMessage) {
+        const triggerInteraction = this.getTriggerInteractionFromTriggerEvent(
+            this.triggerEventsStorage.getLastTriggerEventByTabID(message.tabID)
+        )
+        if (!AuthUtil.instance.isSsoSession()) {
+            telemetry.amazonq_messageResponseError.emit({
+                result: 'Failed',
+                cwsprChatConversationType: 'Chat',
+                cwsprChatRequestLength: message.message?.length ?? 0,
+                cwsprChatResponseCode: 401,
+                cwsprChatTriggerInteraction: triggerInteraction,
+                reason: 'AuthenticationError',
+                reasonDesc: 'Inline chat requires SSO authentication, but current session is not',
+                credentialStartUrl: AuthUtil.instance.connection?.startUrl,
+            })
+            throw new ToolkitError('Inline chat is only available with SSO authentication')
+        }
+
         return this.editorContextExtractor
             .extractContextForTrigger('ChatMessage')
             .then((context) => {
@@ -143,7 +189,7 @@ export class InlineChatProvider {
     private async generateResponse(
         triggerPayload: TriggerPayload & { projectContextQueryLatencyMs?: number },
         triggerID: string
-    ) {
+    ): Promise<GenerateAssistantResponseCommandOutput | undefined> {
         const triggerEvent = this.triggerEventsStorage.getTriggerEvent(triggerID)
         if (triggerEvent === undefined) {
             return
@@ -182,7 +228,12 @@ export class InlineChatProvider {
         let response: GenerateAssistantResponseCommandOutput | undefined = undefined
         session.createNewTokenSource()
         try {
-            response = await session.chatSso(request)
+            if (AuthUtil.instance.isSsoSession()) {
+                response = await session.chatSso(request)
+            } else {
+                // Call sendMessage because Q Developer Streaming Client does not have generateAssistantResponse
+                throw new ToolkitError('Inline chat is only available with SSO authentication')
+            }
             getLogger().info(
                 `response to tab: ${tabID} conversationID: ${session.sessionIdentifier} requestID: ${response.$metadata.requestId} metadata: %O`,
                 response.$metadata

packages/amazonq/src/lsp/client.ts

@@ -30,6 +30,8 @@ import {
     ResponseError,
     LSPErrorCodes,
     updateConfigurationRequestType,
+    GetMfaCodeParams,
+    GetMfaCodeResult,
 } from '@aws/language-server-runtimes/protocol'
 import {
     AuthUtil,
@@ -56,7 +58,7 @@ import { processUtils } from 'aws-core-vscode/shared'
 import { activate as activateChat } from './chat/activation'
 import { activate as activeInlineChat } from '../inlineChat/activation'
 import { AmazonQResourcePaths } from './lspInstaller'
-import { auth2 } from 'aws-core-vscode/auth'
+import { auth2, getMfaTokenFromUser, getMfaSerialFromUser } from 'aws-core-vscode/auth'
 import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
 import { telemetry } from 'aws-core-vscode/telemetry'
 import { SessionManager } from '../app/inline/sessionManager'
@@ -164,6 +166,9 @@ export async function startLanguageServer(
             },
             credentials: {
                 providesBearerToken: true,
+                // Add IAM credentials support
+                providesIamCredentials: true,
+                supportsAssumeRole: true,
             },
         },
         /**
@@ -211,9 +216,10 @@ export async function startLanguageServer(
 
         /** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */
         AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
+            const activeProfile = AuthUtil.instance.regionProfileManager.activeRegionProfile
             void pushConfigUpdate(client, {
                 type: 'profile',
-                profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
+                profileArn: activeProfile?.arn,
             })
         })
 
@@ -333,6 +339,24 @@ async function postStartLanguageServer(
         }
     )
 
+    // Handler for when Flare needs to assume a role with MFA code
+    client.onRequest(
+        auth2.notificationTypes.getMfaCode.method,
+        async (params: GetMfaCodeParams): Promise<GetMfaCodeResult> => {
+            if (params.mfaSerial) {
+                globals.globalState.update('recentMfaSerial', { mfaSerial: params.mfaSerial })
+            }
+            const defaultMfaSerial = globals.globalState.tryGet('recentMfaSerial', Object, {
+                mfaSerial: '',
+            }).mfaSerial
+            let mfaSerial = await getMfaSerialFromUser(defaultMfaSerial, params.profileName)
+            mfaSerial = mfaSerial.trim()
+            globals.globalState.update('recentMfaSerial', { mfaSerial: mfaSerial })
+            const mfaCode = await getMfaTokenFromUser(mfaSerial, params.profileName)
+            return { code: mfaCode ?? '', mfaSerial: mfaSerial ?? '' }
+        }
+    )
+
     const sendProfileToLsp = async () => {
         try {
             const result = await client.sendRequest(updateConfigurationRequestType.method, {

packages/amazonq/test/e2e/amazonq/utils/setup.ts

@@ -22,5 +22,5 @@ export async function loginToIdC() {
         )
     }
 
-    await AuthUtil.instance.login(startUrl, region)
+    await AuthUtil.instance.loginSso(startUrl, region)
 }

packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts

@@ -26,11 +26,11 @@ describe('RegionProfileManager', async function () {
 
     async function setupConnection(type: 'builderId' | 'idc') {
         if (type === 'builderId') {
-            await AuthUtil.instance.login(constants.builderIdStartUrl, region)
+            await AuthUtil.instance.loginSso(constants.builderIdStartUrl, region)
             assert.ok(AuthUtil.instance.isSsoSession())
             assert.ok(AuthUtil.instance.isBuilderIdConnection())
         } else if (type === 'idc') {
-            await AuthUtil.instance.login(enterpriseSsoStartUrl, region)
+            await AuthUtil.instance.loginSso(enterpriseSsoStartUrl, region)
             assert.ok(AuthUtil.instance.isSsoSession())
             assert.ok(AuthUtil.instance.isIdcConnection())
         }