Skip to content

fix(amazonq): Bug fixes for LSP auth on agentic mode #7231

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
May 6, 2025
Merged

fix(amazonq): Bug fixes for LSP auth on agentic mode #7231

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
e5883dc
fix(amazonq): Bug fixes for LSP auth on agentic mode
opieter-aws May 6, 2025
d2f642f
Remove commented code
opieter-aws May 6, 2025
8ac648f
Fix linter
opieter-aws May 6, 2025
39c62a5
Address comments
opieter-aws May 6, 2025
4c4556f
simplify ShowDocumentRequest handler
opieter-aws May 6, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 5 additions & 18 deletions packages/amazonq/src/extension.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,8 @@
* SPDX-License-Identifier: Apache-2.0
*/

import { CredentialsStore, LoginManager, authUtils, initializeAuth } from 'aws-core-vscode/auth'
import { Auth, AuthUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
import {
activate as activateCodeWhisperer,
AuthUtil,
shutdown as shutdownCodeWhisperer,
} from 'aws-core-vscode/codewhisperer'
import { authUtils, CredentialsStore, LoginManager, initializeAuth } from 'aws-core-vscode/auth'
import { activate as activateCodeWhisperer, shutdown as shutdownCodeWhisperer } from 'aws-core-vscode/codewhisperer'
import { makeEndpointsProvider, registerGenericCommands } from 'aws-core-vscode'
import { CommonAuthWebview } from 'aws-core-vscode/login'
import {
Expand Down Expand Up @@ -39,7 +34,6 @@ import {
Experiments,
isSageMaker,
Commands,
isAmazonInternalOs,
} from 'aws-core-vscode/shared'
import { ExtStartUpSources } from 'aws-core-vscode/telemetry'
import { VSCODE_EXTENSION_ID } from 'aws-core-vscode/utils'
Expand All @@ -50,8 +44,6 @@ import { registerCommands } from './commands'
import { focusAmazonQPanel } from 'aws-core-vscode/codewhispererChat'
import { activate as activateAmazonqLsp } from './lsp/activation'
import { activate as activateInlineCompletion } from './app/inline/activation'
import { isAmazonInternalOs } from 'aws-core-vscode/shared'
import { hasGlibcPatch } from './lsp/client'

export const amazonQContextPrefix = 'amazonq'

Expand Down Expand Up @@ -127,17 +119,12 @@ export async function activateAmazonQCommon(context: vscode.ExtensionContext, is
extensionContext: context,
}

// Auth is dependent on LSP, needs to be activated before CW and Inline
await activateAmazonqLsp(context)

// This contains every lsp agnostic things (auth, security scan, code scan)
await activateCodeWhisperer(extContext as ExtContext)

if (
(Experiments.instance.get('amazonqLSP', true) || AuthUtil.instance.isInternalAmazonUser()) &&
(!isAmazonInternalOs() || (await hasGlibcPatch()))
) {
// start the Amazon Q LSP for internal users first
// for AL2, start LSP if glibc patch is found
await activateAmazonqLsp(context)
}
Comment on lines -133 to -140
Copy link
Contributor Author

@opieter-aws opieter-aws May 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This setup doesn't work with auth on LSP, since the LSP needs to start for auth to work

if (!Experiments.instance.get('amazonqLSPInline', false)) {
await activateInlineCompletion()
}
Expand Down
3 changes: 0 additions & 3 deletions packages/amazonq/src/extensionNode.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ import api from './api'
import { activate as activateCWChat } from './app/chat/activation'
import { activate as activateInlineChat } from './inlineChat/activation'
import { beta } from 'aws-core-vscode/dev'
import * as amazonq from 'aws-core-vscode/amazonq'
import { activate as activateNotifications, NotificationsController } from 'aws-core-vscode/notifications'
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
import { telemetry, AuthUserState } from 'aws-core-vscode/telemetry'
Expand Down Expand Up @@ -68,8 +67,6 @@ async function activateAmazonQNode(context: vscode.ExtensionContext) {
await activateTransformationHub(extContext as ExtContext)
activateInlineChat(context)

context.subscriptions.push(amazonq.focusAmazonQPanel.register(), amazonq.focusAmazonQPanelKeybinding.register())

Comment on lines -71 to -72
Copy link
Contributor Author

@opieter-aws opieter-aws May 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Now part of

aws-toolkit-vscode/packages/amazonq/src/app/chat/activation.ts

Line 31 in f1f5a36

amazonq.focusAmazonQPanel.register(),

const authProvider = new CommonAuthViewProvider(
context,
amazonQContextPrefix,
Expand Down
6 changes: 2 additions & 4 deletions packages/amazonq/src/lsp/activation.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,19 +4,17 @@
*/

import vscode from 'vscode'
import { clientId, encryptionKey, startLanguageServer } from './client'
import { startLanguageServer } from './client'
import { AmazonQLspInstaller } from './lspInstaller'
import { lspSetupStage, ToolkitError, messages } from 'aws-core-vscode/shared'
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
import { auth2 } from 'aws-core-vscode/auth'

export async function activate(ctx: vscode.ExtensionContext) {
try {
const client = await lspSetupStage('all', async () => {
await lspSetupStage('all', async () => {
const installResult = await new AmazonQLspInstaller().resolve()
return await lspSetupStage('launch', () => startLanguageServer(ctx, installResult.resourcePaths))
})
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
Copy link
Contributor Author

@opieter-aws opieter-aws May 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move instantiation of the AuthUtil singleton into startLanguageServer. The function uses several AuthUtil calls, which fail if it is not instantiated yet

await AuthUtil.instance.restore()
} catch (err) {
const e = err as ToolkitError
Expand Down
20 changes: 0 additions & 20 deletions packages/amazonq/src/lsp/chat/messages.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,6 @@ import {
ShowSaveFileDialogParams,
LSPErrorCodes,
tabBarActionRequestType,
ShowDocumentParams,
ShowDocumentResult,
ShowDocumentRequest,
contextCommandsNotificationType,
ContextCommandParams,
openFileDiffNotificationType,
Expand Down Expand Up @@ -423,23 +420,6 @@ export function registerMessageListeners(
}
})

languageClient.onRequest<ShowDocumentParams, ShowDocumentResult>(
ShowDocumentRequest.method,
async (params: ShowDocumentParams): Promise<ShowDocumentParams | ResponseError<ShowDocumentResult>> => {
try {
const uri = vscode.Uri.parse(params.uri)
const doc = await vscode.workspace.openTextDocument(uri)
await vscode.window.showTextDocument(doc, { preview: false })
return params
} catch (e) {
return new ResponseError(
LSPErrorCodes.RequestFailed,
`Failed to open document: ${(e as Error).message}`
)
}
}
)

languageClient.onNotification(contextCommandsNotificationType.method, (params: ContextCommandParams) => {
void provider.webview?.postMessage({
command: contextCommandsNotificationType.method,
Expand Down
2 changes: 1 addition & 1 deletion packages/amazonq/src/lsp/chat/webviewProvider.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,7 +137,7 @@ export class AmazonQChatViewProvider implements WebviewViewProvider {
let qChat = undefined
const init = () => {
const vscodeApi = acquireVsCodeApi()
const hybridChatConnector = new HybridChatAdapter(${(await AuthUtil.instance.getChatAuthState()).amazonQ === 'connected'},${featureConfigData},${welcomeCount},${disclaimerAcknowledged},${regionProfileString},${disabledCommands},${isSMUS},${isSM},vscodeApi.postMessage)
const hybridChatConnector = new HybridChatAdapter(${AuthUtil.instance.isConnected()},${featureConfigData},${welcomeCount},${disclaimerAcknowledged},${regionProfileString},${disabledCommands},${isSMUS},${isSM},vscodeApi.postMessage)
const commands = [hybridChatConnector.initialQuickActions[0]]
qChat = amazonQChat.createChat(vscodeApi, {disclaimerAcknowledged: ${disclaimerAcknowledged}, pairProgrammingAcknowledged: ${pairProgrammingAcknowledged}, agenticMode: true, quickActionCommands: commands}, hybridChatConnector, ${JSON.stringify(featureConfigData)});
}
Expand Down
53 changes: 33 additions & 20 deletions packages/amazonq/src/lsp/client.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,12 +23,14 @@ import {
GetSsoTokenProgressToken,
GetSsoTokenProgressType,
MessageActionItem,
ShowDocumentParams,
ShowDocumentRequest,
ShowDocumentResult,
ShowMessageRequest,
ShowMessageRequestParams,
ConnectionMetadata,
ShowDocumentRequest,
ShowDocumentParams,
ShowDocumentResult,
ResponseError,
LSPErrorCodes,
} from '@aws/language-server-runtimes/protocol'
import { AuthUtil, CodeWhispererSettings, getSelectedCustomization } from 'aws-core-vscode/codewhisperer'
import {
Expand All @@ -37,17 +39,17 @@ import {
globals,
Experiments,
Commands,
openUrl,
validateNodeExe,
getLogger,
undefinedIfEmpty,
getOptOutPreference,
isAmazonInternalOs,
fs,
oidcClientName,
openUrl,
} from 'aws-core-vscode/shared'
import { processUtils } from 'aws-core-vscode/shared'
import { activate } from './chat/activation'
import { activate as activateChat } from './chat/activation'
import { AmazonQResourcePaths } from './lspInstaller'
import { auth2 } from 'aws-core-vscode/auth'
import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config'
Expand Down Expand Up @@ -144,9 +146,6 @@ export async function startLanguageServer(
notifications: true,
showSaveFileDialog: true,
},
q: {
Copy link
Contributor

@jpinkney-aws jpinkney-aws May 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't think this can be dropped

Copy link
Contributor Author

@opieter-aws opieter-aws May 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The merge with mainline made it duplicate, see line 140

developerProfiles: true,
},
},
logLevel: toAmazonQLSPLogLevel(globals.logOutputChannel.logLevel),
},
Expand All @@ -173,6 +172,7 @@ export async function startLanguageServer(
toDispose.push(disposable)

await client.onReady()
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))

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

client.onRequest<ShowDocumentResult, Error>(ShowDocumentRequest.method, async (params: ShowDocumentParams) => {
try {
return { success: await openUrl(vscode.Uri.parse(params.uri), lspName) }
} catch (err: any) {
getLogger().error(`Failed to open document for LSP: ${lspName}, error: %s`, err)
return { success: false }
}
})

client.onRequest<MessageActionItem | null, Error>(
ShowMessageRequest.method,
async (params: ShowMessageRequestParams) => {
Expand All @@ -201,6 +192,30 @@ export async function startLanguageServer(
}
)

client.onRequest<ShowDocumentParams, ShowDocumentResult>(
ShowDocumentRequest.method,
async (params: ShowDocumentParams): Promise<ShowDocumentParams | ResponseError<ShowDocumentResult>> => {
const uri = vscode.Uri.parse(params.uri)
getLogger().info(`Processing ShowDocumentRequest for URI scheme: ${uri.scheme}`)
try {
if (uri.scheme.startsWith('http')) {
getLogger().info('Opening URL...')
await openUrl(vscode.Uri.parse(params.uri))
} else {
getLogger().info('Opening text document...')
const doc = await vscode.workspace.openTextDocument(uri)
await vscode.window.showTextDocument(doc, { preview: false })
}
return params
} catch (e) {
return new ResponseError(
LSPErrorCodes.RequestFailed,
`Failed to process ShowDocumentRequest: ${(e as Error).message}`
)
}
}
)

const sendProfileToLsp = async () => {
try {
const result = await client.sendRequest(updateConfigurationRequestType.method, {
Expand Down Expand Up @@ -266,9 +281,7 @@ export async function startLanguageServer(
)
}

if (Experiments.instance.get('amazonqChatLSP', false)) {
activate(client, encryptionKey, resourcePaths.ui)
}
await activateChat(client, encryptionKey, resourcePaths.ui)

toDispose.push(
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(sendProfileToLsp),
Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/codewhisperer/util/authUtil.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,9 @@ export class AuthUtil implements IAuthProvider {

async restore() {
await this.session.restore()
if (!this.isConnected()) {
await this.refreshState()
}
Comment on lines +91 to +93
Copy link
Contributor Author

@opieter-aws opieter-aws May 6, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This ensures that a session is restored successfully upon start of the extension with expired SSO. If SSO is expired upon start of extension, the LSP throws an error. Since there is no state change in that case (notConnected > notConnected), we can't rely on the state change handler to trigger a login flow

}

async login(startUrl: string, region: string) {
Expand Down
Loading