Skip to content

Commit 276f9ef

Browse files
floralphfloralph
authored
Merge pull request #7667 from aws/floralph/P261194666
fix(amazonq): Q chat stopped using IAM creds in the Amazon Q v1.63.0 release
2 parents d50d38e + 2d47f6e commit 276f9ef

File tree

9 files changed

+818
-15
lines changed

9 files changed

+818
-15
lines changed

‎P261194666.md

Lines changed: 630 additions & 0 deletions
Large diffs are not rendered by default.

‎packages/amazonq/src/lsp/auth.ts

Lines changed: 57 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import {
77
bearerCredentialsUpdateRequestType,
8+
iamCredentialsUpdateRequestType,
89
ConnectionMetadata,
910
NotificationType,
1011
RequestType,
@@ -17,8 +18,8 @@ import { LanguageClient } from 'vscode-languageclient'
1718
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
1819
import { Writable } from 'stream'
1920
import { onceChanged } from 'aws-core-vscode/utils'
20-
import { getLogger, oneMinute } from 'aws-core-vscode/shared'
21-
import { isSsoConnection } from 'aws-core-vscode/auth'
21+
import { getLogger, oneMinute, isSageMaker } from 'aws-core-vscode/shared'
22+
import { isSsoConnection, isIamConnection } from 'aws-core-vscode/auth'
2223

2324
export const encryptionKey = crypto.randomBytes(32)
2425

@@ -78,10 +79,16 @@ export class AmazonQLspAuth {
7879
*/
7980
async refreshConnection(force: boolean = false) {
8081
const activeConnection = this.authUtil.conn
81-
if (this.authUtil.isConnectionValid() && isSsoConnection(activeConnection)) {
82-
// send the token to the language server
83-
const token = await this.authUtil.getBearerToken()
84-
await (force ? this._updateBearerToken(token) : this.updateBearerToken(token))
82+
if (this.authUtil.isConnectionValid()) {
83+
if (isSsoConnection(activeConnection)) {
84+
// Existing SSO path
85+
const token = await this.authUtil.getBearerToken()
86+
await (force ? this._updateBearerToken(token) : this.updateBearerToken(token))
87+
} else if (isSageMaker() && isIamConnection(activeConnection)) {
88+
// New SageMaker IAM path
89+
const credentials = await this.authUtil.getCredentials()
90+
await (force ? this._updateIamCredentials(credentials) : this.updateIamCredentials(credentials))
91+
}
8592
}
8693
}
8794

@@ -92,9 +99,7 @@ export class AmazonQLspAuth {
9299

93100
public updateBearerToken = onceChanged(this._updateBearerToken.bind(this))
94101
private async _updateBearerToken(token: string) {
95-
const request = await this.createUpdateCredentialsRequest({
96-
token,
97-
})
102+
const request = await this.createUpdateBearerCredentialsRequest(token)
98103

99104
// "aws/credentials/token/update"
100105
// https://github.com/aws/language-servers/blob/44d81f0b5754747d77bda60b40cc70950413a737/core/aws-lsp-core/src/credentials/credentialsProvider.ts#L27
@@ -103,15 +108,36 @@ export class AmazonQLspAuth {
103108
this.client.info(`UpdateBearerToken: ${JSON.stringify(request)}`)
104109
}
105110

111+
public updateIamCredentials = onceChanged(this._updateIamCredentials.bind(this))
112+
private async _updateIamCredentials(credentials: any) {
113+
getLogger().info(
114+
`[SageMaker Debug] Updating IAM credentials - credentials received: ${credentials ? 'YES' : 'NO'}`
115+
)
116+
if (credentials) {
117+
getLogger().info(
118+
`[SageMaker Debug] IAM credentials structure: accessKeyId=${credentials.accessKeyId ? 'present' : 'missing'}, secretAccessKey=${credentials.secretAccessKey ? 'present' : 'missing'}, sessionToken=${credentials.sessionToken ? 'present' : 'missing'}`
119+
)
120+
}
121+
122+
const request = await this.createUpdateIamCredentialsRequest(credentials)
123+
124+
// "aws/credentials/iam/update"
125+
await this.client.sendRequest(iamCredentialsUpdateRequestType.method, request)
126+
127+
this.client.info(`UpdateIamCredentials: ${JSON.stringify(request)}`)
128+
getLogger().info(`[SageMaker Debug] IAM credentials update request sent successfully`)
129+
}
130+
106131
public startTokenRefreshInterval(pollingTime: number = oneMinute / 2) {
107132
const interval = setInterval(async () => {
108133
await this.refreshConnection().catch((e) => this.logRefreshError(e))
109134
}, pollingTime)
110135
return interval
111136
}
112137

113-
private async createUpdateCredentialsRequest(data: any): Promise<UpdateCredentialsParams> {
114-
const payload = new TextEncoder().encode(JSON.stringify({ data }))
138+
private async createUpdateBearerCredentialsRequest(token: string): Promise<UpdateCredentialsParams> {
139+
const bearerCredentials = { token }
140+
const payload = new TextEncoder().encode(JSON.stringify({ data: bearerCredentials }))
115141

116142
const jwt = await new jose.CompactEncrypt(payload)
117143
.setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
@@ -127,4 +153,24 @@ export class AmazonQLspAuth {
127153
encrypted: true,
128154
}
129155
}
156+
157+
private async createUpdateIamCredentialsRequest(credentials: any): Promise<UpdateCredentialsParams> {
158+
// Extract IAM credentials structure
159+
const iamCredentials = {
160+
accessKeyId: credentials.accessKeyId,
161+
secretAccessKey: credentials.secretAccessKey,
162+
sessionToken: credentials.sessionToken,
163+
}
164+
const payload = new TextEncoder().encode(JSON.stringify({ data: iamCredentials }))
165+
166+
const jwt = await new jose.CompactEncrypt(payload)
167+
.setProtectedHeader({ alg: 'dir', enc: 'A256GCM' })
168+
.encrypt(encryptionKey)
169+
170+
return {
171+
data: jwt,
172+
// Omit metadata for IAM credentials since startUrl is undefined for non-SSO connections
173+
encrypted: true,
174+
}
175+
}
130176
}

‎packages/amazonq/src/lsp/chat/messages.ts

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -295,6 +295,31 @@ export function registerMessageListeners(
295295
}
296296

297297
const chatRequest = await encryptRequest<ChatParams>(chatParams, encryptionKey)
298+
299+
// Add detailed logging for SageMaker debugging
300+
if (process.env.USE_IAM_AUTH === 'true') {
301+
languageClient.info(`[SageMaker Debug] Making chat request with IAM auth`)
302+
languageClient.info(`[SageMaker Debug] Chat request method: ${chatRequestType.method}`)
303+
languageClient.info(
304+
`[SageMaker Debug] Original chat params: ${JSON.stringify(
305+
{
306+
tabId: chatParams.tabId,
307+
prompt: chatParams.prompt,
308+
// Don't log full textDocument content, just metadata
309+
textDocument: chatParams.textDocument
310+
? { uri: chatParams.textDocument.uri }
311+
: undefined,
312+
context: chatParams.context ? `${chatParams.context.length} context items` : undefined,
313+
},
314+
undefined,
315+
2
316+
)}`
317+
)
318+
languageClient.info(
319+
`[SageMaker Debug] Environment context: USE_IAM_AUTH=${process.env.USE_IAM_AUTH}, AWS_REGION=${process.env.AWS_REGION}`
320+
)
321+
}
322+
298323
try {
299324
const chatResult = await languageClient.sendRequest<string | ChatResult>(
300325
chatRequestType.method,
@@ -304,6 +329,26 @@ export function registerMessageListeners(
304329
},
305330
cancellationToken.token
306331
)
332+
333+
// Add response content logging for SageMaker debugging
334+
if (process.env.USE_IAM_AUTH === 'true') {
335+
languageClient.info(`[SageMaker Debug] Chat response received - type: ${typeof chatResult}`)
336+
if (typeof chatResult === 'string') {
337+
languageClient.info(
338+
`[SageMaker Debug] Chat response (string): ${chatResult.substring(0, 200)}...`
339+
)
340+
} else if (chatResult && typeof chatResult === 'object') {
341+
languageClient.info(
342+
`[SageMaker Debug] Chat response (object keys): ${Object.keys(chatResult)}`
343+
)
344+
if ('message' in chatResult) {
345+
languageClient.info(
346+
`[SageMaker Debug] Chat response message: ${JSON.stringify(chatResult.message).substring(0, 200)}...`
347+
)
348+
}
349+
}
350+
}
351+
307352
await handleCompleteResult<ChatResult>(
308353
chatResult,
309354
encryptionKey,

‎packages/amazonq/src/lsp/client.ts

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import {
1616
RenameFilesParams,
1717
ResponseMessage,
1818
WorkspaceFolder,
19+
ConnectionMetadata,
1920
} from '@aws/language-server-runtimes/protocol'
2021
import {
2122
AuthUtil,
@@ -184,10 +185,11 @@ export async function startLanguageServer(
184185
contextConfiguration: {
185186
workspaceIdentifier: extensionContext.storageUri?.path,
186187
},
187-
logLevel: toAmazonQLSPLogLevel(globals.logOutputChannel.logLevel),
188+
logLevel: isSageMaker() ? 'debug' : toAmazonQLSPLogLevel(globals.logOutputChannel.logLevel),
188189
},
189190
credentials: {
190191
providesBearerToken: true,
192+
providesIam: isSageMaker(), // Enable IAM credentials for SageMaker environments
191193
},
192194
},
193195
/**
@@ -213,6 +215,32 @@ export async function startLanguageServer(
213215
toDispose.push(disposable)
214216
await client.onReady()
215217

218+
// Set up connection metadata handler
219+
client.onRequest<ConnectionMetadata, Error>(notificationTypes.getConnectionMetadata.method, () => {
220+
// For IAM auth, provide a default startUrl
221+
if (process.env.USE_IAM_AUTH === 'true') {
222+
getLogger().info(
223+
`[SageMaker Debug] Connection metadata requested - returning hardcoded startUrl for IAM auth`
224+
)
225+
return {
226+
sso: {
227+
// TODO P261194666 Replace with correct startUrl once identified
228+
startUrl: 'https://amzn.awsapps.com/start', // Default for IAM auth
229+
},
230+
}
231+
}
232+
233+
// For SSO auth, use the actual startUrl
234+
getLogger().info(
235+
`[SageMaker Debug] Connection metadata requested - returning actual startUrl for SSO auth: ${AuthUtil.instance.auth.startUrl}`
236+
)
237+
return {
238+
sso: {
239+
startUrl: AuthUtil.instance.auth.startUrl,
240+
},
241+
}
242+
})
243+
216244
const auth = await initializeAuth(client)
217245

218246
await onLanguageServerReady(extensionContext, auth, client, resourcePaths, toDispose)

‎packages/core/package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -455,7 +455,7 @@
455455
"webpackDev": "webpack --mode development",
456456
"serveVue": "ts-node ./scripts/build/checkServerPort.ts && webpack serve --port 8080 --config-name vue --mode development",
457457
"watch": "npm run clean && npm run buildScripts && npm run compileOnly -- --watch",
458-
"lint": "ts-node ./scripts/lint/testLint.ts",
458+
"lint": "node --max-old-space-size=8192 -r ts-node/register ./scripts/lint/testLint.ts",
459459
"generateClients": "ts-node ./scripts/build/generateServiceClient.ts ",
460460
"generateIcons": "ts-node ../../scripts/generateIcons.ts",
461461
"generateTelemetry": "node ../../node_modules/@aws-toolkits/telemetry/lib/generateTelemetry.js --extraInput=src/shared/telemetry/vscodeTelemetry.json --output=src/shared/telemetry/telemetry.gen.ts"

‎packages/core/scripts/lint/testLint.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,9 @@ void (async () => {
99
try {
1010
console.log('Running linting tests...')
1111

12-
const mocha = new Mocha()
12+
const mocha = new Mocha({
13+
timeout: 5000,
14+
})
1315

1416
const testFiles = await glob('dist/src/testLint/**/*.test.js')
1517
for (const file of testFiles) {

‎packages/core/src/auth/auth.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,6 +1032,16 @@ export class Auth implements AuthService, ConnectionManager {
10321032
}
10331033
}
10341034
}
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+
}
10351045
}
10361046

10371047
/**

‎packages/core/src/shared/lsp/utils/platform.ts

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,18 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6+
import * as vscode from 'vscode'
67
import { ToolkitError } from '../../errors'
78
import { Logger } from '../../logger/logger'
89
import { ChildProcess } from '../../utilities/processUtils'
910
import { waitUntil } from '../../utilities/timeoutUtils'
1011
import { isDebugInstance } from '../../vscode/env'
12+
import { isSageMaker } from '../../extensionUtilities'
13+
import { getLogger } from '../../logger/logger'
14+
15+
interface SagemakerCookie {
16+
authMode?: 'Sso' | 'Iam'
17+
}
1118

1219
export function getNodeExecutableName(): string {
1320
return process.platform === 'win32' ? 'node.exe' : 'node'
@@ -101,7 +108,40 @@ export function createServerOptions({
101108
args.unshift('--inspect=6080')
102109
}
103110

104-
const lspProcess = new ChildProcess(bin, args, { warnThresholds })
111+
// Set USE_IAM_AUTH environment variable for SageMaker environments based on cookie detection
112+
// This tells the language server to use IAM authentication mode instead of SSO mode
113+
const env = { ...process.env }
114+
if (isSageMaker()) {
115+
try {
116+
// The command `sagemaker.parseCookies` is registered in VS Code SageMaker environment
117+
const result = (await vscode.commands.executeCommand('sagemaker.parseCookies')) as SagemakerCookie
118+
if (result.authMode !== 'Sso') {
119+
env.USE_IAM_AUTH = 'true'
120+
getLogger().info(
121+
`[SageMaker Debug] Setting USE_IAM_AUTH=true for language server process (authMode: ${result.authMode})`
122+
)
123+
} else {
124+
getLogger().info(`[SageMaker Debug] Using SSO auth mode, not setting USE_IAM_AUTH`)
125+
}
126+
} catch (err) {
127+
getLogger().warn(`[SageMaker Debug] Failed to parse SageMaker cookies, defaulting to IAM auth: ${err}`)
128+
env.USE_IAM_AUTH = 'true'
129+
}
130+
131+
// Log important environment variables for debugging
132+
getLogger().info(`[SageMaker Debug] Environment variables for language server:`)
133+
getLogger().info(`[SageMaker Debug] USE_IAM_AUTH: ${env.USE_IAM_AUTH}`)
134+
getLogger().info(
135+
`[SageMaker Debug] AWS_CONTAINER_CREDENTIALS_RELATIVE_URI: ${env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI}`
136+
)
137+
getLogger().info(`[SageMaker Debug] AWS_DEFAULT_REGION: ${env.AWS_DEFAULT_REGION}`)
138+
getLogger().info(`[SageMaker Debug] AWS_REGION: ${env.AWS_REGION}`)
139+
}
140+
141+
const lspProcess = new ChildProcess(bin, args, {
142+
warnThresholds,
143+
spawnOptions: { env },
144+
})
105145

106146
// this is a long running process, awaiting it will never resolve
107147
void lspProcess.run()

‎packages/core/src/testLint/eslint.test.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ describe('eslint', function () {
1212
it('passes eslint', function () {
1313
const result = runCmd(
1414
[
15+
'node',
16+
'--max-old-space-size=8192',
1517
'../../node_modules/.bin/eslint',
1618
'-c',
1719
'../../.eslintrc.js',

0 commit comments

Comments
 (0)