diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.test.ts index 7d6e9258fb..5ad0522a0b 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.test.ts @@ -6,12 +6,21 @@ import { CancellationToken, CredentialsType, InitializeParams, + PartialInitializeResult, Server, UpdateConfigurationParams, } from '@aws/language-server-runtimes/server-interface' -import { AmazonQServiceServerFactory } from './amazonQServer' +import { + AMAZON_Q_SERVICE_SERVER_IAM_NAME, + AMAZON_Q_SERVICE_SERVER_TOKEN_NAME, + AmazonQServiceServerFactory, + AmazonQServiceServerIAM, + AmazonQServiceServerToken, +} from './amazonQServer' import { BaseAmazonQServiceManager } from './amazonQServiceManager/BaseAmazonQServiceManager' +const TEST_SERVER_NAME = 'Test Amazon Q Server' + describe('AmazonQServiceServer', () => { let features: TestFeatures let server: Server @@ -23,7 +32,7 @@ describe('AmazonQServiceServer', () => { initBaseTestServiceManagerSpy = sinon.spy(initBaseTestServiceManager) TestAmazonQServiceManager.resetInstance() - server = AmazonQServiceServerFactory(() => initBaseTestServiceManagerSpy(features)) + server = AmazonQServiceServerFactory(() => initBaseTestServiceManagerSpy(features), TEST_SERVER_NAME) }) afterEach(() => { @@ -43,6 +52,55 @@ describe('AmazonQServiceServer', () => { sinon.assert.calledOnce(initBaseTestServiceManagerSpy) }) + it('declares serverInfo so the runtime can deliver notifications to the client', async () => { + server(features) + + // Invoke the registered initializer directly: doSendInitializeRequest returns void, so the + // result is only reachable through the handler the server registered. + const initializer = features.lsp.addInitializer.args[0]?.[0] + const result = (await initializer({} as InitializeParams, {} as CancellationToken)) as PartialInitializeResult + + // The runtime only builds a notification router for servers that declare serverInfo, and + // notification.showNotification() is a silent no-op without one. + expect(result.serverInfo?.name).to.equal(TEST_SERVER_NAME) + }) + + it('gives the IAM and token servers distinct serverInfo names', async () => { + // Regression guard, asserted against the real exported servers rather than the constants, so + // it also catches the same name being passed to both factory calls. + // + // Runtimes such as agent-standalone register BOTH of these servers, and the runtime rejects + // initialize with `Duplicate servers defined` when two servers report the same name -- which + // fails the entire language server, not just the duplicate. A shared name here made every such + // runtime fall back to whatever server the client had bundled, visible only as a client-side + // warning, so Q kept working while silently running a different server. + const names: (string | undefined)[] = [] + + for (const qServer of [AmazonQServiceServerIAM, AmazonQServiceServerToken]) { + const serverFeatures = new TestFeatures() + try { + // The service managers refuse to initialize before the LSP connection has, so the + // client params have to be in place before the initializer runs. + serverFeatures.setClientParams({} as InitializeParams) + qServer(serverFeatures) + + const initializer = serverFeatures.lsp.addInitializer.args[0]?.[0] + const result = (await initializer( + {} as InitializeParams, + {} as CancellationToken + )) as PartialInitializeResult + + names.push(result.serverInfo?.name) + } finally { + serverFeatures.dispose() + TestAmazonQServiceManager.resetInstance() + } + } + + expect(names).to.deep.equal([AMAZON_Q_SERVICE_SERVER_IAM_NAME, AMAZON_Q_SERVICE_SERVER_TOKEN_NAME]) + expect(new Set(names).size, `server names must be unique: ${names.join(', ')}`).to.equal(names.length) + }) + it('hooks handleDidChangeConfiguration to didChangeConfiguration and onInitialized handlers', async () => { const handleDidChangeConfigurationSpy = sinon.spy( BaseAmazonQServiceManager.prototype, @@ -111,7 +169,7 @@ describe('AmazonQServiceServer', () => { throw new Error('Service manager initialization failed') } - const errorServer = AmazonQServiceServerFactory(errorFactory) + const errorServer = AmazonQServiceServerFactory(errorFactory, TEST_SERVER_NAME) expect(() => { errorServer(features) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts index e9c83b0a95..8ec4e19372 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts @@ -17,9 +17,16 @@ import { const LOGGING_PREFIX = '[AMAZON Q SERVER]: ' +/** + * @param serverName Value for the server's `serverInfo.name`. Required, and required to be unique + * across every server registered in the same runtime: the runtime rejects `initialize` outright with + * `Duplicate servers defined` if two servers report the same name, which takes down the whole language + * server rather than the offending one. This factory is instantiated more than once per runtime (IAM + * and token), so a shared constant here breaks every bundle that registers both. + */ export const AmazonQServiceServerFactory = - (serviceManager: (features: QServiceManagerFeatures) => AmazonQBaseServiceManager): Server => - ({ credentialsProvider, lsp, workspace, logging, runtime, sdkInitializator }) => { + (serviceManager: (features: QServiceManagerFeatures) => AmazonQBaseServiceManager, serverName: string): Server => + ({ credentialsProvider, lsp, workspace, logging, runtime, sdkInitializator, notification }) => { let amazonQServiceManager: AmazonQBaseServiceManager const log = (message: string) => { @@ -39,6 +46,11 @@ export const AmazonQServiceServerFactory = logging, runtime, sdkInitializator, + // Required for the service manager to be able to surface anything to the client. It is + // optional on QServiceManagerFeatures so that existing constructions (including test + // fixtures) keep compiling, which means omitting it here does not fail the build -- it + // just silently disables client-facing reporting. + notification, }) /* @@ -58,6 +70,14 @@ export const AmazonQServiceServerFactory = return { capabilities: {}, awsServerCapabilities: {}, + // Required for anything to reach the client through the notification feature. The + // runtime only builds a notification router when a server declares serverInfo, and + // showNotification is a silent no-op without one. Not exposed to clients; it is used + // internally to route notification followups back to the originating server, so the + // name must stay stable. + serverInfo: { + name: serverName, + }, } }) @@ -110,5 +130,18 @@ export const AmazonQServiceServerFactory = return () => {} } -export const AmazonQServiceServerIAM = AmazonQServiceServerFactory(initBaseIAMServiceManager) -export const AmazonQServiceServerToken = AmazonQServiceServerFactory(initBaseTokenServiceManager) +// Must stay distinct from each other, and from every other server registered in the same runtime, or +// the runtime fails initialize with `Duplicate servers defined`. Must also stay stable over time: they +// are encoded into the id of every notification the client echoes back, so renaming strands followups +// for notifications already on screen. Exported so the uniqueness can be asserted in tests. +export const AMAZON_Q_SERVICE_SERVER_IAM_NAME = 'AWS Language Server for Amazon Q Developer (IAM)' +export const AMAZON_Q_SERVICE_SERVER_TOKEN_NAME = 'AWS Language Server for Amazon Q Developer (Token)' + +export const AmazonQServiceServerIAM = AmazonQServiceServerFactory( + initBaseIAMServiceManager, + AMAZON_Q_SERVICE_SERVER_IAM_NAME +) +export const AmazonQServiceServerToken = AmazonQServiceServerFactory( + initBaseTokenServiceManager, + AMAZON_Q_SERVICE_SERVER_TOKEN_NAME +) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts index ace2ae9943..c93552a7d7 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts @@ -238,6 +238,15 @@ describe('AmazonQTokenServiceManager', () => { assert(codewhispererServiceStub.generateSuggestions.calledOnce) }) + it('gives the streaming client an access-blocked observer', async () => { + const streamingClient = amazonQTokenServiceManager.getStreamingClient() + + // Chat runs through the streaming client, so without this it is the one surface where a + // blocked identity shows up and the one place nothing observes it. The token client is + // a stub in this harness, so only the streaming side is asserted here. + assert.strictEqual(typeof streamingClient.onAccessBlocked, 'function') + }) + it('should initialize service with region set by client', async () => { features.setClientParams({ processId: 0, diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts index 47b43d0d81..608d4ae2ed 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts @@ -76,6 +76,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< private region?: string private endpoint?: string private regionChangeListeners: Array<(region: string) => void> = [] + private cachedQDevAccessBlockedNotifier?: (error: unknown) => void /** * Internal state of Service connection, based on status of bearer token and Amazon Q Developer profile selection. @@ -525,12 +526,34 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< return this.cachedStreamingClient } + /** + * One notifier shared by the token and streaming clients, so a blocked identity produces a single + * notification regardless of which client observes it first (the notifier dedupes per instance). + * + * Deliberately scoped to the current service generation and cleared by + * {@link resetCodewhispererService}: signing out and back in with another blocked identity must + * notify again, which a manager-lifetime notifier would suppress. + */ + private getQDevAccessBlockedNotifier(): ((error: unknown) => void) | undefined { + if (!this.features.notification) { + return undefined + } + + this.cachedQDevAccessBlockedNotifier ??= createQDevAccessBlockedNotifier( + this.features.notification, + this.features.logging + ) + + return this.cachedQDevAccessBlockedNotifier + } + private resetCodewhispererService() { this.logging.log('Resetting Q-only services') this.cachedCodewhispererService?.abortInflightRequests() this.cachedCodewhispererService = undefined this.cachedStreamingClient?.abortInflightRequests() this.cachedStreamingClient = undefined + this.cachedQDevAccessBlockedNotifier = undefined this.activeIdcProfile = undefined this.region = undefined this.endpoint = undefined @@ -597,9 +620,7 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< service.customizationArn = this.configurationCache.getProperty('customizationArn') service.profileArn = this.activeIdcProfile?.arn - if (this.features.notification) { - service.onAccessBlocked = createQDevAccessBlockedNotifier(this.features.notification, this.features.logging) - } + service.onAccessBlocked = this.getQDevAccessBlockedNotifier() service.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty( 'shareCodeWhispererContentWithAWS' ) @@ -638,6 +659,9 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< endpoint, this.getCustomUserAgent() ) + // Chat runs through the streaming client, so it needs the same observer as the token client -- + // otherwise a blocked identity goes unnoticed unless some other client happens to be called. + streamingClient.onAccessBlocked = this.getQDevAccessBlockedNotifier() streamingClient.profileArn = this.activeIdcProfile?.arn streamingClient.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty( 'shareCodeWhispererContentWithAWS' diff --git a/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.test.ts b/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.test.ts index 1914680154..858fe3d8cc 100644 --- a/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.test.ts @@ -7,7 +7,7 @@ import { MessageType } from '@aws/language-server-runtimes/protocol' import { Logging, Notification } from '@aws/language-server-runtimes/server-interface' import * as assert from 'assert' import * as sinon from 'sinon' -import { createQDevAccessBlockedNotifier } from './qDevAccessBlockedNotifier' +import { createQDevAccessBlockedNotifier, Q_DEV_ACCESS_BLOCKED_NOTIFICATION_ID } from './qDevAccessBlockedNotifier' describe('createQDevAccessBlockedNotifier', function () { let showNotification: sinon.SinonStub @@ -40,6 +40,8 @@ describe('createQDevAccessBlockedNotifier', function () { const params = showNotification.firstCall.args[0] assert.strictEqual(params.type, MessageType.Error) assert.strictEqual(params.content.text, serviceMessage) + // Clients key off the id rather than the message text, which is service-owned copy. + assert.strictEqual(params.id, Q_DEV_ACCESS_BLOCKED_NOTIFICATION_ID) }) it('notifies at most once, since every request from a blocked identity fails', function () { diff --git a/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.ts b/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.ts index 7f61207fb5..6da7b4b38a 100644 --- a/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.ts +++ b/server/aws-lsp-codewhisperer/src/shared/qDevAccessBlockedNotifier.ts @@ -16,6 +16,12 @@ const FALLBACK_TEXT = 'Amazon Q Developer is not available for this account.' const TITLE = 'Amazon Q Developer' +/** + * Stable identifier so clients can recognise this notification without inspecting its text. Clients + * must not match on the message: it is the service's own copy and is expected to change. + */ +export const Q_DEV_ACCESS_BLOCKED_NOTIFICATION_ID = 'qDevPluginAccessBlocked' + /** * Builds the reaction to RTS blocking Q Developer plugin access for the current identity, for use as * {@link CodeWhispererServiceToken.onAccessBlocked}. @@ -56,6 +62,7 @@ export function createQDevAccessBlockedNotifier( try { logging.warn(`Q Developer plugin access is blocked for this identity: ${text}`) notification.showNotification({ + id: Q_DEV_ACCESS_BLOCKED_NOTIFICATION_ID, type: MessageType.Error, content: { title: TITLE, diff --git a/server/aws-lsp-codewhisperer/src/shared/streamingClientService.test.ts b/server/aws-lsp-codewhisperer/src/shared/streamingClientService.test.ts index 5b1f08d851..382052cce7 100644 --- a/server/aws-lsp-codewhisperer/src/shared/streamingClientService.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/streamingClientService.test.ts @@ -142,6 +142,32 @@ describe('StreamingClientServiceToken', () => { expect(streamingClientServiceDefault['shareCodeWhispererContentWithAWS']).to.be.undefined }) + describe('access-blocked observer', () => { + it('registers exactly one named middleware on the initialize step', () => { + // Guards the wiring rather than the callback: the harness stubs + // CodeWhispererStreaming.prototype.sendMessage, which bypasses the middleware stack, so a + // behavioural test here would pass without the middleware existing at all. + // + // The name matters. Without it a second registration stacks another observer instead of + // replacing the first, which would report the same block twice. + const registered = streamingClientService.client.middlewareStack + .identify() + .filter(entry => entry.includes('detectQDevPluginAccessBlocked')) + + expect(registered).to.have.lengthOf(1) + expect(registered[0]).to.contain('initialize') + }) + + it('exposes a settable observer for the service manager to assign', () => { + // Assigned after construction, so the middleware has to read it at call time. If this + // became readonly or were dropped, chat-time blocks would go unobserved. + const observer = () => {} + streamingClientService.onAccessBlocked = observer + + expect(streamingClientService.onAccessBlocked).to.equal(observer) + }) + }) + describe('generateAssistantResponse', () => { const MOCKED_GENERATE_RESPONSE_REQUEST = { conversationState: { diff --git a/server/aws-lsp-codewhisperer/src/shared/streamingClientService.ts b/server/aws-lsp-codewhisperer/src/shared/streamingClientService.ts index 373c2b54c4..dbb6aa98c3 100644 --- a/server/aws-lsp-codewhisperer/src/shared/streamingClientService.ts +++ b/server/aws-lsp-codewhisperer/src/shared/streamingClientService.ts @@ -18,7 +18,7 @@ import { Logging, IamCredentials, } from '@aws/language-server-runtimes/server-interface' -import { getBearerTokenFromProvider, isUsageLimitError } from './utils' +import { getBearerTokenFromProvider, isUsageLimitError, isQDevPluginAccessBlockedError } from './utils' import { CLIENT_TIMEOUT_MS, MAX_REQUEST_ATTEMPTS } from '../language-server/agenticChat/constants/constants' import { AmazonQUsageLimitError } from './amazonQServiceManager/errors' @@ -72,6 +72,12 @@ export abstract class StreamingClientServiceBase { } export class StreamingClientServiceToken extends StreamingClientServiceBase { + /** + * Observes Q Developer plugin access-blocked rejections. Assigned by the service manager after + * construction, so the middleware below reads it at call time rather than capturing it. + */ + public onAccessBlocked?: (error: unknown) => void + client: CodeWhispererStreaming public profileArn?: string private retryClassifier: QRetryClassifier @@ -132,6 +138,37 @@ export class StreamingClientServiceToken extends StreamingClientServiceBase { step: 'build', } ) + + // Observe access-blocked rejections without altering behaviour: the error is always rethrown + // so callers see it exactly as before. Registered on the outermost (initialize) step so it + // fires once per operation, after the SDK's retries are exhausted, rather than once per + // attempt. Chat runs through this client, so without it a blocked identity goes unnoticed + // unless some other client happens to be called. + this.client.middlewareStack.add( + next => async args => { + try { + return await next(args) + } catch (e) { + if (isQDevPluginAccessBlockedError(e)) { + try { + this.onAccessBlocked?.(e) + } catch (observerError) { + logging.debug( + `onAccessBlocked observer threw, ignoring: ${(observerError as Error)?.message}` + ) + } + } + throw e + } + }, + { + step: 'initialize', + // Named so a second registration replaces this one rather than stacking another + // observer, and so the middleware is identifiable in SDK stack introspection. Matches + // the token client's registration. + name: 'detectQDevPluginAccessBlocked', + } + ) } public async sendMessage(