fix: Q chat stopped using IAM creds in the Amazon Q v1.63.0 release

floralph · floralph · commit d0923ff71749 · 2025-07-14T14:44:26.000-07:00
diff --git a/P261194666.md b/P261194666.md
diff --git a/packages/amazonq/src/lsp/auth.ts b/packages/amazonq/src/lsp/auth.ts
@@ -5,6 +5,7 @@
 
 import {
     bearerCredentialsUpdateRequestType,
+    iamCredentialsUpdateRequestType,
     ConnectionMetadata,
     NotificationType,
     RequestType,
@@ -17,8 +18,8 @@ import { LanguageClient } from 'vscode-languageclient'
 import { AuthUtil } from 'aws-core-vscode/codewhisperer'
 import { Writable } from 'stream'
 import { onceChanged } from 'aws-core-vscode/utils'
-import { getLogger, oneMinute } from 'aws-core-vscode/shared'
-import { isSsoConnection } from 'aws-core-vscode/auth'
+import { getLogger, oneMinute, isSageMaker } from 'aws-core-vscode/shared'
+import { isSsoConnection, isIamConnection } from 'aws-core-vscode/auth'
 
 export const encryptionKey = crypto.randomBytes(32)
 
@@ -78,10 +79,16 @@ export class AmazonQLspAuth {
      */
     async refreshConnection(force: boolean = false) {
         const activeConnection = this.authUtil.conn
-        if (this.authUtil.isConnectionValid() && isSsoConnection(activeConnection)) {
-            // send the token to the language server
-            const token = await this.authUtil.getBearerToken()
-            await (force ? this._updateBearerToken(token) : this.updateBearerToken(token))
+        if (this.authUtil.isConnectionValid()) {
+            if (isSsoConnection(activeConnection)) {
+                // Existing SSO path
+                const token = await this.authUtil.getBearerToken()
+                await (force ? this._updateBearerToken(token) : this.updateBearerToken(token))
+            } else if (isSageMaker() && isIamConnection(activeConnection)) {
+                // New SageMaker IAM path
+                const credentials = await this.authUtil.getCredentials()
+                await (force ? this._updateIamCredentials(credentials) : this.updateIamCredentials(credentials))
+            }
         }
     }
 
@@ -92,9 +99,7 @@ export class AmazonQLspAuth {
 
     public updateBearerToken = onceChanged(this._updateBearerToken.bind(this))
     private async _updateBearerToken(token: string) {
-        const request = await this.createUpdateCredentialsRequest({
-            token,
-        })
+        const request = await this.createUpdateBearerCredentialsRequest(token)
 
         // "aws/credentials/token/update"
         // https://github.com/aws/language-servers/blob/44d81f0b5754747d77bda60b40cc70950413a737/core/aws-lsp-core/src/credentials/credentialsProvider.ts#L27
@@ -103,15 +108,36 @@ export class AmazonQLspAuth {
         this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)
     }
 
+    public updateIamCredentials = onceChanged(this._updateIamCredentials.bind(this))
+    private async _updateIamCredentials(credentials: any) {
+        getLogger().info(
+            `[SageMaker Debug] Updating IAM credentials - credentials received: ${credentials ? 'YES' : 'NO'}`
+        )
+        if (credentials) {
+            getLogger().info(
+                `[SageMaker Debug] IAM credentials structure: accessKeyId=${credentials.accessKeyId ? 'present' : 'missing'}, secretAccessKey=${credentials.secretAccessKey ? 'present' : 'missing'}, sessionToken=${credentials.sessionToken ? 'present' : 'missing'}`
+            )
+        }
+
+        const request = await this.createUpdateIamCredentialsRequest(credentials)
+
+        // "aws/credentials/iam/update"
+        await this.client.sendRequest(iamCredentialsUpdateRequestType.method, request)
+
+        this.client.info(`UpdateIamCredentials: ${JSON.stringify(request)}`)
+        getLogger().info(`[SageMaker Debug] IAM credentials update request sent successfully`)
+    }
+
     public startTokenRefreshInterval(pollingTime: number = oneMinute / 2) {
         const interval = setInterval(async () => {
             await this.refreshConnection().catch((e) => this.logRefreshError(e))
         }, pollingTime)
         return interval
     }
 
-    private async createUpdateCredentialsRequest(data: any): Promise<UpdateCredentialsParams> {
-        const payload = new TextEncoder().encode(JSON.stringify({ data }))
+    private async createUpdateBearerCredentialsRequest(token: string): Promise<UpdateCredentialsParams> {
+        const bearerCredentials = { token }
+        const payload = new TextEncoder().encode(JSON.stringify({ data: bearerCredentials }))
 
         const jwt = await new jose.CompactEncrypt(payload)
             .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
@@ -127,4 +153,24 @@ export class AmazonQLspAuth {
             encrypted: true,
         }
     }
+
+    private async createUpdateIamCredentialsRequest(credentials: any): Promise<UpdateCredentialsParams> {
+        // Extract IAM credentials structure
+        const iamCredentials = {
+            accessKeyId: credentials.accessKeyId,
+            secretAccessKey: credentials.secretAccessKey,
+            sessionToken: credentials.sessionToken,
+        }
+        const payload = new TextEncoder().encode(JSON.stringify({ data: iamCredentials }))
+
+        const jwt = await new jose.CompactEncrypt(payload)
+            .setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
+            .encrypt(encryptionKey)
+
+        return {
+            data: jwt,
+            // Omit metadata for IAM credentials since startUrl is undefined for non-SSO connections
+            encrypted: true,
+        }
+    }
 }
diff --git a/packages/amazonq/src/lsp/chat/messages.ts b/packages/amazonq/src/lsp/chat/messages.ts
@@ -285,6 +285,31 @@ export function registerMessageListeners(
                 }
 
                 const chatRequest = await encryptRequest<ChatParams>(chatParams, encryptionKey)
+
+                // Add detailed logging for SageMaker debugging
+                if (process.env.USE_IAM_AUTH === 'true') {
+                    languageClient.info(`[SageMaker Debug] Making chat request with IAM auth`)
+                    languageClient.info(`[SageMaker Debug] Chat request method: ${chatRequestType.method}`)
+                    languageClient.info(
+                        `[SageMaker Debug] Original chat params: ${JSON.stringify(
+                            {
+                                tabId: chatParams.tabId,
+                                prompt: chatParams.prompt,
+                                // Don't log full textDocument content, just metadata
+                                textDocument: chatParams.textDocument
+                                    ? { uri: chatParams.textDocument.uri }
+                                    : undefined,
+                                context: chatParams.context ? `${chatParams.context.length} context items` : undefined,
+                            },
+                            null,
+                            2
+                        )}`
+                    )
+                    languageClient.info(
+                        `[SageMaker Debug] Environment context: USE_IAM_AUTH=${process.env.USE_IAM_AUTH}, AWS_REGION=${process.env.AWS_REGION}`
+                    )
+                }
+
                 try {
                     const chatResult = await languageClient.sendRequest<string | ChatResult>(
                         chatRequestType.method,
@@ -294,6 +319,26 @@ export function registerMessageListeners(
                         },
                         cancellationToken.token
                     )
+
+                    // Add response content logging for SageMaker debugging
+                    if (process.env.USE_IAM_AUTH === 'true') {
+                        languageClient.info(`[SageMaker Debug] Chat response received - type: ${typeof chatResult}`)
+                        if (typeof chatResult === 'string') {
+                            languageClient.info(
+                                `[SageMaker Debug] Chat response (string): ${chatResult.substring(0, 200)}...`
+                            )
+                        } else if (chatResult && typeof chatResult === 'object') {
+                            languageClient.info(
+                                `[SageMaker Debug] Chat response (object keys): ${Object.keys(chatResult)}`
+                            )
+                            if ('message' in chatResult) {
+                                languageClient.info(
+                                    `[SageMaker Debug] Chat response message: ${JSON.stringify(chatResult.message).substring(0, 200)}...`
+                                )
+                            }
+                        }
+                    }
+
                     await handleCompleteResult<ChatResult>(
                         chatResult,
                         encryptionKey,
diff --git a/packages/amazonq/src/lsp/client.ts b/packages/amazonq/src/lsp/client.ts
@@ -16,6 +16,7 @@ import {
     RenameFilesParams,
     ResponseMessage,
     WorkspaceFolder,
+    ConnectionMetadata,
 } from '@aws/language-server-runtimes/protocol'
 import {
     AuthUtil,
@@ -181,10 +182,11 @@ export async function startLanguageServer(
                 contextConfiguration: {
                     workspaceIdentifier: extensionContext.storageUri?.path,
                 },
-                logLevel: toAmazonQLSPLogLevel(globals.logOutputChannel.logLevel),
+                logLevel: isSageMaker() ? 'debug' : toAmazonQLSPLogLevel(globals.logOutputChannel.logLevel),
             },
             credentials: {
                 providesBearerToken: true,
+                providesIam: isSageMaker(), // Enable IAM credentials for SageMaker environments
             },
         },
         /**
@@ -210,6 +212,32 @@ export async function startLanguageServer(
     toDispose.push(disposable)
     await client.onReady()
 
+    // Set up connection metadata handler
+    client.onRequest<ConnectionMetadata, Error>(notificationTypes.getConnectionMetadata.method, () => {
+        // For IAM auth, provide a default startUrl
+        if (process.env.USE_IAM_AUTH === 'true') {
+            getLogger().info(
+                `[SageMaker Debug] Connection metadata requested - returning hardcoded startUrl for IAM auth`
+            )
+            return {
+                sso: {
+                    // TODO P261194666 Replace with correct startUrl once identified
+                    startUrl: 'https://amzn.awsapps.com/start', // Default for IAM auth
+                },
+            }
+        }
+
+        // For SSO auth, use the actual startUrl
+        getLogger().info(
+            `[SageMaker Debug] Connection metadata requested - returning actual startUrl for SSO auth: ${AuthUtil.instance.auth.startUrl}`
+        )
+        return {
+            sso: {
+                startUrl: AuthUtil.instance.auth.startUrl,
+            },
+        }
+    })
+
     const auth = await initializeAuth(client)
 
     await onLanguageServerReady(extensionContext, auth, client, resourcePaths, toDispose)
diff --git a/packages/core/src/auth/auth.ts b/packages/core/src/auth/auth.ts
@@ -1032,6 +1032,16 @@ export class Auth implements AuthService, ConnectionManager {
                 }
             }
         }
+
+        // Add conditional auto-login logic for SageMaker (jmkeyes@ guidance)
+        if (hasVendedIamCredentials() && isSageMaker()) {
+            // SageMaker auto-login logic - use 'ec2' source since SageMaker uses EC2-like instance credentials
+            const sagemakerProfileId = asString({ credentialSource: 'ec2', credentialTypeId: 'sagemaker-instance' })
+            if ((await tryConnection(sagemakerProfileId)) === true) {
+                getLogger().info(`auth: automatically connected with SageMaker credentials`)
+                return
+            }
+        }
     }
 
     /**
diff --git a/packages/core/src/shared/lsp/utils/platform.ts b/packages/core/src/shared/lsp/utils/platform.ts
@@ -3,11 +3,18 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import * as vscode from 'vscode'
 import { ToolkitError } from '../../errors'
 import { Logger } from '../../logger/logger'
 import { ChildProcess } from '../../utilities/processUtils'
 import { waitUntil } from '../../utilities/timeoutUtils'
 import { isDebugInstance } from '../../vscode/env'
+import { isSageMaker } from '../../extensionUtilities'
+import { getLogger } from '../../logger'
+
+interface SagemakerCookie {
+    authMode?: 'Sso' | 'Iam'
+}
 
 export function getNodeExecutableName(): string {
     return process.platform === 'win32' ? 'node.exe' : 'node'
@@ -101,7 +108,44 @@ export function createServerOptions({
             args.unshift('--inspect=6080')
         }
 
-        const lspProcess = new ChildProcess(bin, args, { warnThresholds })
+        // Set USE_IAM_AUTH environment variable for SageMaker environments based on cookie detection
+        // This tells the language server to use IAM authentication mode instead of SSO mode
+        const env = { ...process.env }
+        if (isSageMaker()) {
+            try {
+                // The command `sagemaker.parseCookies` is registered in VS Code SageMaker environment
+                const result = (await vscode.commands.executeCommand('sagemaker.parseCookies')) as SagemakerCookie
+                if (result.authMode !== 'Sso') {
+                    env.USE_IAM_AUTH = 'true'
+                    getLogger().info(
+                        `[SageMaker Debug] Setting USE_IAM_AUTH=true for language server process (authMode: ${result.authMode})`
+                    )
+                } else {
+                    getLogger().info(`[SageMaker Debug] Using SSO auth mode, not setting USE_IAM_AUTH`)
+                }
+            } catch (err) {
+                getLogger().warn(`[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: ${err}`)
+                env.USE_IAM_AUTH = 'true'
+            }
+
+            // Enable verbose logging for Mynah backend to help debug the generic error response
+            env.RUST_LOG = 'debug'
+
+            // Log important environment variables for debugging
+            getLogger().info(`[SageMaker Debug] Environment variables for language server:`)
+            getLogger().info(`[SageMaker Debug] USE_IAM_AUTH: ${env.USE_IAM_AUTH}`)
+            getLogger().info(
+                `[SageMaker Debug] AWS_CONTAINER_CREDENTIALS_RELATIVE_URI: ${env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI}`
+            )
+            getLogger().info(`[SageMaker Debug] AWS_DEFAULT_REGION: ${env.AWS_DEFAULT_REGION}`)
+            getLogger().info(`[SageMaker Debug] AWS_REGION: ${env.AWS_REGION}`)
+            getLogger().info(`[SageMaker Debug] RUST_LOG: ${env.RUST_LOG}`)
+        }
+
+        const lspProcess = new ChildProcess(bin, args, {
+            warnThresholds,
+            spawnOptions: { env },
+        })
 
         // this is a long running process, awaiting it will never resolve
         void lspProcess.run()

Original file line number	Diff line number	Diff line change
`@@ -1032,6 +1032,16 @@ export class Auth implements AuthService, ConnectionManager {`
`1032`	`1032`	`}`
`1033`	`1033`	`}`
`1034`	`1034`	`}`
	`1035`	`+`
	`1036`	`+ // Add conditional auto-login logic for SageMaker (jmkeyes@ guidance)`
	`1037`	`+ if (hasVendedIamCredentials() && isSageMaker()) {`
	`1038`	`+ // SageMaker auto-login logic - use 'ec2' source since SageMaker uses EC2-like instance credentials`
	`1039`	`+ const sagemakerProfileId = asString({ credentialSource: 'ec2', credentialTypeId: 'sagemaker-instance' })`
	`1040`	`+ if ((await tryConnection(sagemakerProfileId)) === true) {`
	`1041`	+ getLogger().info(`auth: automatically connected with SageMaker credentials`)
	`1042`	`+ return`
	`1043`	`+ }`
	`1044`	`+ }`
`1035`	`1045`	`}`
`1036`	`1046`
`1037`	`1047`	`/**`