Skip to content

Commit 8787af3

Browse files
opieter-awsopieter-aws
authored
fix(amazonq): Bug fixes for LSP auth on agentic mode (#7231)
## Problem Agentic chat introduced code changes that break auth on LSP ## Solution Fixes for the auth flow: * Move instantiation of `AuthUtil` into `startLanguageServer` sinc ethe function uses several AuthUtil calls, which fail if it is not instantiated yet * Remove the dependency of LSP on glibc patch, since auth doesn't work if LSP is not started * Add URL opening capability to the new `ShowDocumentRequest` handler, and remove the old handler from `client.ts`. * Fix bug where session isn't restored upon start of extension when SSO is expired. --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 0363ba4 commit 8787af3

File tree

7 files changed

+44
-66
lines changed

7 files changed

+44
-66
lines changed

‎packages/amazonq/src/extension.ts‎

Lines changed: 5 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,8 @@
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

6-
import { CredentialsStore, LoginManager, authUtils, initializeAuth } from 'aws-core-vscode/auth'
7-
import { Auth, AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
8-
import {
9-
activate as activateCodeWhisperer,
10-
AuthUtil,
11-
shutdown as shutdownCodeWhisperer,
12-
} from 'aws-core-vscode/codewhisperer'
6+
import { authUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
7+
import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
138
import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
149
import { CommonAuthWebview } from 'aws-core-vscode/login'
1510
import {
@@ -39,7 +34,6 @@ import {
3934
Experiments,
4035
isSageMaker,
4136
Commands,
42-
isAmazonInternalOs,
4337
} from 'aws-core-vscode/shared'
4438
import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
4539
import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
@@ -50,8 +44,6 @@ import { registerCommands } from './commands'
5044
import { focusAmazonQPanel } from 'aws-core-vscode/codewhispererChat'
5145
import { activate as activateAmazonqLsp } from './lsp/activation'
5246
import { activate as activateInlineCompletion } from './app/inline/activation'
53-
import { isAmazonInternalOs } from 'aws-core-vscode/shared'
54-
import { hasGlibcPatch } from './lsp/client'
5547

5648
export const amazonQContextPrefix = 'amazonq'
5749

@@ -127,17 +119,12 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
127119
extensionContext: context,
128120
}
129121

122+
// Auth is dependent on LSP, needs to be activated before CW and Inline
123+
await activateAmazonqLsp(context)
124+
130125
// This contains every lsp agnostic things (auth, security scan, code scan)
131126
await activateCodeWhisperer(extContext as ExtContext)
132127

133-
if (
134-
(Experiments.instance.get('amazonqLSP', true) || AuthUtil.instance.isInternalAmazonUser()) &&
135-
(!isAmazonInternalOs() || (await hasGlibcPatch()))
136-
) {
137-
// start the Amazon Q LSP for internal users first
138-
// for AL2, start LSP if glibc patch is found
139-
await activateAmazonqLsp(context)
140-
}
141128
if (!Experiments.instance.get('amazonqLSPInline', false)) {
142129
await activateInlineCompletion()
143130
}

‎packages/amazonq/src/extensionNode.ts‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@ import api from './api'
1919
import { activate as activateCWChat } from './app/chat/activation'
2020
import { activate as activateInlineChat } from './inlineChat/activation'
2121
import { beta } from 'aws-core-vscode/dev'
22-
import * as amazonq from 'aws-core-vscode/amazonq'
2322
import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications'
2423
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
2524
import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
@@ -68,8 +67,6 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
6867
await activateTransformationHub(extContext as ExtContext)
6968
activateInlineChat(context)
7069

71-
context.subscriptions.push(amazonq.focusAmazonQPanel.register(), amazonq.focusAmazonQPanelKeybinding.register())
72-
7370
const authProvider = new CommonAuthViewProvider(
7471
context,
7572
amazonQContextPrefix,

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

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,19 +4,17 @@
44
*/
55

66
import vscode from 'vscode'
7-
import { clientId, encryptionKey, startLanguageServer } from './client'
7+
import { startLanguageServer } from './client'
88
import { AmazonQLspInstaller } from './lspInstaller'
99
import { lspSetupStage, ToolkitError, messages } from 'aws-core-vscode/shared'
1010
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
11-
import { auth2 } from 'aws-core-vscode/auth'
1211

1312
export async function activate(ctx: vscode.ExtensionContext) {
1413
try {
15-
const client = await lspSetupStage('all', async () => {
14+
await lspSetupStage('all', async () => {
1615
const installResult = await new AmazonQLspInstaller().resolve()
1716
return await lspSetupStage('launch', () => startLanguageServer(ctx, installResult.resourcePaths))
1817
})
19-
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
2018
await AuthUtil.instance.restore()
2119
} catch (err) {
2220
const e = err as ToolkitError

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

Lines changed: 0 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,6 @@ import {
3636
ShowSaveFileDialogParams,
3737
LSPErrorCodes,
3838
tabBarActionRequestType,
39-
ShowDocumentParams,
40-
ShowDocumentResult,
41-
ShowDocumentRequest,
4239
contextCommandsNotificationType,
4340
ContextCommandParams,
4441
openFileDiffNotificationType,
@@ -423,23 +420,6 @@ export function registerMessageListeners(
423420
}
424421
})
425422

426-
languageClient.onRequest<ShowDocumentParams, ShowDocumentResult>(
427-
ShowDocumentRequest.method,
428-
async (params: ShowDocumentParams): Promise<ShowDocumentParams | ResponseError<ShowDocumentResult>> => {
429-
try {
430-
const uri = vscode.Uri.parse(params.uri)
431-
const doc = await vscode.workspace.openTextDocument(uri)
432-
await vscode.window.showTextDocument(doc, { preview: false })
433-
return params
434-
} catch (e) {
435-
return new ResponseError(
436-
LSPErrorCodes.RequestFailed,
437-
`Failed to open document: ${(e as Error).message}`
438-
)
439-
}
440-
}
441-
)
442-
443423
languageClient.onNotification(contextCommandsNotificationType.method, (params: ContextCommandParams) => {
444424
void provider.webview?.postMessage({
445425
command: contextCommandsNotificationType.method,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -137,7 +137,7 @@ export class AmazonQChatViewProvider implements WebviewViewProvider {
137137
let qChat = undefined
138138
const init = () => {
139139
const vscodeApi = acquireVsCodeApi()
140-
const hybridChatConnector = new HybridChatAdapter(${(await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'},${featureConfigData},${welcomeCount},${disclaimerAcknowledged},${regionProfileString},${disabledCommands},${isSMUS},${isSM},vscodeApi.postMessage)
140+
const hybridChatConnector = new HybridChatAdapter(${AuthUtil.instance.isConnected()},${featureConfigData},${welcomeCount},${disclaimerAcknowledged},${regionProfileString},${disabledCommands},${isSMUS},${isSM},vscodeApi.postMessage)
141141
const commands = [hybridChatConnector.initialQuickActions[0]]
142142
qChat = amazonQChat.createChat(vscodeApi, {disclaimerAcknowledged: ${disclaimerAcknowledged}, pairProgrammingAcknowledged: ${pairProgrammingAcknowledged}, agenticMode: true, quickActionCommands: commands}, hybridChatConnector, ${JSON.stringify(featureConfigData)});
143143
}

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

Lines changed: 33 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -23,12 +23,14 @@ import {
2323
GetSsoTokenProgressToken,
2424
GetSsoTokenProgressType,
2525
MessageActionItem,
26-
ShowDocumentParams,
27-
ShowDocumentRequest,
28-
ShowDocumentResult,
2926
ShowMessageRequest,
3027
ShowMessageRequestParams,
3128
ConnectionMetadata,
29+
ShowDocumentRequest,
30+
ShowDocumentParams,
31+
ShowDocumentResult,
32+
ResponseError,
33+
LSPErrorCodes,
3234
} from '@aws/language-server-runtimes/protocol'
3335
import { AuthUtil, CodeWhispererSettings, getSelectedCustomization } from 'aws-core-vscode/codewhisperer'
3436
import {
@@ -37,17 +39,17 @@ import {
3739
globals,
3840
Experiments,
3941
Commands,
40-
openUrl,
4142
validateNodeExe,
4243
getLogger,
4344
undefinedIfEmpty,
4445
getOptOutPreference,
4546
isAmazonInternalOs,
4647
fs,
4748
oidcClientName,
49+
openUrl,
4850
} from 'aws-core-vscode/shared'
4951
import { processUtils } from 'aws-core-vscode/shared'
50-
import { activate } from './chat/activation'
52+
import { activate as activateChat } from './chat/activation'
5153
import { AmazonQResourcePaths } from './lspInstaller'
5254
import { auth2 } from 'aws-core-vscode/auth'
5355
import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config'
@@ -144,9 +146,6 @@ export async function startLanguageServer(
144146
notifications: true,
145147
showSaveFileDialog: true,
146148
},
147-
q: {
148-
developerProfiles: true,
149-
},
150149
},
151150
logLevel: toAmazonQLSPLogLevel(globals.logOutputChannel.logLevel),
152151
},
@@ -173,6 +172,7 @@ export async function startLanguageServer(
173172
toDispose.push(disposable)
174173

175174
await client.onReady()
175+
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
176176

177177
// Request handler for when the server wants to know about the clients auth connnection. Must be registered before the initial auth init call
178178
client.onRequest<ConnectionMetadata, Error>(auth2.notificationTypes.getConnectionMetadata.method, () => {
@@ -183,15 +183,6 @@ export async function startLanguageServer(
183183
}
184184
})
185185

186-
client.onRequest<ShowDocumentResult, Error>(ShowDocumentRequest.method, async (params: ShowDocumentParams) => {
187-
try {
188-
return { success: await openUrl(vscode.Uri.parse(params.uri), lspName) }
189-
} catch (err: any) {
190-
getLogger().error(`Failed to open document for LSP: ${lspName}, error: %s`, err)
191-
return { success: false }
192-
}
193-
})
194-
195186
client.onRequest<MessageActionItem | null, Error>(
196187
ShowMessageRequest.method,
197188
async (params: ShowMessageRequestParams) => {
@@ -201,6 +192,30 @@ export async function startLanguageServer(
201192
}
202193
)
203194

195+
client.onRequest<ShowDocumentParams, ShowDocumentResult>(
196+
ShowDocumentRequest.method,
197+
async (params: ShowDocumentParams): Promise<ShowDocumentParams | ResponseError<ShowDocumentResult>> => {
198+
const uri = vscode.Uri.parse(params.uri)
199+
getLogger().info(`Processing ShowDocumentRequest for URI scheme: ${uri.scheme}`)
200+
try {
201+
if (uri.scheme.startsWith('http')) {
202+
getLogger().info('Opening URL...')
203+
await openUrl(vscode.Uri.parse(params.uri))
204+
} else {
205+
getLogger().info('Opening text document...')
206+
const doc = await vscode.workspace.openTextDocument(uri)
207+
await vscode.window.showTextDocument(doc, { preview: false })
208+
}
209+
return params
210+
} catch (e) {
211+
return new ResponseError(
212+
LSPErrorCodes.RequestFailed,
213+
`Failed to process ShowDocumentRequest: ${(e as Error).message}`
214+
)
215+
}
216+
}
217+
)
218+
204219
const sendProfileToLsp = async () => {
205220
try {
206221
const result = await client.sendRequest(updateConfigurationRequestType.method, {
@@ -266,9 +281,7 @@ export async function startLanguageServer(
266281
)
267282
}
268283

269-
if (Experiments.instance.get('amazonqChatLSP', false)) {
270-
activate(client, encryptionKey, resourcePaths.ui)
271-
}
284+
await activateChat(client, encryptionKey, resourcePaths.ui)
272285

273286
toDispose.push(
274287
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp),

‎packages/core/src/codewhisperer/util/authUtil.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,9 @@ export class AuthUtil implements IAuthProvider {
8888

8989
async restore() {
9090
await this.session.restore()
91+
if (!this.isConnected()) {
92+
await this.refreshState()
93+
}
9194
}
9295

9396
async login(startUrl: string, region: string) {

0 commit comments

Comments
 (0)