Skip to content

Commit f849b3f

Browse files
authored
Merge pull request #7302 from nkomonen-amazon/customizationFix
fix(amazonq): Customization/Profiles not being sent to LSP on startup
2 parents 4e1d059 + acd6478 commit f849b3f

File tree

5 files changed

+105
-68
lines changed

5 files changed

+105
-68
lines changed

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

Lines changed: 34 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,8 @@ import { registerLanguageServerEventListener, registerMessageListeners } from '.
1111
import { Commands, getLogger, globals, undefinedIfEmpty } from 'aws-core-vscode/shared'
1212
import { activate as registerLegacyChatListeners } from '../../app/chat/activation'
1313
import { DefaultAmazonQAppInitContext } from 'aws-core-vscode/amazonq'
14-
import { AuthUtil, getSelectedCustomization } from 'aws-core-vscode/codewhisperer'
15-
import {
16-
DidChangeConfigurationNotification,
17-
updateConfigurationRequestType,
18-
} from '@aws/language-server-runtimes/protocol'
14+
import { AuthUtil, getSelectedCustomization, notifyNewCustomizations } from 'aws-core-vscode/codewhisperer'
15+
import { pushConfigUpdate } from '../config'
1916

2017
export async function activate(languageClient: LanguageClient, encryptionKey: Buffer, mynahUIPath: string) {
2118
const disposables = globals.context.subscriptions
@@ -25,11 +22,8 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
2522
type: 'profile',
2623
profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
2724
})
28-
// We need to push the cached customization on startup explicitly
29-
await pushConfigUpdate(languageClient, {
30-
type: 'customization',
31-
customization: getSelectedCustomization(),
32-
})
25+
26+
await initializeCustomizations()
3327

3428
const provider = new AmazonQChatViewProvider(mynahUIPath)
3529

@@ -79,17 +73,10 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
7973

8074
disposables.push(
8175
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
82-
void pushConfigUpdate(languageClient, {
83-
type: 'profile',
84-
profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
85-
})
8676
await provider.refreshWebview()
8777
}),
8878
Commands.register('aws.amazonq.updateCustomizations', () => {
89-
void pushConfigUpdate(languageClient, {
90-
type: 'customization',
91-
customization: undefinedIfEmpty(getSelectedCustomization().arn),
92-
})
79+
pushCustomizationToServer(languageClient)
9380
}),
9481
globals.logOutputChannel.onDidChangeLogLevel((logLevel) => {
9582
getLogger('amazonqLsp').info(`Local log level changed to ${logLevel}, notifying LSP`)
@@ -98,46 +85,35 @@ export async function activate(languageClient: LanguageClient, encryptionKey: Bu
9885
})
9986
})
10087
)
101-
}
10288

103-
/**
104-
* Push a config value to the language server, effectively updating it with the
105-
* latest configuration from the client.
106-
*
107-
* The issue is we need to push certain configs to different places, since there are
108-
* different handlers for specific configs. So this determines the correct place to
109-
* push the given config.
110-
*/
111-
async function pushConfigUpdate(client: LanguageClient, config: QConfigs) {
112-
switch (config.type) {
113-
case 'profile':
114-
await client.sendRequest(updateConfigurationRequestType.method, {
115-
section: 'aws.q',
116-
settings: { profileArn: config.profileArn },
117-
})
118-
break
119-
case 'customization':
120-
client.sendNotification(DidChangeConfigurationNotification.type.method, {
121-
section: 'aws.q',
122-
settings: { customization: config.customization },
123-
})
124-
break
125-
case 'logLevel':
126-
client.sendNotification(DidChangeConfigurationNotification.type.method, {
127-
section: 'aws.logLevel',
128-
})
129-
break
89+
/**
90+
* Initialize customizations on extension startup
91+
*/
92+
async function initializeCustomizations() {
93+
/**
94+
* Even though this function is called "notify", it has a side effect that first restores the
95+
* cached customization. So for {@link getSelectedCustomization()} to work as expected, we must
96+
* call {@link notifyNewCustomizations} first.
97+
*
98+
* TODO: Separate restoring and notifying, or just rename the function to something better
99+
*/
100+
if (AuthUtil.instance.isIdcConnection() && AuthUtil.instance.isConnected()) {
101+
await notifyNewCustomizations()
102+
}
103+
104+
/**
105+
* HACK: We must explicitly push the customization here since restoring the customization from cache
106+
* does not currently trigger a push to server.
107+
*
108+
* TODO: Always push to server whenever restoring from cache.
109+
*/
110+
pushCustomizationToServer(languageClient)
111+
}
112+
113+
function pushCustomizationToServer(languageClient: LanguageClient) {
114+
void pushConfigUpdate(languageClient, {
115+
type: 'customization',
116+
customization: undefinedIfEmpty(getSelectedCustomization().arn),
117+
})
130118
}
131119
}
132-
type ProfileConfig = {
133-
type: 'profile'
134-
profileArn: string | undefined
135-
}
136-
type CustomizationConfig = {
137-
type: 'customization'
138-
customization: string | undefined
139-
}
140-
type LogLevelConfig = {
141-
type: 'logLevel'
142-
}
143-
type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig

packages/amazonq/src/lsp/client.ts

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ import { processUtils } from 'aws-core-vscode/shared'
5454
import { activate as activateChat } from './chat/activation'
5555
import { AmazonQResourcePaths } from './lspInstaller'
5656
import { auth2 } from 'aws-core-vscode/auth'
57-
import { ConfigSection, isValidConfigSection, toAmazonQLSPLogLevel } from './config'
57+
import { ConfigSection, isValidConfigSection, pushConfigUpdate, toAmazonQLSPLogLevel } from './config'
5858
import { telemetry } from 'aws-core-vscode/telemetry'
5959

6060
const localize = nls.loadMessageBundle()
@@ -179,15 +179,31 @@ export async function startLanguageServer(
179179

180180
await client.onReady()
181181

182-
// IMPORTANT: This sets up Auth and must be called before anything attempts to use it
183-
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
184-
// Ideally this would be part of AuthUtil.create() as it restores the existing Auth connection, but we have some
185-
// future work which will need to delay this call
186-
await AuthUtil.instance.restore()
182+
/**
183+
* We use the Flare Auth language server, and our Auth client depends on it.
184+
* Because of this we initialize our Auth client **immediately** after the language server is ready.
185+
* Doing this removes the chance of something else attempting to use the Auth client before it is ready.
186+
*/
187+
await initializeAuth(client)
187188

188189
await postStartLanguageServer(client, resourcePaths, toDispose)
189190

190191
return client
192+
193+
async function initializeAuth(client: LanguageClient) {
194+
AuthUtil.create(new auth2.LanguageClientAuth(client, clientId, encryptionKey))
195+
196+
/** All must be setup before {@link AuthUtil.restore} otherwise they may not trigger when expected */
197+
AuthUtil.instance.regionProfileManager.onDidChangeRegionProfile(async () => {
198+
void pushConfigUpdate(client, {
199+
type: 'profile',
200+
profileArn: AuthUtil.instance.regionProfileManager.activeRegionProfile?.arn,
201+
})
202+
})
203+
204+
// Try and restore a cached connection if exists
205+
await AuthUtil.instance.restore()
206+
}
191207
}
192208

193209
async function postStartLanguageServer(

packages/amazonq/src/lsp/config.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,11 @@
55
import * as vscode from 'vscode'
66
import { DevSettings, getServiceEnvVarConfig } from 'aws-core-vscode/shared'
77
import { LspConfig } from 'aws-core-vscode/amazonq'
8+
import { LanguageClient } from 'vscode-languageclient'
9+
import {
10+
DidChangeConfigurationNotification,
11+
updateConfigurationRequestType,
12+
} from '@aws/language-server-runtimes/protocol'
813

914
export interface ExtendedAmazonQLSPConfig extends LspConfig {
1015
ui?: string
@@ -54,3 +59,45 @@ export function getAmazonQLspConfig(): ExtendedAmazonQLSPConfig {
5459
export function toAmazonQLSPLogLevel(logLevel: vscode.LogLevel): LspLogLevel {
5560
return lspLogLevelMapping.get(logLevel) ?? 'info'
5661
}
62+
63+
/**
64+
* Request/Notify a config value to the language server, effectively updating it with the
65+
* latest configuration from the client.
66+
*
67+
* The issue is we need to push certain configs to different places, since there are
68+
* different handlers for specific configs. So this determines the correct place to
69+
* push the given config.
70+
*/
71+
export async function pushConfigUpdate(client: LanguageClient, config: QConfigs) {
72+
switch (config.type) {
73+
case 'profile':
74+
await client.sendRequest(updateConfigurationRequestType.method, {
75+
section: 'aws.q',
76+
settings: { profileArn: config.profileArn },
77+
})
78+
break
79+
case 'customization':
80+
client.sendNotification(DidChangeConfigurationNotification.type.method, {
81+
section: 'aws.q',
82+
settings: { customization: config.customization },
83+
})
84+
break
85+
case 'logLevel':
86+
client.sendNotification(DidChangeConfigurationNotification.type.method, {
87+
section: 'aws.logLevel',
88+
})
89+
break
90+
}
91+
}
92+
type ProfileConfig = {
93+
type: 'profile'
94+
profileArn: string | undefined
95+
}
96+
type CustomizationConfig = {
97+
type: 'customization'
98+
customization: string | undefined
99+
}
100+
type LogLevelConfig = {
101+
type: 'logLevel'
102+
}
103+
type QConfigs = ProfileConfig | CustomizationConfig | LogLevelConfig

packages/core/src/codewhisperer/activation.ts

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ import { AuthUtil } from './util/authUtil'
7272
import { ImportAdderProvider } from './service/importAdderProvider'
7373
import { TelemetryHelper } from './util/telemetryHelper'
7474
import { openUrl } from '../shared/utilities/vsCodeUtils'
75-
import { notifyNewCustomizations, onProfileChangedListener } from './util/customizationUtil'
75+
import { onProfileChangedListener } from './util/customizationUtil'
7676
import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands'
7777
import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider'
7878
import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider'
@@ -354,9 +354,6 @@ export async function activate(context: ExtContext): Promise<void> {
354354
{ emit: false, functionId: { name: 'activateCwCore' } }
355355
)
356356

357-
if (AuthUtil.instance.isIdcConnection() && AuthUtil.instance.isConnected()) {
358-
await notifyNewCustomizations()
359-
}
360357
if (AuthUtil.instance.isBuilderIdConnection()) {
361358
await CodeScansState.instance.setScansEnabled(false)
362359
}

packages/core/src/codewhisperer/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ export {
107107
baseCustomization,
108108
onProfileChangedListener,
109109
CustomizationProvider,
110+
notifyNewCustomizations,
110111
} from './util/customizationUtil'
111112
export { Container } from './service/serviceContainer'
112113
export * from './util/gitUtil'

0 commit comments

Comments
 (0)