Skip to content

Commit 0fa84d8

Browse files
ctlai95ctlai95
committed
Merge branch 'feature/hybridChat' into lsp-open-diff
2 parents 1910cbf + 7bb66e1 commit 0fa84d8

File tree

11 files changed

+307
-106
lines changed

11 files changed

+307
-106
lines changed

‎packages/amazonq/.changes/next-release/Feature-2a696d00-e8c8-44a4-ab2c-2204b4d8e31d.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Feature",
3+
"description": "Support selecting customizations across all Q profiles with automatic profile switching for enterprise users"
4+
}

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

Lines changed: 57 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,14 +41,17 @@ import {
4141
ContextCommandParams,
4242
openFileDiffNotificationType,
4343
OpenFileDiffParams,
44+
LINK_CLICK_NOTIFICATION_METHOD,
45+
LinkClickParams,
46+
INFO_LINK_CLICK_NOTIFICATION_METHOD,
4447
} from '@aws/language-server-runtimes/protocol'
4548
import { v4 as uuidv4 } from 'uuid'
4649
import * as vscode from 'vscode'
4750
import { Disposable, LanguageClient, Position, TextDocumentIdentifier } from 'vscode-languageclient'
4851
import * as jose from 'jose'
4952
import { AmazonQChatViewProvider } from './webviewProvider'
5053
import { AuthUtil } from 'aws-core-vscode/codewhisperer'
51-
import { AmazonQPromptSettings, messages } from 'aws-core-vscode/shared'
54+
import { AmazonQPromptSettings, messages, openUrl } from 'aws-core-vscode/shared'
5255
import { DefaultAmazonQAppInitContext, messageDispatcher, EditorContentController } from 'aws-core-vscode/amazonq'
5356

5457
export function registerLanguageServerEventListener(languageClient: LanguageClient, provider: AmazonQChatViewProvider) {
@@ -174,11 +177,31 @@ export function registerMessageListeners(
174177
}
175178
break
176179
}
180+
case INFO_LINK_CLICK_NOTIFICATION_METHOD:
181+
case LINK_CLICK_NOTIFICATION_METHOD: {
182+
const linkParams = message.params as LinkClickParams
183+
void openUrl(vscode.Uri.parse(linkParams.link))
184+
break
185+
}
177186
case chatRequestType.method: {
178-
const chatParams = { ...message.params } as ChatParams
187+
const chatParams: ChatParams = { ...message.params }
179188
const partialResultToken = uuidv4()
180-
const chatDisposable = languageClient.onProgress(chatRequestType, partialResultToken, (partialResult) =>
181-
handlePartialResult<ChatResult>(partialResult, encryptionKey, provider, chatParams.tabId)
189+
let lastPartialResult: ChatResult | undefined
190+
const chatDisposable = languageClient.onProgress(
191+
chatRequestType,
192+
partialResultToken,
193+
(partialResult) => {
194+
// Store the latest partial result
195+
if (typeof partialResult === 'string' && encryptionKey) {
196+
void decodeRequest<ChatResult>(partialResult, encryptionKey).then(
197+
(decoded) => (lastPartialResult = decoded)
198+
)
199+
} else {
200+
lastPartialResult = partialResult as ChatResult
201+
}
202+
203+
void handlePartialResult<ChatResult>(partialResult, encryptionKey, provider, chatParams.tabId)
204+
}
182205
)
183206

184207
const editor =
@@ -190,17 +213,36 @@ export function registerMessageListeners(
190213
}
191214

192215
const chatRequest = await encryptRequest<ChatParams>(chatParams, encryptionKey)
193-
const chatResult = (await languageClient.sendRequest(chatRequestType.method, {
194-
...chatRequest,
195-
partialResultToken,
196-
})) as string | ChatResult
197-
void handleCompleteResult<ChatResult>(
198-
chatResult,
199-
encryptionKey,
200-
provider,
201-
chatParams.tabId,
202-
chatDisposable
203-
)
216+
try {
217+
const chatResult = await languageClient.sendRequest<string | ChatResult>(chatRequestType.method, {
218+
...chatRequest,
219+
partialResultToken,
220+
})
221+
await handleCompleteResult<ChatResult>(
222+
chatResult,
223+
encryptionKey,
224+
provider,
225+
chatParams.tabId,
226+
chatDisposable
227+
)
228+
} catch (e) {
229+
languageClient.info(`Error occurred during chat request: ${e}`)
230+
// Use the last partial result if available, append error message
231+
const errorResult: ChatResult = {
232+
...lastPartialResult,
233+
body: lastPartialResult?.body
234+
? `${lastPartialResult.body}\n\n ❌ Error: Request failed to complete`
235+
: '❌ An error occurred while processing your request',
236+
}
237+
238+
await handleCompleteResult<ChatResult>(
239+
errorResult,
240+
encryptionKey,
241+
provider,
242+
chatParams.tabId,
243+
chatDisposable
244+
)
245+
}
204246
break
205247
}
206248
case quickActionRequestType.method: {

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ export async function startLanguageServer(
127127
},
128128
awsClientCapabilities: {
129129
q: {
130-
developerProfiles: true,
130+
developerProfiles: false,
131131
},
132132
window: {
133133
notifications: true,

‎packages/amazonq/test/unit/codewhisperer/region/regionProfileManager.test.ts‎

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ describe('RegionProfileManager', function () {
6565
const mockClient = {
6666
listAvailableProfiles: listProfilesStub,
6767
}
68-
const createClientStub = sinon.stub(sut, 'createQClient').resolves(mockClient)
68+
const createClientStub = sinon.stub(sut, '_createQClient').resolves(mockClient)
6969

7070
const r = await sut.listRegionProfile()
7171

@@ -234,13 +234,65 @@ describe('RegionProfileManager', function () {
234234
})
235235

236236
describe('createQClient', function () {
237+
it(`should configure the endpoint and region from a profile`, async function () {
238+
await setupConnection('idc')
239+
240+
const iadClient = await sut.createQClient({
241+
name: 'foo',
242+
region: 'us-east-1',
243+
arn: 'arn',
244+
description: 'description',
245+
})
246+
247+
assert.deepStrictEqual(iadClient.config.region, 'us-east-1')
248+
assert.deepStrictEqual(iadClient.endpoint.href, 'https://q.us-east-1.amazonaws.com/')
249+
250+
const fraClient = await sut.createQClient({
251+
name: 'bar',
252+
region: 'eu-central-1',
253+
arn: 'arn',
254+
description: 'description',
255+
})
256+
257+
assert.deepStrictEqual(fraClient.config.region, 'eu-central-1')
258+
assert.deepStrictEqual(fraClient.endpoint.href, 'https://q.eu-central-1.amazonaws.com/')
259+
})
260+
261+
it(`should throw if the region is not supported or recognizable by Q`, async function () {
262+
await setupConnection('idc')
263+
264+
await assert.rejects(
265+
async () => {
266+
await sut.createQClient({
267+
name: 'foo',
268+
region: 'ap-east-1',
269+
arn: 'arn',
270+
description: 'description',
271+
})
272+
},
273+
{ message: /trying to initiatize Q client with unrecognizable region/ }
274+
)
275+
276+
await assert.rejects(
277+
async () => {
278+
await sut.createQClient({
279+
name: 'foo',
280+
region: 'unknown-somewhere',
281+
arn: 'arn',
282+
description: 'description',
283+
})
284+
},
285+
{ message: /trying to initiatize Q client with unrecognizable region/ }
286+
)
287+
})
288+
237289
it(`should configure the endpoint and region correspondingly`, async function () {
238290
await setupConnection('idc')
239291
await sut.switchRegionProfile(profileFoo, 'user')
240292
assert.deepStrictEqual(sut.activeRegionProfile, profileFoo)
241293
const conn = authUtil.conn as SsoConnection
242294

243-
const client = await sut.createQClient('eu-central-1', 'https://amazon.com/', conn)
295+
const client = await sut._createQClient('eu-central-1', 'https://amazon.com/', conn)
244296

245297
assert.deepStrictEqual(client.config.region, 'eu-central-1')
246298
assert.deepStrictEqual(client.endpoint.href, 'https://amazon.com/')

‎packages/core/src/codewhisperer/activation.ts‎

Lines changed: 2 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -72,12 +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 {
76-
getAvailableCustomizationsList,
77-
getSelectedCustomization,
78-
notifyNewCustomizations,
79-
switchToBaseCustomizationAndNotify,
80-
} from './util/customizationUtil'
75+
import { notifyNewCustomizations, onProfileChangedListener } from './util/customizationUtil'
8176
import { CodeWhispererCommandBackend, CodeWhispererCommandDeclarations } from './commands/gettingStartedPageCommands'
8277
import { SecurityIssueHoverProvider } from './service/securityIssueHoverProvider'
8378
import { SecurityIssueCodeActionProvider } from './service/securityIssueCodeActionProvider'
@@ -343,26 +338,7 @@ export async function activate(context: ExtContext): Promise<void> {
343338
SecurityIssueCodeActionProvider.instance
344339
),
345340
vscode.commands.registerCommand('aws.amazonq.openEditorAtRange', openEditorAtRange),
346-
auth.regionProfileManager.onDidChangeRegionProfile(() => {
347-
// Validate user still has access to the selected customization.
348-
const selectedCustomization = getSelectedCustomization()
349-
// No need to validate base customization which has empty arn.
350-
if (selectedCustomization.arn.length > 0) {
351-
getAvailableCustomizationsList()
352-
.then(async (customizations) => {
353-
const r = customizations.find((it) => it.arn === selectedCustomization.arn)
354-
if (!r) {
355-
await switchToBaseCustomizationAndNotify()
356-
}
357-
})
358-
.catch((e) => {
359-
getLogger().error(
360-
`encounter error while validating selected customization on profile change: %s`,
361-
(e as Error).message
362-
)
363-
})
364-
}
365-
})
341+
auth.regionProfileManager.onDidChangeRegionProfile(onProfileChangedListener)
366342
)
367343

368344
// run the auth startup code with context for telemetry

‎packages/core/src/codewhisperer/client/codewhisperer.ts‎

Lines changed: 1 addition & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,17 @@ import { AWSError, Credentials, Service } from 'aws-sdk'
77
import globals from '../../shared/extensionGlobals'
88
import * as CodeWhispererClient from './codewhispererclient'
99
import * as CodeWhispererUserClient from './codewhispereruserclient'
10-
import { ListAvailableCustomizationsResponse, SendTelemetryEventRequest } from './codewhispereruserclient'
10+
import { SendTelemetryEventRequest } from './codewhispereruserclient'
1111
import { ServiceOptions } from '../../shared/awsClientBuilder'
1212
import { hasVendedIamCredentials } from '../../auth/auth'
1313
import { CodeWhispererSettings } from '../util/codewhispererSettings'
1414
import { PromiseResult } from 'aws-sdk/lib/request'
1515
import { AuthUtil } from '../util/authUtil'
1616
import { isSsoConnection } from '../../auth/connection'
17-
import { pageableToCollection } from '../../shared/utilities/collectionUtils'
1817
import apiConfig = require('./service-2.json')
1918
import userApiConfig = require('./user-service-2.json')
2019
import { session } from '../util/codeWhispererSession'
2120
import { getLogger } from '../../shared/logger/logger'
22-
import { indent } from '../../shared/utilities/textUtilities'
2321
import { getClientId, getOptOutPreference, getOperatingSystem } from '../../shared/telemetry/util'
2422
import { extensionVersion, getServiceEnvVarConfig } from '../../shared/vscode/env'
2523
import { DevSettings } from '../../shared/settings'
@@ -219,28 +217,6 @@ export class DefaultCodeWhispererClient {
219217
.promise()
220218
}
221219

222-
public async listAvailableCustomizations(): Promise<ListAvailableCustomizationsResponse[]> {
223-
const client = await this.createUserSdkClient()
224-
const profile = AuthUtil.instance.regionProfileManager.activeRegionProfile
225-
const requester = async (request: CodeWhispererUserClient.ListAvailableCustomizationsRequest) =>
226-
client.listAvailableCustomizations(request).promise()
227-
return pageableToCollection(requester, { profileArn: profile?.arn }, 'nextToken')
228-
.promise()
229-
.then((resps) => {
230-
let logStr = 'amazonq: listAvailableCustomizations API request:'
231-
for (const resp of resps) {
232-
const requestId = resp.$response.requestId
233-
logStr += `\n${indent('RequestID: ', 4)}${requestId},\n${indent('Customizations:', 4)}`
234-
for (const [index, c] of resp.customizations.entries()) {
235-
const entry = `${index.toString().padStart(2, '0')}: ${c.name?.trim()}`
236-
logStr += `\n${indent(entry, 8)}`
237-
}
238-
}
239-
getLogger().debug(logStr)
240-
return resps
241-
})
242-
}
243-
244220
public async sendTelemetryEvent(request: SendTelemetryEventRequest) {
245221
const requestWithCommonFields: SendTelemetryEventRequest = {
246222
...request,

‎packages/core/src/codewhisperer/index.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,13 @@ export * as diagnosticsProvider from './service/diagnosticsProvider'
9999
export * from './ui/codeWhispererNodes'
100100
export { SecurityScanError, SecurityScanTimedOutError } from '../codewhisperer/models/errors'
101101
export * as CodeWhispererConstants from '../codewhisperer/models/constants'
102-
export { getSelectedCustomization, setSelectedCustomization, baseCustomization } from './util/customizationUtil'
102+
export {
103+
getSelectedCustomization,
104+
setSelectedCustomization,
105+
baseCustomization,
106+
onProfileChangedListener,
107+
CustomizationProvider,
108+
} from './util/customizationUtil'
103109
export { Container } from './service/serviceContainer'
104110
export * from './util/gitUtil'
105111
export * from './ui/prompters'

‎packages/core/src/codewhisperer/region/regionProfileManager.ts‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -47,12 +47,17 @@ const endpoints = createConstantMap({
4747
* 'update' -> plugin auto select the profile on users' behalf as there is only 1 profile
4848
* 'reload' -> on plugin restart, plugin will try to reload previous selected profile
4949
*/
50-
export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload'
50+
export type ProfileSwitchIntent = 'user' | 'auth' | 'update' | 'reload' | 'customization'
51+
52+
export type ProfileChangedEvent = {
53+
profile: RegionProfile | undefined
54+
intent: ProfileSwitchIntent
55+
}
5156

5257
export class RegionProfileManager {
5358
private static logger = getLogger()
5459
private _activeRegionProfile: RegionProfile | undefined
55-
private _onDidChangeRegionProfile = new vscode.EventEmitter<RegionProfile | undefined>()
60+
private _onDidChangeRegionProfile = new vscode.EventEmitter<ProfileChangedEvent>()
5661
public readonly onDidChangeRegionProfile = this._onDidChangeRegionProfile.event
5762

5863
// Store the last API results (for UI propuse) so we don't need to call service again if doesn't require "latest" result
@@ -112,7 +117,7 @@ export class RegionProfileManager {
112117
}
113118
const availableProfiles: RegionProfile[] = []
114119
for (const [region, endpoint] of endpoints.entries()) {
115-
const client = await this.createQClient(region, endpoint, conn as SsoConnection)
120+
const client = await this._createQClient(region, endpoint, conn as SsoConnection)
116121
const requester = async (request: CodeWhispererUserClient.ListAvailableProfilesRequest) =>
117122
client.listAvailableProfiles(request).promise()
118123
const request: CodeWhispererUserClient.ListAvailableProfilesRequest = {}
@@ -162,7 +167,7 @@ export class RegionProfileManager {
162167
const ssoConn = this.connectionProvider() as SsoConnection
163168

164169
// only prompt to users when users switch from A profile to B profile
165-
if (this.activeRegionProfile !== undefined && regionProfile !== undefined) {
170+
if (source !== 'customization' && this.activeRegionProfile !== undefined && regionProfile !== undefined) {
166171
const response = await showConfirmationMessage({
167172
prompt: localize(
168173
'AWS.amazonq.profile.confirmation',
@@ -204,13 +209,16 @@ export class RegionProfileManager {
204209
})
205210
}
206211

207-
await this._switchRegionProfile(regionProfile)
212+
await this._switchRegionProfile(regionProfile, source)
208213
}
209214

210-
private async _switchRegionProfile(regionProfile: RegionProfile | undefined) {
215+
private async _switchRegionProfile(regionProfile: RegionProfile | undefined, source: ProfileSwitchIntent) {
211216
this._activeRegionProfile = regionProfile
212217

213-
this._onDidChangeRegionProfile.fire(regionProfile)
218+
this._onDidChangeRegionProfile.fire({
219+
profile: regionProfile,
220+
intent: source,
221+
})
214222
// dont show if it's a default (fallback)
215223
if (regionProfile && this.profiles.length > 1) {
216224
void vscode.window.showInformationMessage(`You are using the ${regionProfile.name} profile for Q.`).then()
@@ -343,7 +351,21 @@ export class RegionProfileManager {
343351
}
344352
}
345353

346-
async createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
354+
// TODO: Should maintain sdk client in a better way
355+
async createQClient(profile: RegionProfile): Promise<CodeWhispererUserClient> {
356+
const conn = this.connectionProvider()
357+
if (conn === undefined || !isSsoConnection(conn)) {
358+
throw new Error('No valid SSO connection')
359+
}
360+
const endpoint = endpoints.get(profile.region)
361+
if (!endpoint) {
362+
throw new Error(`trying to initiatize Q client with unrecognizable region ${profile.region}`)
363+
}
364+
return this._createQClient(profile.region, endpoint, conn)
365+
}
366+
367+
// Visible for testing only, do not use this directly, please use createQClient(profile)
368+
async _createQClient(region: string, endpoint: string, conn: SsoConnection): Promise<CodeWhispererUserClient> {
347369
const token = (await conn.getToken()).accessToken
348370
const serviceOption: ServiceOptions = {
349371
apiConfig: userApiConfig,

0 commit comments

Comments
 (0)