From 5f8bb0e3e0c00f3eb342731a88667c95a0206215 Mon Sep 17 00:00:00 2001 From: liramon1 Date: Fri, 25 Jul 2025 13:12:35 -0400 Subject: [PATCH 1/8] refactor: split codewhisperer service into multiple files (#1974) --- .../src/shared/codeWhispererService.test.ts | 398 ++++++++++++++++++ 1 file changed, 398 insertions(+) create mode 100644 server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts diff --git a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts new file mode 100644 index 0000000000..2c095b5dba --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts @@ -0,0 +1,398 @@ +/*! + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + CredentialsProvider, + CredentialsType, + Workspace, + Logging, + SDKInitializator, +} from '@aws/language-server-runtimes/server-interface' +import { ConfigurationOptions } from 'aws-sdk' +import * as sinon from 'sinon' +import * as assert from 'assert' +import { + CodeWhispererServiceBase, + GenerateSuggestionsRequest, + GenerateSuggestionsResponse, +} from './codeWhispererService/codeWhispererServiceBase' +import { CodeWhispererServiceToken } from './codeWhispererService/codeWhispererServiceToken' +import { CodeWhispererServiceIAM } from './codeWhispererService/codeWhispererServiceIAM' + +describe('CodeWhispererService', function () { + let sandbox: sinon.SinonSandbox + let mockCredentialsProvider: sinon.SinonStubbedInstance + let mockWorkspace: sinon.SinonStubbedInstance + let mockLogging: sinon.SinonStubbedInstance + let mockSDKInitializator: sinon.SinonStubbedInstance + + beforeEach(function () { + sandbox = sinon.createSandbox() + + mockCredentialsProvider = { + getCredentials: sandbox.stub(), + hasCredentials: sandbox.stub(), + refresh: sandbox.stub(), + } as any + + mockWorkspace = { + getWorkspaceFolder: sandbox.stub(), + getWorkspaceFolders: sandbox.stub(), + } as any + + mockLogging = { + debug: sandbox.stub(), + error: sandbox.stub(), + info: sandbox.stub(), + warn: sandbox.stub(), + log: sandbox.stub(), + } + + mockSDKInitializator = { + initialize: sandbox.stub(), + } as any + }) + + afterEach(function () { + sandbox.restore() + }) + + describe('CodeWhispererServiceBase', function () { + let service: CodeWhispererServiceBase + + beforeEach(function () { + // Create a concrete implementation for testing abstract class + class TestCodeWhispererService extends CodeWhispererServiceBase { + client: any = {} + + getCredentialsType(): CredentialsType { + return 'iam' + } + + // Add public getters for protected properties + get testCodeWhispererRegion() { + return this.codeWhispererRegion + } + + get testCodeWhispererEndpoint() { + return this.codeWhispererEndpoint + } + + async generateCompletionsAndEdits(): Promise { + return { + suggestions: [], + responseContext: { requestId: 'test', codewhispererSessionId: 'test' }, + } + } + + async generateSuggestions(): Promise { + return { + suggestions: [], + responseContext: { requestId: 'test', codewhispererSessionId: 'test' }, + } + } + + clearCachedSuggestions(): void {} + } + + service = new TestCodeWhispererService('us-east-1', 'https://codewhisperer.us-east-1.amazonaws.com') + }) + + describe('constructor', function () { + it('should initialize with region and endpoint', function () { + assert.strictEqual((service as any).testCodeWhispererRegion, 'us-east-1') + assert.strictEqual( + (service as any).testCodeWhispererEndpoint, + 'https://codewhisperer.us-east-1.amazonaws.com' + ) + }) + }) + + describe('request tracking', function () { + it('should track inflight requests', function () { + const mockRequest = { + abort: sandbox.stub(), + } as any + + service.trackRequest(mockRequest) + assert.strictEqual(service.inflightRequests.size, 1) + assert.strictEqual(service.inflightRequests.has(mockRequest), true) + }) + + it('should complete and remove tracked requests', function () { + const mockRequest = { + abort: sandbox.stub(), + } as any + + service.trackRequest(mockRequest) + service.completeRequest(mockRequest) + + assert.strictEqual(service.inflightRequests.size, 0) + assert.strictEqual(service.inflightRequests.has(mockRequest), false) + }) + + it('should abort all inflight requests', function () { + const mockRequest1 = { abort: sandbox.stub() } as any + const mockRequest2 = { abort: sandbox.stub() } as any + + service.trackRequest(mockRequest1) + service.trackRequest(mockRequest2) + + service.abortInflightRequests() + + assert.strictEqual(mockRequest1.abort.calledOnce, true) + assert.strictEqual(mockRequest2.abort.calledOnce, true) + assert.strictEqual(service.inflightRequests.size, 0) + }) + }) + + describe('updateClientConfig', function () { + it('should update client configuration', function () { + const mockClient = { + config: { + update: sandbox.stub(), + }, + // Add minimal required properties to satisfy the interface + createCodeScan: sandbox.stub(), + createCodeScanUploadUrl: sandbox.stub(), + createProfile: sandbox.stub(), + deleteProfile: sandbox.stub(), + generateCompletions: sandbox.stub(), + generateSuggestions: sandbox.stub(), + getCodeAnalysis: sandbox.stub(), + getCodeScan: sandbox.stub(), + listCodeAnalysisFindings: sandbox.stub(), + listCodeScans: sandbox.stub(), + listFeatureEvaluations: sandbox.stub(), + listProfiles: sandbox.stub(), + sendTelemetryEvent: sandbox.stub(), + startCodeAnalysis: sandbox.stub(), + stopCodeAnalysis: sandbox.stub(), + updateProfile: sandbox.stub(), + } as any + service.client = mockClient + + const options: ConfigurationOptions = { region: 'us-west-2' } + service.updateClientConfig(options) + + assert.strictEqual(mockClient.config.update.calledOnceWith(options), true) + }) + }) + + describe('generateItemId', function () { + it('should generate unique item IDs', function () { + const id1 = service.generateItemId() + const id2 = service.generateItemId() + + assert.strictEqual(typeof id1, 'string') + assert.strictEqual(typeof id2, 'string') + assert.notStrictEqual(id1, id2) + }) + }) + }) + + describe('CodeWhispererServiceIAM', function () { + let service: CodeWhispererServiceIAM + + beforeEach(function () { + // Mock the createCodeWhispererSigv4Client function to avoid real client creation + const mockClient = { + generateRecommendations: sandbox.stub().returns({ + promise: sandbox.stub().resolves({ + recommendations: [], + $response: { + requestId: 'test-request-id', + httpResponse: { + headers: { 'x-amzn-sessionid': 'test-session-id' }, + }, + }, + }), + }), + setupRequestListeners: sandbox.stub(), + config: { + update: sandbox.stub(), + }, + } + + // Mock the client creation + const createClientStub = sandbox.stub( + require('../client/sigv4/codewhisperer'), + 'createCodeWhispererSigv4Client' + ) + createClientStub.returns(mockClient) + + service = new CodeWhispererServiceIAM( + mockCredentialsProvider as any, + {} as any, // workspace parameter + mockLogging as any, + 'us-east-1', + 'https://codewhisperer.us-east-1.amazonaws.com', + mockSDKInitializator as any + ) + }) + + describe('getCredentialsType', function () { + it('should return iam credentials type', function () { + assert.strictEqual(service.getCredentialsType(), 'iam') + }) + }) + + describe('generateSuggestions', function () { + it('should call client.generateRecommendations and process response', async function () { + const mockRequest: GenerateSuggestionsRequest = { + fileContext: { + filename: 'test.js', + programmingLanguage: { languageName: 'javascript' }, + leftFileContent: 'const x = ', + rightFileContent: '', + }, + maxResults: 5, + } + + const result = await service.generateSuggestions(mockRequest) + + assert.strictEqual(Array.isArray(result.suggestions), true) + assert.strictEqual(typeof result.responseContext.requestId, 'string') + assert.strictEqual(typeof result.responseContext.codewhispererSessionId, 'string') + }) + + it('should add customizationArn to request if set', async function () { + service.customizationArn = 'test-arn' + + const mockRequest: GenerateSuggestionsRequest = { + fileContext: { + filename: 'test.js', + programmingLanguage: { languageName: 'javascript' }, + leftFileContent: 'const x = ', + rightFileContent: '', + }, + maxResults: 5, + } + + await service.generateSuggestions(mockRequest) + + // Verify that the client was called with the customizationArn + const clientCall = (service.client.generateRecommendations as sinon.SinonStub).getCall(0) + assert.strictEqual(clientCall.args[0].customizationArn, 'test-arn') + }) + }) + }) + + describe('CodeWhispererServiceToken', function () { + let service: CodeWhispererServiceToken + let mockClient: any + + beforeEach(function () { + // Mock the token client + mockClient = { + generateCompletions: sandbox.stub().returns({ + promise: sandbox.stub().resolves({ + completions: [ + { + content: 'console.log("hello");', + references: [], + }, + ], + $response: { + requestId: 'test-request-id', + httpResponse: { + headers: { 'x-amzn-sessionid': 'test-session-id' }, + }, + }, + }), + }), + config: { + update: sandbox.stub(), + }, + } + + // Mock the client creation + const createTokenClientStub = sandbox.stub( + require('../client/token/codewhisperer'), + 'createCodeWhispererTokenClient' + ) + createTokenClientStub.returns(mockClient) + + // Mock bearer credentials + mockCredentialsProvider.getCredentials.returns({ + token: 'mock-bearer-token', + }) + + service = new CodeWhispererServiceToken( + mockCredentialsProvider as any, + mockWorkspace as any, + mockLogging as any, + 'us-east-1', + 'https://codewhisperer.us-east-1.amazonaws.com', + mockSDKInitializator as any + ) + }) + + describe('getCredentialsType', function () { + it('should return bearer credentials type', function () { + assert.strictEqual(service.getCredentialsType(), 'bearer') + }) + }) + + describe('generateSuggestions', function () { + it('should call client.generateCompletions and process response', async function () { + const mockRequest: GenerateSuggestionsRequest = { + fileContext: { + filename: 'test.js', + programmingLanguage: { languageName: 'javascript' }, + leftFileContent: 'const x = ', + rightFileContent: '', + }, + maxResults: 5, + } + + const result = await service.generateSuggestions(mockRequest) + + assert.strictEqual(mockClient.generateCompletions.calledOnce, true) + assert.strictEqual(Array.isArray(result.suggestions), true) + assert.strictEqual(typeof result.responseContext.requestId, 'string') + assert.strictEqual(typeof result.responseContext.codewhispererSessionId, 'string') + }) + + it('should add customizationArn to request if set', async function () { + service.customizationArn = 'test-arn' + + const mockRequest: GenerateSuggestionsRequest = { + fileContext: { + filename: 'test.js', + programmingLanguage: { languageName: 'javascript' }, + leftFileContent: 'const x = ', + rightFileContent: '', + }, + maxResults: 5, + } + + await service.generateSuggestions(mockRequest) + + const clientCall = mockClient.generateCompletions.getCall(0) + assert.strictEqual(clientCall.args[0].customizationArn, 'test-arn') + }) + + it('should process profile ARN with withProfileArn method', async function () { + const mockRequest: GenerateSuggestionsRequest = { + fileContext: { + filename: 'test.js', + programmingLanguage: { languageName: 'javascript' }, + leftFileContent: 'const x = ', + rightFileContent: '', + }, + maxResults: 5, + } + + const withProfileArnStub = sandbox.stub(service, 'withProfileArn' as any) + withProfileArnStub.returns(mockRequest) + + await service.generateSuggestions(mockRequest) + + assert.strictEqual(withProfileArnStub.calledOnceWith(mockRequest), true) + }) + }) + }) +}) From fcaf3da3cb781c31a2e528c1895a0c075db65af1 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Thu, 24 Jul 2025 14:26:23 -0400 Subject: [PATCH 2/8] feat: add abstract token methods to CodeWhispererServiceBase --- .../src/shared/codeWhispererService.test.ts | 130 +++++++++++++++++- 1 file changed, 129 insertions(+), 1 deletion(-) diff --git a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts index 2c095b5dba..0b20f90ce5 100644 --- a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts @@ -10,7 +10,7 @@ import { Logging, SDKInitializator, } from '@aws/language-server-runtimes/server-interface' -import { ConfigurationOptions } from 'aws-sdk' +import { AWSError, ConfigurationOptions } from 'aws-sdk' import * as sinon from 'sinon' import * as assert from 'assert' import { @@ -20,6 +20,38 @@ import { } from './codeWhispererService/codeWhispererServiceBase' import { CodeWhispererServiceToken } from './codeWhispererService/codeWhispererServiceToken' import { CodeWhispererServiceIAM } from './codeWhispererService/codeWhispererServiceIAM' +import { PromiseResult } from 'aws-sdk/lib/request' +import { + CreateUploadUrlRequest, + CreateUploadUrlResponse, + StartTransformationRequest, + StartTransformationResponse, + StopTransformationRequest, + StopTransformationResponse, + GetTransformationRequest, + GetTransformationResponse, + GetTransformationPlanRequest, + GetTransformationPlanResponse, + StartCodeAnalysisRequest, + StartCodeAnalysisResponse, + GetCodeAnalysisRequest, + GetCodeAnalysisResponse, + ListCodeAnalysisFindingsRequest, + ListCodeAnalysisFindingsResponse, + ListAvailableCustomizationsRequest, + ListAvailableCustomizationsResponse, + ListAvailableProfilesRequest, + SendTelemetryEventRequest, + SendTelemetryEventResponse, + CreateWorkspaceRequest, + CreateWorkspaceResponse, + ListWorkspaceMetadataRequest, + ListWorkspaceMetadataResponse, + DeleteWorkspaceRequest, + DeleteWorkspaceResponse, + ListFeatureEvaluationsRequest, + ListFeatureEvaluationsResponse, +} from '../client/token/codewhispererbearertokenclient' describe('CodeWhispererService', function () { let sandbox: sinon.SinonSandbox @@ -95,6 +127,102 @@ describe('CodeWhispererService', function () { } clearCachedSuggestions(): void {} + + override codeModernizerCreateUploadUrl( + request: CreateUploadUrlRequest + ): Promise { + throw new Error('Method not implemented.') + } + + override codeModernizerStartCodeTransformation( + request: StartTransformationRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override codeModernizerStopCodeTransformation( + request: StopTransformationRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override codeModernizerGetCodeTransformation( + request: GetTransformationRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override codeModernizerGetCodeTransformationPlan( + request: GetTransformationPlanRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override createUploadUrl( + request: CreateUploadUrlRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override startCodeAnalysis( + request: StartCodeAnalysisRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override getCodeAnalysis( + request: GetCodeAnalysisRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override listCodeAnalysisFindings( + request: ListCodeAnalysisFindingsRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override listAvailableCustomizations( + request: ListAvailableCustomizationsRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override listAvailableProfiles( + request: ListAvailableProfilesRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override sendTelemetryEvent( + request: SendTelemetryEventRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override createWorkspace( + request: CreateWorkspaceRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override listWorkspaceMetadata( + request: ListWorkspaceMetadataRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override deleteWorkspace( + request: DeleteWorkspaceRequest + ): Promise> { + throw new Error('Method not implemented.') + } + + override listFeatureEvaluations( + request: ListFeatureEvaluationsRequest + ): Promise> { + throw new Error('Method not implemented.') + } } service = new TestCodeWhispererService('us-east-1', 'https://codewhisperer.us-east-1.amazonaws.com') From c83c0b7e0b572b3fe094392eed91fc9ea9822435 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 25 Jul 2025 14:23:37 -0400 Subject: [PATCH 3/8] fix: remove unnecessary file --- .../src/shared/codeWhispererService.test.ts | 526 ------------------ 1 file changed, 526 deletions(-) delete mode 100644 server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts diff --git a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts deleted file mode 100644 index 0b20f90ce5..0000000000 --- a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService.test.ts +++ /dev/null @@ -1,526 +0,0 @@ -/*! - * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -import { - CredentialsProvider, - CredentialsType, - Workspace, - Logging, - SDKInitializator, -} from '@aws/language-server-runtimes/server-interface' -import { AWSError, ConfigurationOptions } from 'aws-sdk' -import * as sinon from 'sinon' -import * as assert from 'assert' -import { - CodeWhispererServiceBase, - GenerateSuggestionsRequest, - GenerateSuggestionsResponse, -} from './codeWhispererService/codeWhispererServiceBase' -import { CodeWhispererServiceToken } from './codeWhispererService/codeWhispererServiceToken' -import { CodeWhispererServiceIAM } from './codeWhispererService/codeWhispererServiceIAM' -import { PromiseResult } from 'aws-sdk/lib/request' -import { - CreateUploadUrlRequest, - CreateUploadUrlResponse, - StartTransformationRequest, - StartTransformationResponse, - StopTransformationRequest, - StopTransformationResponse, - GetTransformationRequest, - GetTransformationResponse, - GetTransformationPlanRequest, - GetTransformationPlanResponse, - StartCodeAnalysisRequest, - StartCodeAnalysisResponse, - GetCodeAnalysisRequest, - GetCodeAnalysisResponse, - ListCodeAnalysisFindingsRequest, - ListCodeAnalysisFindingsResponse, - ListAvailableCustomizationsRequest, - ListAvailableCustomizationsResponse, - ListAvailableProfilesRequest, - SendTelemetryEventRequest, - SendTelemetryEventResponse, - CreateWorkspaceRequest, - CreateWorkspaceResponse, - ListWorkspaceMetadataRequest, - ListWorkspaceMetadataResponse, - DeleteWorkspaceRequest, - DeleteWorkspaceResponse, - ListFeatureEvaluationsRequest, - ListFeatureEvaluationsResponse, -} from '../client/token/codewhispererbearertokenclient' - -describe('CodeWhispererService', function () { - let sandbox: sinon.SinonSandbox - let mockCredentialsProvider: sinon.SinonStubbedInstance - let mockWorkspace: sinon.SinonStubbedInstance - let mockLogging: sinon.SinonStubbedInstance - let mockSDKInitializator: sinon.SinonStubbedInstance - - beforeEach(function () { - sandbox = sinon.createSandbox() - - mockCredentialsProvider = { - getCredentials: sandbox.stub(), - hasCredentials: sandbox.stub(), - refresh: sandbox.stub(), - } as any - - mockWorkspace = { - getWorkspaceFolder: sandbox.stub(), - getWorkspaceFolders: sandbox.stub(), - } as any - - mockLogging = { - debug: sandbox.stub(), - error: sandbox.stub(), - info: sandbox.stub(), - warn: sandbox.stub(), - log: sandbox.stub(), - } - - mockSDKInitializator = { - initialize: sandbox.stub(), - } as any - }) - - afterEach(function () { - sandbox.restore() - }) - - describe('CodeWhispererServiceBase', function () { - let service: CodeWhispererServiceBase - - beforeEach(function () { - // Create a concrete implementation for testing abstract class - class TestCodeWhispererService extends CodeWhispererServiceBase { - client: any = {} - - getCredentialsType(): CredentialsType { - return 'iam' - } - - // Add public getters for protected properties - get testCodeWhispererRegion() { - return this.codeWhispererRegion - } - - get testCodeWhispererEndpoint() { - return this.codeWhispererEndpoint - } - - async generateCompletionsAndEdits(): Promise { - return { - suggestions: [], - responseContext: { requestId: 'test', codewhispererSessionId: 'test' }, - } - } - - async generateSuggestions(): Promise { - return { - suggestions: [], - responseContext: { requestId: 'test', codewhispererSessionId: 'test' }, - } - } - - clearCachedSuggestions(): void {} - - override codeModernizerCreateUploadUrl( - request: CreateUploadUrlRequest - ): Promise { - throw new Error('Method not implemented.') - } - - override codeModernizerStartCodeTransformation( - request: StartTransformationRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override codeModernizerStopCodeTransformation( - request: StopTransformationRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override codeModernizerGetCodeTransformation( - request: GetTransformationRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override codeModernizerGetCodeTransformationPlan( - request: GetTransformationPlanRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override createUploadUrl( - request: CreateUploadUrlRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override startCodeAnalysis( - request: StartCodeAnalysisRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override getCodeAnalysis( - request: GetCodeAnalysisRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override listCodeAnalysisFindings( - request: ListCodeAnalysisFindingsRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override listAvailableCustomizations( - request: ListAvailableCustomizationsRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override listAvailableProfiles( - request: ListAvailableProfilesRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override sendTelemetryEvent( - request: SendTelemetryEventRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override createWorkspace( - request: CreateWorkspaceRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override listWorkspaceMetadata( - request: ListWorkspaceMetadataRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override deleteWorkspace( - request: DeleteWorkspaceRequest - ): Promise> { - throw new Error('Method not implemented.') - } - - override listFeatureEvaluations( - request: ListFeatureEvaluationsRequest - ): Promise> { - throw new Error('Method not implemented.') - } - } - - service = new TestCodeWhispererService('us-east-1', 'https://codewhisperer.us-east-1.amazonaws.com') - }) - - describe('constructor', function () { - it('should initialize with region and endpoint', function () { - assert.strictEqual((service as any).testCodeWhispererRegion, 'us-east-1') - assert.strictEqual( - (service as any).testCodeWhispererEndpoint, - 'https://codewhisperer.us-east-1.amazonaws.com' - ) - }) - }) - - describe('request tracking', function () { - it('should track inflight requests', function () { - const mockRequest = { - abort: sandbox.stub(), - } as any - - service.trackRequest(mockRequest) - assert.strictEqual(service.inflightRequests.size, 1) - assert.strictEqual(service.inflightRequests.has(mockRequest), true) - }) - - it('should complete and remove tracked requests', function () { - const mockRequest = { - abort: sandbox.stub(), - } as any - - service.trackRequest(mockRequest) - service.completeRequest(mockRequest) - - assert.strictEqual(service.inflightRequests.size, 0) - assert.strictEqual(service.inflightRequests.has(mockRequest), false) - }) - - it('should abort all inflight requests', function () { - const mockRequest1 = { abort: sandbox.stub() } as any - const mockRequest2 = { abort: sandbox.stub() } as any - - service.trackRequest(mockRequest1) - service.trackRequest(mockRequest2) - - service.abortInflightRequests() - - assert.strictEqual(mockRequest1.abort.calledOnce, true) - assert.strictEqual(mockRequest2.abort.calledOnce, true) - assert.strictEqual(service.inflightRequests.size, 0) - }) - }) - - describe('updateClientConfig', function () { - it('should update client configuration', function () { - const mockClient = { - config: { - update: sandbox.stub(), - }, - // Add minimal required properties to satisfy the interface - createCodeScan: sandbox.stub(), - createCodeScanUploadUrl: sandbox.stub(), - createProfile: sandbox.stub(), - deleteProfile: sandbox.stub(), - generateCompletions: sandbox.stub(), - generateSuggestions: sandbox.stub(), - getCodeAnalysis: sandbox.stub(), - getCodeScan: sandbox.stub(), - listCodeAnalysisFindings: sandbox.stub(), - listCodeScans: sandbox.stub(), - listFeatureEvaluations: sandbox.stub(), - listProfiles: sandbox.stub(), - sendTelemetryEvent: sandbox.stub(), - startCodeAnalysis: sandbox.stub(), - stopCodeAnalysis: sandbox.stub(), - updateProfile: sandbox.stub(), - } as any - service.client = mockClient - - const options: ConfigurationOptions = { region: 'us-west-2' } - service.updateClientConfig(options) - - assert.strictEqual(mockClient.config.update.calledOnceWith(options), true) - }) - }) - - describe('generateItemId', function () { - it('should generate unique item IDs', function () { - const id1 = service.generateItemId() - const id2 = service.generateItemId() - - assert.strictEqual(typeof id1, 'string') - assert.strictEqual(typeof id2, 'string') - assert.notStrictEqual(id1, id2) - }) - }) - }) - - describe('CodeWhispererServiceIAM', function () { - let service: CodeWhispererServiceIAM - - beforeEach(function () { - // Mock the createCodeWhispererSigv4Client function to avoid real client creation - const mockClient = { - generateRecommendations: sandbox.stub().returns({ - promise: sandbox.stub().resolves({ - recommendations: [], - $response: { - requestId: 'test-request-id', - httpResponse: { - headers: { 'x-amzn-sessionid': 'test-session-id' }, - }, - }, - }), - }), - setupRequestListeners: sandbox.stub(), - config: { - update: sandbox.stub(), - }, - } - - // Mock the client creation - const createClientStub = sandbox.stub( - require('../client/sigv4/codewhisperer'), - 'createCodeWhispererSigv4Client' - ) - createClientStub.returns(mockClient) - - service = new CodeWhispererServiceIAM( - mockCredentialsProvider as any, - {} as any, // workspace parameter - mockLogging as any, - 'us-east-1', - 'https://codewhisperer.us-east-1.amazonaws.com', - mockSDKInitializator as any - ) - }) - - describe('getCredentialsType', function () { - it('should return iam credentials type', function () { - assert.strictEqual(service.getCredentialsType(), 'iam') - }) - }) - - describe('generateSuggestions', function () { - it('should call client.generateRecommendations and process response', async function () { - const mockRequest: GenerateSuggestionsRequest = { - fileContext: { - filename: 'test.js', - programmingLanguage: { languageName: 'javascript' }, - leftFileContent: 'const x = ', - rightFileContent: '', - }, - maxResults: 5, - } - - const result = await service.generateSuggestions(mockRequest) - - assert.strictEqual(Array.isArray(result.suggestions), true) - assert.strictEqual(typeof result.responseContext.requestId, 'string') - assert.strictEqual(typeof result.responseContext.codewhispererSessionId, 'string') - }) - - it('should add customizationArn to request if set', async function () { - service.customizationArn = 'test-arn' - - const mockRequest: GenerateSuggestionsRequest = { - fileContext: { - filename: 'test.js', - programmingLanguage: { languageName: 'javascript' }, - leftFileContent: 'const x = ', - rightFileContent: '', - }, - maxResults: 5, - } - - await service.generateSuggestions(mockRequest) - - // Verify that the client was called with the customizationArn - const clientCall = (service.client.generateRecommendations as sinon.SinonStub).getCall(0) - assert.strictEqual(clientCall.args[0].customizationArn, 'test-arn') - }) - }) - }) - - describe('CodeWhispererServiceToken', function () { - let service: CodeWhispererServiceToken - let mockClient: any - - beforeEach(function () { - // Mock the token client - mockClient = { - generateCompletions: sandbox.stub().returns({ - promise: sandbox.stub().resolves({ - completions: [ - { - content: 'console.log("hello");', - references: [], - }, - ], - $response: { - requestId: 'test-request-id', - httpResponse: { - headers: { 'x-amzn-sessionid': 'test-session-id' }, - }, - }, - }), - }), - config: { - update: sandbox.stub(), - }, - } - - // Mock the client creation - const createTokenClientStub = sandbox.stub( - require('../client/token/codewhisperer'), - 'createCodeWhispererTokenClient' - ) - createTokenClientStub.returns(mockClient) - - // Mock bearer credentials - mockCredentialsProvider.getCredentials.returns({ - token: 'mock-bearer-token', - }) - - service = new CodeWhispererServiceToken( - mockCredentialsProvider as any, - mockWorkspace as any, - mockLogging as any, - 'us-east-1', - 'https://codewhisperer.us-east-1.amazonaws.com', - mockSDKInitializator as any - ) - }) - - describe('getCredentialsType', function () { - it('should return bearer credentials type', function () { - assert.strictEqual(service.getCredentialsType(), 'bearer') - }) - }) - - describe('generateSuggestions', function () { - it('should call client.generateCompletions and process response', async function () { - const mockRequest: GenerateSuggestionsRequest = { - fileContext: { - filename: 'test.js', - programmingLanguage: { languageName: 'javascript' }, - leftFileContent: 'const x = ', - rightFileContent: '', - }, - maxResults: 5, - } - - const result = await service.generateSuggestions(mockRequest) - - assert.strictEqual(mockClient.generateCompletions.calledOnce, true) - assert.strictEqual(Array.isArray(result.suggestions), true) - assert.strictEqual(typeof result.responseContext.requestId, 'string') - assert.strictEqual(typeof result.responseContext.codewhispererSessionId, 'string') - }) - - it('should add customizationArn to request if set', async function () { - service.customizationArn = 'test-arn' - - const mockRequest: GenerateSuggestionsRequest = { - fileContext: { - filename: 'test.js', - programmingLanguage: { languageName: 'javascript' }, - leftFileContent: 'const x = ', - rightFileContent: '', - }, - maxResults: 5, - } - - await service.generateSuggestions(mockRequest) - - const clientCall = mockClient.generateCompletions.getCall(0) - assert.strictEqual(clientCall.args[0].customizationArn, 'test-arn') - }) - - it('should process profile ARN with withProfileArn method', async function () { - const mockRequest: GenerateSuggestionsRequest = { - fileContext: { - filename: 'test.js', - programmingLanguage: { languageName: 'javascript' }, - leftFileContent: 'const x = ', - rightFileContent: '', - }, - maxResults: 5, - } - - const withProfileArnStub = sandbox.stub(service, 'withProfileArn' as any) - withProfileArnStub.returns(mockRequest) - - await service.generateSuggestions(mockRequest) - - assert.strictEqual(withProfileArnStub.calledOnceWith(mockRequest), true) - }) - }) - }) -}) From 86614a6f500c0c43feab11f6ede55ed73e69aa9c Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 23 Jul 2025 10:27:59 -0400 Subject: [PATCH 4/8] refactor: combine IAM and SSO implementations of AmazonQServiceManager --- .../agenticChat/agenticChatController.test.ts | 14 +- .../AmazonQServiceManager.test.ts | 1099 +++++++++++++++++ .../AmazonQServiceManager.ts | 654 ++++++++++ .../src/shared/testUtils.ts | 20 +- 4 files changed, 1782 insertions(+), 5 deletions(-) create mode 100644 server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts create mode 100644 server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts index 3ee23089b2..bb9ced1528 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts @@ -35,7 +35,11 @@ import { } from '@aws/language-server-runtimes/server-interface' import { TestFeatures } from '@aws/language-server-runtimes/testing' import * as assert from 'assert' -import { createIterableResponse, setCredentialsForAmazonQTokenServiceManagerFactory } from '../../shared/testUtils' +import { + createIterableResponse, + setCredentialsForAmazonQTokenServiceManagerFactory, + setIamCredentialsForAmazonQServiceManagerFactory, +} from '../../shared/testUtils' import sinon from 'ts-sinon' import { AgenticChatController } from './agenticChatController' import { ChatSessionManagementService } from '../chat/chatSessionManagementService' @@ -179,7 +183,8 @@ describe('AgenticChatController', () => { let getMessagesStub: sinon.SinonStub let addMessageStub: sinon.SinonStub - const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures) + const setSsoCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures) + const setIamCredentials = setIamCredentialsForAmazonQServiceManagerFactory(() => testFeatures) beforeEach(() => { // Override the response timeout for tests to avoid long waits @@ -272,7 +277,7 @@ describe('AgenticChatController', () => { } testFeatures.lsp.window.showDocument = sinon.stub() testFeatures.setClientParams(cachedInitializeParams) - setCredentials('builderId') + setSsoCredentials('builderId') activeTabSpy = sinon.spy(ChatTelemetryController.prototype, 'activeTabId', ['get', 'set']) removeConversationSpy = sinon.spy(ChatTelemetryController.prototype, 'removeConversation') @@ -3144,6 +3149,9 @@ ${' '.repeat(8)}} // Reset the singleton instance ChatSessionManagementService.reset() + // Store IAM credentials + setIamCredentials() + // Create IAM service manager AmazonQIAMServiceManager.resetInstance() iamServiceManager = AmazonQIAMServiceManager.initInstance(testFeatures) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts new file mode 100644 index 0000000000..fcec94ac7d --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts @@ -0,0 +1,1099 @@ +import * as assert from 'assert' +import sinon, { StubbedInstance, stubInterface } from 'ts-sinon' +import { AmazonQServiceManager } from './AmazonQServiceManager' +import { TestFeatures } from '@aws/language-server-runtimes/testing' +import { GenerateSuggestionsRequest } from '../codeWhispererService/codeWhispererServiceBase' +import { CodeWhispererServiceToken } from '../codeWhispererService/codeWhispererServiceToken' +import { CodeWhispererServiceIAM } from '../codeWhispererService/codeWhispererServiceIAM' +import { + AmazonQServiceInitializationError, + AmazonQServicePendingProfileError, + AmazonQServicePendingProfileUpdateError, + AmazonQServicePendingSigninError, +} from './errors' +import { + CancellationToken, + InitializeParams, + LSPErrorCodes, + ResponseError, +} from '@aws/language-server-runtimes/protocol' +import { + AWS_Q_ENDPOINT_URL_ENV_VAR, + AWS_Q_ENDPOINTS, + AWS_Q_REGION_ENV_VAR, + DEFAULT_AWS_Q_ENDPOINT_URL, + DEFAULT_AWS_Q_REGION, +} from '../constants' +import * as qDeveloperProfilesFetcherModule from './qDeveloperProfiles' +import { + setTokenCredentialsForAmazonQServiceManagerFactory, + setIamCredentialsForAmazonQServiceManagerFactory, +} from '../testUtils' +import { StreamingClientService } from '../streamingClientService' +import { generateSingletonInitializationTests } from './testUtils' +import * as utils from '../utils' + +export const mockedProfiles: qDeveloperProfilesFetcherModule.AmazonQDeveloperProfile[] = [ + { + arn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + name: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + identityDetails: { + region: 'us-east-1', + }, + }, + { + arn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2', + name: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2', + identityDetails: { + region: 'us-east-1', + }, + }, + { + arn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', + name: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', + identityDetails: { + region: 'eu-central-1', + }, + }, +] + +const TEST_ENDPOINT_US_EAST_1 = 'http://amazon-q-in-us-east-1-endpoint' +const TEST_ENDPOINT_EU_CENTRAL_1 = 'http://amazon-q-in-eu-central-1-endpoint' + +describe('Token', () => { + let codewhispererServiceStub: StubbedInstance + let codewhispererStubFactory: sinon.SinonStub> + let sdkInitializatorSpy: sinon.SinonSpy + let getListAllAvailableProfilesHandlerStub: sinon.SinonStub + + let amazonQServiceManager: AmazonQServiceManager + let features: TestFeatures + + beforeEach(() => { + // Override endpoints for testing + AWS_Q_ENDPOINTS.set('us-east-1', TEST_ENDPOINT_US_EAST_1) + AWS_Q_ENDPOINTS.set('eu-central-1', TEST_ENDPOINT_EU_CENTRAL_1) + + getListAllAvailableProfilesHandlerStub = sinon + .stub() + .resolves( + Promise.resolve(mockedProfiles).then(() => + new Promise(resolve => setTimeout(resolve, 1)).then(() => mockedProfiles) + ) + ) + + sinon + .stub(qDeveloperProfilesFetcherModule, 'getListAllAvailableProfilesHandler') + .returns(getListAllAvailableProfilesHandlerStub) + + AmazonQServiceManager.resetInstance() + + features = new TestFeatures() + + sdkInitializatorSpy = Object.assign(sinon.spy(features.sdkInitializator), { + v2: sinon.spy(features.sdkInitializator.v2), + }) + + codewhispererServiceStub = stubInterface() + // @ts-ignore + codewhispererServiceStub.client = sinon.stub() + codewhispererServiceStub.customizationArn = undefined + codewhispererServiceStub.shareCodeWhispererContentWithAWS = false + codewhispererServiceStub.profileArn = undefined + + // Initialize the class with mocked dependencies + codewhispererStubFactory = sinon.stub().returns(codewhispererServiceStub) + }) + + afterEach(() => { + AmazonQServiceManager.resetInstance() + features.dispose() + sinon.restore() + }) + + const setupServiceManager = (enableProfiles = false) => { + // @ts-ignore + const cachedInitializeParams: InitializeParams = { + initializationOptions: { + aws: { + awsClientCapabilities: { + q: { + developerProfiles: enableProfiles, + }, + }, + }, + }, + } + features.setClientParams(cachedInitializeParams) + + AmazonQServiceManager.initInstance(features) + amazonQServiceManager = AmazonQServiceManager.getInstance() + amazonQServiceManager.setServiceFactory(codewhispererStubFactory) + } + + const setCredentials = setTokenCredentialsForAmazonQServiceManagerFactory(() => features) + + const clearCredentials = () => { + features.credentialsProvider.hasCredentials.returns(false) + features.credentialsProvider.getCredentials.returns(undefined) + features.credentialsProvider.getConnectionType.returns('none') + } + + const setupServiceManagerWithProfile = async ( + profileArn = 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' + ): Promise => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: profileArn, + }, + }, + {} as CancellationToken + ) + + const service = amazonQServiceManager.getCodewhispererService() + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + + return service + } + + describe('Initialization process', () => { + generateSingletonInitializationTests(AmazonQServiceManager) + }) + + describe('Client is not connected', () => { + it('should be in PENDING_CONNECTION state when bearer token is not set', () => { + setupServiceManager() + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + clearCredentials() + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingSigninError) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'none') + }) + }) + + describe('Clear state upon bearer token deletion', () => { + let cancelActiveProfileChangeTokenSpy: sinon.SinonSpy + + beforeEach(() => { + setupServiceManager() + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + cancelActiveProfileChangeTokenSpy = sinon.spy( + amazonQServiceManager as any, + 'cancelActiveProfileChangeToken' + ) + + setCredentials('builderId') + }) + + it('should clear local state variables on receiving bearer token deletion event', () => { + amazonQServiceManager.getCodewhispererService() + + amazonQServiceManager.handleOnCredentialsDeleted('bearer') + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'none') + assert.strictEqual((amazonQServiceManager as any)['cachedCodewhispererService'], undefined) + assert.strictEqual((amazonQServiceManager as any)['cachedStreamingClient'], undefined) + assert.strictEqual((amazonQServiceManager as any)['activeIdcProfile'], undefined) + sinon.assert.calledOnce(cancelActiveProfileChangeTokenSpy) + }) + + it('should not clear local state variables on receiving iam token deletion event', () => { + amazonQServiceManager.getCodewhispererService() + + amazonQServiceManager.handleOnCredentialsDeleted('iam') + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') + assert(!(amazonQServiceManager['cachedCodewhispererService'] === undefined)) + assert.strictEqual((amazonQServiceManager as any)['activeIdcProfile'], undefined) + sinon.assert.notCalled(cancelActiveProfileChangeTokenSpy) + }) + }) + + describe('BuilderId support', () => { + const testRegion = 'some-region' + const testEndpoint = 'http://some-endpoint-in-some-region' + + beforeEach(() => { + setupServiceManager() + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('builderId') + + AWS_Q_ENDPOINTS.set(testRegion, testEndpoint) + + features.lsp.getClientInitializeParams.reset() + }) + + it('should be INITIALIZED with BuilderId Connection', async () => { + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') + + assert(streamingClient instanceof StreamingClientService) + assert(codewhispererServiceStub.generateSuggestions.calledOnce) + }) + + it('should initialize service with region set by client', async () => { + features.setClientParams({ + processId: 0, + rootUri: 'some-root-uri', + capabilities: {}, + initializationOptions: { + aws: { + region: testRegion, + }, + }, + }) + + amazonQServiceManager.getCodewhispererService() + assert(codewhispererStubFactory.calledOnceWithExactly(testRegion, testEndpoint)) + + const streamingClient = amazonQServiceManager.getStreamingClient() + assert.strictEqual(await streamingClient.client.config.region(), testRegion) + assert.strictEqual( + (await streamingClient.client.config.endpoint()).hostname, + 'some-endpoint-in-some-region' + ) + }) + + it('should initialize service with region set by runtime if not set by client', async () => { + features.runtime.getConfiguration.withArgs(AWS_Q_REGION_ENV_VAR).returns('eu-central-1') + features.runtime.getConfiguration.withArgs(AWS_Q_ENDPOINT_URL_ENV_VAR).returns(TEST_ENDPOINT_EU_CENTRAL_1) + + amazonQServiceManager.getCodewhispererService() + assert(codewhispererStubFactory.calledOnceWithExactly('eu-central-1', TEST_ENDPOINT_EU_CENTRAL_1)) + + const streamingClient = amazonQServiceManager.getStreamingClient() + assert.strictEqual(await streamingClient.client.config.region(), 'eu-central-1') + assert.strictEqual( + (await streamingClient.client.config.endpoint()).hostname, + 'amazon-q-in-eu-central-1-endpoint' + ) + }) + + it('should initialize service with default region if not set by client and runtime', async () => { + amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + + assert(codewhispererStubFactory.calledOnceWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) + + assert.strictEqual(await streamingClient.client.config.region(), DEFAULT_AWS_Q_REGION) + assert.strictEqual( + (await streamingClient.client.config.endpoint()).hostname, + 'codewhisperer.us-east-1.amazonaws.com' + ) + }) + }) + + describe('IdentityCenter support', () => { + describe('Developer Profiles Support is disabled', () => { + it('should be INITIALIZED with IdentityCenter Connection', async () => { + setupServiceManager() + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert(codewhispererServiceStub.generateSuggestions.calledOnce) + + assert(streamingClient instanceof StreamingClientService) + }) + }) + + describe('Developer Profiles Support is enabled', () => { + it('should not throw when receiving null profile arn in PENDING_CONNECTION state', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + await assert.doesNotReject( + amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: null, + }, + }, + {} as CancellationToken + ) + ) + + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + }) + + it('should initialize to PENDING_Q_PROFILE state when IdentityCenter Connection is set', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingProfileError) + assert.throws(() => amazonQServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + }) + + it('handles Profile configuration request for valid profile and initializes to INITIALIZED state', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') + }) + + it('handles Profile configuration request for valid profile & cancels the old in-flight update request', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + assert.strictEqual((amazonQServiceManager as any)['profileChangeTokenSource'], undefined) + + let firstRequestStarted = false + const originalHandleProfileChange = amazonQServiceManager['handleProfileChange'] + amazonQServiceManager['handleProfileChange'] = async (...args) => { + firstRequestStarted = true + return originalHandleProfileChange.apply(amazonQServiceManager, args) + } + const firstUpdate = amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + while (!firstRequestStarted) { + await new Promise(resolve => setTimeout(resolve, 1)) + } + const secondUpdate = amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + const results = await Promise.allSettled([firstUpdate, secondUpdate]) + + assert.strictEqual((amazonQServiceManager as any)['profileChangeTokenSource'], undefined) + const service = amazonQServiceManager.getCodewhispererService() + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + + assert.strictEqual(results[0].status, 'fulfilled') + assert.strictEqual(results[1].status, 'fulfilled') + }) + + it('handles Profile configuration change to valid profile in same region', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient1 = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' + ) + + assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) + assert(streamingClient1 instanceof StreamingClientService) + assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1') + + // Profile change + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2', + }, + }, + {} as CancellationToken + ) + await service.generateSuggestions({} as GenerateSuggestionsRequest) + const streamingClient2 = amazonQServiceManager.getStreamingClient() + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2' + ) + + // CodeWhisperer Service was not recreated + assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) + + assert(streamingClient2 instanceof StreamingClientService) + assert.strictEqual(streamingClient1, streamingClient2) + assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1') + }) + + it('handles Profile configuration change to valid profile in different region', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient1 = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' + ) + assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) + + assert(streamingClient1 instanceof StreamingClientService) + assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1') + + // Profile change + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + await service.generateSuggestions({} as GenerateSuggestionsRequest) + const streamingClient2 = amazonQServiceManager.getStreamingClient() + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ' + ) + + // CodeWhisperer Service was recreated + assert(codewhispererStubFactory.calledTwice) + assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, [ + 'eu-central-1', + TEST_ENDPOINT_EU_CENTRAL_1, + ]) + + // Streaming Client was recreated + assert(streamingClient2 instanceof StreamingClientService) + assert.notStrictEqual(streamingClient1, streamingClient2) + assert.strictEqual(await streamingClient2.client.config.region(), 'eu-central-1') + }) + + // As we're not validating profile at this moment, there is no "invalid" profile + it.skip('handles Profile configuration change from valid to invalid profile', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + let service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' + ) + assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') + + // Profile change to invalid profile + + await assert.rejects( + amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/invalid-profile-arn', + }, + }, + {} as CancellationToken + ), + new ResponseError(LSPErrorCodes.RequestFailed, 'Requested Amazon Q Profile does not exist', { + awsErrorCode: 'E_AMAZON_Q_INVALID_PROFILE', + }) + ) + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingProfileError) + assert.throws(() => amazonQServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + // CodeWhisperer Service was not recreated + assert(codewhispererStubFactory.calledOnce) + assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1]) + }) + + // As we're not validating profile at this moment, there is no "non-existing" profile + it.skip('handles non-existing profile selection', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await assert.rejects( + amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/invalid-profile-arn', + }, + }, + {} as CancellationToken + ), + new ResponseError(LSPErrorCodes.RequestFailed, 'Requested Amazon Q Profile does not exist', { + awsErrorCode: 'E_AMAZON_Q_INVALID_PROFILE', + }) + ) + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingProfileError) + assert.throws(() => amazonQServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(codewhispererStubFactory.notCalled) + }) + + it('prevents service usage while profile change is inflight when profile was not set', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingProfileError) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + + amazonQServiceManager.setState('PENDING_Q_PROFILE_UPDATE') + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') + + assert.throws( + () => amazonQServiceManager.getCodewhispererService(), + AmazonQServicePendingProfileUpdateError + ) + assert.throws(() => amazonQServiceManager.getStreamingClient(), AmazonQServicePendingProfileUpdateError) + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ' + ) + assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, [ + 'eu-central-1', + TEST_ENDPOINT_EU_CENTRAL_1, + ]) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'eu-central-1') + }) + + it('prevents service usage while profile change is inflight when profile was set before', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingProfileError) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + const service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual( + amazonQServiceManager.getActiveProfileArn(), + 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' + ) + assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1]) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') + + // Updaing profile + amazonQServiceManager.setState('PENDING_Q_PROFILE_UPDATE') + assert.throws( + () => amazonQServiceManager.getCodewhispererService(), + AmazonQServicePendingProfileUpdateError + ) + assert.throws(() => amazonQServiceManager.getStreamingClient(), AmazonQServicePendingProfileUpdateError) + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') + }) + + it('resets to PENDING_PROFILE from INITIALIZED when receiving null profileArn', async () => { + await setupServiceManagerWithProfile() + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: null, + }, + }, + {} as CancellationToken + ) + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests) + }) + + it('resets to PENDING_Q_PROFILE from PENDING_Q_PROFILE_UPDATE when receiving null profileArn', async () => { + await setupServiceManagerWithProfile() + + amazonQServiceManager.setState('PENDING_Q_PROFILE_UPDATE') + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') + + // Null profile arn + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: null, + }, + }, + {} as CancellationToken + ) + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests) + assert.throws(() => amazonQServiceManager.getCodewhispererService()) + }) + + it('cancels on-going profile update when credentials are deleted', async () => { + await setupServiceManagerWithProfile() + + amazonQServiceManager.setState('PENDING_Q_PROFILE_UPDATE') + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') + + amazonQServiceManager.handleOnCredentialsDeleted('bearer') + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests) + assert.throws(() => amazonQServiceManager.getCodewhispererService()) + }) + + // Due to service limitation, validation was removed for the sake of recovering API availability + // When service is ready to take more tps, revert https://github.com/aws/language-servers/pull/1329 to add profile validation + it('should not call service to validate profile and always assume its validness', async () => { + setupServiceManager(true) + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + + setCredentials('identityCenter') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + sinon.assert.notCalled(getListAllAvailableProfilesHandlerStub) + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + }) + }) + }) + + describe('Connection types with no Developer Profiles support', () => { + it('handles reauthentication scenario when connection type is none but profile ARN is provided', async () => { + setupServiceManager(true) + clearCredentials() + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'none') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ) + + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + }) + + it('ignores null profile when connection type is none', async () => { + setupServiceManager(true) + clearCredentials() + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'none') + + await amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: null, + }, + }, + {} as CancellationToken + ) + + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'none') + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') + }) + + it('returns error when profile update is requested and connection type is builderId', async () => { + setupServiceManager(true) + setCredentials('builderId') + + await assert.rejects( + amazonQServiceManager.handleOnUpdateConfiguration( + { + section: 'aws.q', + settings: { + profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', + }, + }, + {} as CancellationToken + ), + new ResponseError( + LSPErrorCodes.RequestFailed, + 'Connection type builderId does not support Developer Profiles feature.', + { + awsErrorCode: 'E_AMAZON_Q_CONNECTION_NO_PROFILE_SUPPORT', + } + ) + ) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') + }) + }) + + describe('Handle connection type changes', () => { + describe('connection changes from BuilderId to IdentityCenter', () => { + it('should initialize service with default region when profile support is disabled', async () => { + setupServiceManager(false) + setCredentials('builderId') + + let service1 = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service1.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') + + setCredentials('identityCenter') + let service2 = amazonQServiceManager.getCodewhispererService() + const streamingClient2 = amazonQServiceManager.getStreamingClient() + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(codewhispererStubFactory.calledTwice) + assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) + + assert(streamingClient2 instanceof StreamingClientService) + assert.strictEqual(await streamingClient2.client.config.region(), DEFAULT_AWS_Q_REGION) + }) + + it('should initialize service to PENDING_Q_PROFILE state when profile support is enabled', async () => { + setupServiceManager(true) + setCredentials('builderId') + + let service = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') + + setCredentials('identityCenter') + + assert.throws(() => amazonQServiceManager.getCodewhispererService(), AmazonQServicePendingProfileError) + assert.throws(() => amazonQServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) + + assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_Q_PROFILE') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(codewhispererStubFactory.calledOnce) + assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) + }) + }) + + describe('connection changes from IdentityCenter to BuilderId', () => { + it('should initialize service in default IAD region', async () => { + setupServiceManager(false) + setCredentials('identityCenter') + + let service1 = amazonQServiceManager.getCodewhispererService() + const streamingClient = amazonQServiceManager.getStreamingClient() + await service1.generateSuggestions({} as GenerateSuggestionsRequest) + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(streamingClient instanceof StreamingClientService) + assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') + + setCredentials('builderId') + let service2 = amazonQServiceManager.getCodewhispererService() + const streamingClient2 = amazonQServiceManager.getStreamingClient() + + assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') + assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') + assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) + + assert(codewhispererStubFactory.calledTwice) + assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) + + assert(streamingClient2 instanceof StreamingClientService) + assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1') + }) + }) + }) + + describe('handle LSP Configuration settings', () => { + it('should initialize codewhisperer service with default configurations when not set by client', async () => { + setupServiceManager() + setCredentials('identityCenter') + + await amazonQServiceManager.handleDidChangeConfiguration() + + const service = amazonQServiceManager.getCodewhispererService() + + assert.strictEqual(service.customizationArn, undefined) + assert.strictEqual(service.shareCodeWhispererContentWithAWS, false) + }) + + it('should returned configured codewhispererService with expected configuration values', async () => { + const getConfigStub = features.lsp.workspace.getConfiguration + getConfigStub.withArgs('aws.q').resolves({ + customization: 'test-customization-arn', + optOutTelemetryPreference: true, + }) + getConfigStub.withArgs('aws.codeWhisperer').resolves({ + includeSuggestionsWithCodeReferences: true, + shareCodeWhispererContentWithAWS: true, + }) + + // Initialize mock server + setupServiceManager() + setCredentials('identityCenter') + + amazonQServiceManager = AmazonQServiceManager.getInstance() + const service = amazonQServiceManager.getCodewhispererService() + + assert.strictEqual(service.customizationArn, undefined) + assert.strictEqual(service.shareCodeWhispererContentWithAWS, false) + + await amazonQServiceManager.handleDidChangeConfiguration() + + // Force next tick to allow async work inside handleDidChangeConfiguration to complete + await Promise.resolve() + + assert.strictEqual(service.customizationArn, 'test-customization-arn') + assert.strictEqual(service.shareCodeWhispererContentWithAWS, true) + }) + }) + + describe('Initialize', () => { + it('should throw when initialize is called before LSP has been initialized with InitializeParams', () => { + features.resetClientParams() + + assert.throws(() => AmazonQServiceManager.initInstance(features), AmazonQServiceInitializationError) + }) + }) +}) + +describe('IAM', () => { + describe('Initialization process', () => { + generateSingletonInitializationTests(AmazonQServiceManager) + }) + + describe('Service caching', () => { + let serviceManager: AmazonQServiceManager + let features: TestFeatures + let updateCachedServiceConfigSpy: sinon.SinonSpy + + const setCredentials = setIamCredentialsForAmazonQServiceManagerFactory(() => features) + + beforeEach(() => { + features = new TestFeatures() + features.lsp.getClientInitializeParams.resolves({}) + + updateCachedServiceConfigSpy = sinon.spy( + AmazonQServiceManager.prototype, + 'updateCachedServiceConfig' as keyof AmazonQServiceManager + ) + + AmazonQServiceManager.resetInstance() + serviceManager = AmazonQServiceManager.initInstance(features) + }) + + afterEach(() => { + AmazonQServiceManager.resetInstance() + features.dispose() + sinon.restore() + }) + + it('should initialize the CodeWhisperer service only once', () => { + setCredentials() + const service = serviceManager.getCodewhispererService() + sinon.assert.calledOnce(updateCachedServiceConfigSpy) + + assert.deepStrictEqual(serviceManager.getCodewhispererService(), service) + sinon.assert.calledOnce(updateCachedServiceConfigSpy) + }) + + it('should initialize the streaming client only once', () => { + setCredentials() + // Mock getIAMCredentialsFromProvider to return dummy credentials + const getIAMCredentialsStub = sinon.stub(utils, 'getIAMCredentialsFromProvider').returns({ + accessKeyId: 'dummy-access-key', + secretAccessKey: 'dummy-secret-key', + sessionToken: 'dummy-session-token', + }) + + const streamingClient = serviceManager.getStreamingClient() + + assert.deepStrictEqual(serviceManager.getStreamingClient(), streamingClient) + + getIAMCredentialsStub.restore() + }) + }) +}) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts new file mode 100644 index 0000000000..e9c1afc2df --- /dev/null +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts @@ -0,0 +1,654 @@ +import { + UpdateConfigurationParams, + ResponseError, + LSPErrorCodes, + SsoConnectionType, + CancellationToken, + CredentialsType, + InitializeParams, + CancellationTokenSource, +} from '@aws/language-server-runtimes/server-interface' +import { CodeWhispererServiceToken } from '../codeWhispererService/codeWhispererServiceToken' +import { CodeWhispererServiceIAM } from '../codeWhispererService/codeWhispererServiceIAM' +import { + AmazonQError, + AmazonQServiceAlreadyInitializedError, + AmazonQServiceInitializationError, + AmazonQServiceInvalidProfileError, + AmazonQServiceNoProfileSupportError, + AmazonQServiceNotInitializedError, + AmazonQServicePendingProfileError, + AmazonQServicePendingProfileUpdateError, + AmazonQServicePendingSigninError, + AmazonQServiceProfileUpdateCancelled, +} from './errors' +import { + AmazonQBaseServiceManager, + BaseAmazonQServiceManager, + QServiceManagerFeatures, +} from './BaseAmazonQServiceManager' +import { AWS_Q_ENDPOINTS, Q_CONFIGURATION_SECTION } from '../constants' +import { AmazonQDeveloperProfile, signalsAWSQDeveloperProfilesEnabled } from './qDeveloperProfiles' +import { isStringOrNull } from '../utils' +import { getAmazonQRegionAndEndpoint } from './configurationUtils' +import { getUserAgent } from '../telemetryUtils' +import { StreamingClientServiceToken, StreamingClientServiceIAM } from '../streamingClientService' +import { parse } from '@aws-sdk/util-arn-parser' + +/** + * AmazonQServiceManager manages state and provides centralized access to + * instance of CodeWhispererService SDK client to any consuming code. + * It ensures that CodeWhispererService is configured to always access correct regionalized Amazon Q Developer API endpoint. + * Regional endppoint is selected based on: + * 1) current SSO auth connection type (BuilderId or IDC). + * 2) selected Amazon Q Developer profile (only for IDC connection type). + * + * @states + * - PENDING_CONNECTION: Initial state when no bearer token is set + * - PENDING_Q_PROFILE: When using Identity Center and waiting for profile selection + * - PENDING_Q_PROFILE_UPDATE: During profile update operation + * - INITIALIZED: Service is ready to handle requests + * + * @connectionTypes + * - none: No active connection + * - builderId: Connected via Builder ID + * - identityCenter: Connected via Identity Center + * + * AmazonQServiceManager is a singleton class, which must be instantiated with Language Server runtimes [Features](https://github.com/aws/language-server-runtimes/blob/21d5d1dc7c73499475b7c88c98d2ce760e5d26c8/runtimes/server-interface/server.ts#L31-L42) + * in the `AmazonQServiceServer` via the `initBaseServiceManager` factory. Dependencies of this class can access the singleton via + * the `getOrThrowBaseServiceManager` factory or `getInstance()` method after the initialized notification has been received during + * the LSP hand shake. + * + */ +export class AmazonQServiceManager extends BaseAmazonQServiceManager { + private static instance: AmazonQServiceManager | null = null + private enableDeveloperProfileSupport?: boolean + private activeIdcProfile?: AmazonQDeveloperProfile + private connectionType?: SsoConnectionType + private profileChangeTokenSource: CancellationTokenSource | undefined + private region?: string + private endpoint?: string + private regionChangeListeners: Array<(region: string) => void> = [] + /** + * Internal state of Service connection, based on status of bearer token and Amazon Q Developer profile selection. + * Supported states: + * PENDING_CONNECTION - Waiting for (Bearer Token and StartURL) or (Access Key and Secret Key) to be passed + * PENDING_Q_PROFILE - (only for identityCenter connection) waiting for setting Developer Profile + * PENDING_Q_PROFILE_UPDATE (only for identityCenter connection) waiting for Developer Profile to complete + * INITIALIZED - Service is initialized + */ + private state: 'PENDING_CONNECTION' | 'PENDING_Q_PROFILE' | 'PENDING_Q_PROFILE_UPDATE' | 'INITIALIZED' = + 'PENDING_CONNECTION' + + private constructor(features: QServiceManagerFeatures) { + super(features) + } + + // @VisibleForTesting, please DO NOT use in production + setState(state: 'PENDING_CONNECTION' | 'PENDING_Q_PROFILE' | 'PENDING_Q_PROFILE_UPDATE' | 'INITIALIZED') { + this.state = state + } + + public static initInstance(features: QServiceManagerFeatures): AmazonQServiceManager { + if (!AmazonQServiceManager.instance) { + AmazonQServiceManager.instance = new AmazonQServiceManager(features) + AmazonQServiceManager.instance.initialize() + + return AmazonQServiceManager.instance + } + + throw new AmazonQServiceAlreadyInitializedError() + } + + public static getInstance(): AmazonQServiceManager { + if (!AmazonQServiceManager.instance) { + throw new AmazonQServiceInitializationError( + 'Amazon Q service has not been initialized yet. Make sure the Amazon Q server is present and properly initialized.' + ) + } + + return AmazonQServiceManager.instance + } + + private initialize(): void { + if (!this.features.lsp.getClientInitializeParams()) { + this.log('AmazonQServiceManager initialized before LSP connection was initialized.') + throw new AmazonQServiceInitializationError( + 'AmazonQServiceManager initialized before LSP connection was initialized.' + ) + } + + // Bind methods that are passed by reference to some handlers to maintain proper scope. + this.serviceFactory = this.serviceFactory.bind(this) + + this.log('Reading enableDeveloperProfileSupport setting from AWSInitializationOptions') + if (this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws) { + const awsOptions = this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws || {} + this.enableDeveloperProfileSupport = signalsAWSQDeveloperProfilesEnabled(awsOptions) + + this.log(`Enabled Q Developer Profile support: ${this.enableDeveloperProfileSupport}`) + } + + this.connectionType = 'none' + this.state = 'PENDING_CONNECTION' + + this.log('Manager instance is initialize') + } + + public handleOnCredentialsDeleted(type: CredentialsType): void { + this.log(`Received credentials delete event for type: ${type}`) + if (type === 'iam') { + return + } + this.cancelActiveProfileChangeToken() + + this.resetCodewhispererService() + this.connectionType = 'none' + this.state = 'PENDING_CONNECTION' + } + + public async handleOnUpdateConfiguration(params: UpdateConfigurationParams, _token: CancellationToken) { + try { + if (params.section === Q_CONFIGURATION_SECTION && params.settings.profileArn !== undefined) { + const profileArn = params.settings.profileArn + const region = params.settings.region + + if (!isStringOrNull(profileArn)) { + throw new Error('Expected params.settings.profileArn to be of either type string or null') + } + + this.log(`Profile update is requested for profile ${profileArn}`) + this.cancelActiveProfileChangeToken() + this.profileChangeTokenSource = new CancellationTokenSource() + + await this.handleProfileChange(profileArn, this.profileChangeTokenSource.token) + } + } catch (error) { + this.log('Error updating profiles: ' + error) + if (error instanceof AmazonQServiceProfileUpdateCancelled) { + throw new ResponseError(LSPErrorCodes.ServerCancelled, error.message, { + awsErrorCode: error.code, + }) + } + if (error instanceof AmazonQError) { + throw new ResponseError(LSPErrorCodes.RequestFailed, error.message, { + awsErrorCode: error.code, + }) + } + + throw new ResponseError(LSPErrorCodes.RequestFailed, 'Failed to update configuration') + } finally { + if (this.profileChangeTokenSource) { + this.profileChangeTokenSource.dispose() + this.profileChangeTokenSource = undefined + } + } + } + + private handleConnectionChange() { + if (this.features.credentialsProvider.hasCredentials('iam')) { + if (!this.cachedCodewhispererService) { + const amazonQRegionAndEndpoint = getAmazonQRegionAndEndpoint( + this.features.runtime, + this.features.logging + ) + this.region = amazonQRegionAndEndpoint.region + this.endpoint = amazonQRegionAndEndpoint.endpoint + this.cachedCodewhispererService = new CodeWhispererService( + this.features.credentialsProvider, + this.features.workspace, + this.features.logging, + this.region, + this.endpoint, + this.features.sdkInitializator + ) + this.updateCachedServiceConfig() + } + this.state = 'INITIALIZED' + return + } else { + this.handleSsoConnectionChange() + return + } + } + + /** + * Validate if Bearer Token Connection type has changed mid-session. + * When connection type change is detected: reinitialize CodeWhispererService class with current connection type. + */ + private handleSsoConnectionChange() { + const newConnectionType = this.features.credentialsProvider.getConnectionType() + + this.logServiceState('Validate State of SSO Connection') + + const noCreds = !this.features.credentialsProvider.hasCredentials('bearer') + const noConnectionType = newConnectionType === 'none' + if (noCreds || noConnectionType) { + // Connection was reset, wait for SSO connection token from client + this.log( + `No active SSO connection is detected: no ${noCreds ? 'credentials' : 'connection type'} provided. Resetting the client` + ) + this.resetCodewhispererService() + this.connectionType = 'none' + this.state = 'PENDING_CONNECTION' + + return + } + + // Connection type hasn't change. + + if (newConnectionType === this.connectionType) { + this.logging.debug(`Connection type did not change: ${this.connectionType}`) + + return + } + + // Connection type changed to 'builderId' + + if (newConnectionType === 'builderId') { + this.log('Detected New connection type: builderId') + this.resetCodewhispererService() + + // For the builderId connection type regional endpoint discovery chain is: + // region set by client -> runtime region -> default region + const clientParams = this.features.lsp.getClientInitializeParams() + + this.createCodewhispererServiceInstances('builderId', clientParams?.initializationOptions?.aws?.region) + this.state = 'INITIALIZED' + this.log('Initialized Amazon Q service with builderId connection') + + return + } + + // Connection type changed to 'identityCenter' + + if (newConnectionType === 'identityCenter') { + this.log('Detected New connection type: identityCenter') + + this.resetCodewhispererService() + + if (this.enableDeveloperProfileSupport) { + this.connectionType = 'identityCenter' + this.state = 'PENDING_Q_PROFILE' + this.logServiceState('Pending profile selection for IDC connection') + + return + } + + this.createCodewhispererServiceInstances('identityCenter') + this.state = 'INITIALIZED' + this.log('Initialized Amazon Q service with identityCenter connection') + + return + } + + this.logServiceState('Unknown Connection state') + } + + private cancelActiveProfileChangeToken() { + this.profileChangeTokenSource?.cancel() + this.profileChangeTokenSource?.dispose() + this.profileChangeTokenSource = undefined + } + + private handleTokenCancellationRequest(token: CancellationToken) { + if (token.isCancellationRequested) { + this.logServiceState('Handling CancellationToken cancellation request') + throw new AmazonQServiceProfileUpdateCancelled('Requested profile update got cancelled') + } + } + + private async handleProfileChange(newProfileArn: string | null, token: CancellationToken): Promise { + if (!this.enableDeveloperProfileSupport) { + this.log('Developer Profiles Support is not enabled') + return + } + + if (typeof newProfileArn === 'string' && newProfileArn.length === 0) { + throw new Error('Received invalid Profile ARN (empty string)') + } + + this.logServiceState('UpdateProfile is requested') + + // Test if connection type changed + this.handleSsoConnectionChange() + + if (this.connectionType === 'none') { + if (newProfileArn !== null) { + // During reauthentication, connection might be temporarily 'none' but user is providing a profile + // Set connection type to identityCenter to proceed with profile setting + this.connectionType = 'identityCenter' + this.state = 'PENDING_Q_PROFILE_UPDATE' + } else { + this.logServiceState('Received null profile while not connected, ignoring request') + return + } + } + + if (this.connectionType !== 'identityCenter') { + this.logServiceState('Q Profile can not be set') + throw new AmazonQServiceNoProfileSupportError( + `Connection type ${this.connectionType} does not support Developer Profiles feature.` + ) + } + + if ((this.state === 'INITIALIZED' && this.activeIdcProfile) || this.state === 'PENDING_Q_PROFILE') { + // Change status to pending to prevent API calls until profile is updated. + // Because `listAvailableProfiles` below can take few seconds to complete, + // there is possibility that client could send requests while profile is changing. + this.state = 'PENDING_Q_PROFILE_UPDATE' + } + + // Client sent an explicit null, indicating they want to reset the assigned profile (if any) + if (newProfileArn === null) { + this.logServiceState('Received null profile, resetting to PENDING_Q_PROFILE state') + this.resetCodewhispererService() + this.state = 'PENDING_Q_PROFILE' + + return + } + + const parsedArn = parse(newProfileArn) + const region = parsedArn.region + const endpoint = AWS_Q_ENDPOINTS.get(region) + if (!endpoint) { + throw new Error('Requested profileArn region is not supported') + } + + // Hack to inject a dummy profile name as it's not used by client IDE for now, if client IDE starts consuming name field then we should also pass both profile name and arn from the IDE + // When service is ready to take more tps, revert https://github.com/aws/language-servers/pull/1329 to add profile validation + const newProfile: AmazonQDeveloperProfile = { + arn: newProfileArn, + name: 'Client provided profile', + identityDetails: { + region: parsedArn.region, + }, + } + + if (!newProfile || !newProfile.identityDetails?.region) { + this.log(`Amazon Q Profile ${newProfileArn} is not valid`) + this.resetCodewhispererService() + this.state = 'PENDING_Q_PROFILE' + + throw new AmazonQServiceInvalidProfileError('Requested Amazon Q Profile does not exist') + } + + this.handleTokenCancellationRequest(token) + + if (!this.activeIdcProfile) { + this.activeIdcProfile = newProfile + this.createCodewhispererServiceInstances('identityCenter', newProfile.identityDetails.region) + this.state = 'INITIALIZED' + this.log( + `Initialized identityCenter connection to region ${newProfile.identityDetails.region} for profile ${newProfile.arn}` + ) + + return + } + + // Profile didn't change + if (this.activeIdcProfile && this.activeIdcProfile.arn === newProfile.arn) { + // Update cached profile fields, keep existing client + this.log(`Profile selection did not change, active profile is ${this.activeIdcProfile.arn}`) + this.activeIdcProfile = newProfile + this.state = 'INITIALIZED' + + return + } + + this.handleTokenCancellationRequest(token) + + // At this point new valid profile is selected. + + const oldRegion = this.activeIdcProfile.identityDetails?.region + const newRegion = newProfile.identityDetails.region + if (oldRegion === newRegion) { + this.log(`New profile is in the same region as old one, keeping exising service.`) + this.log(`New active profile is ${this.activeIdcProfile.arn}, region ${oldRegion}`) + this.activeIdcProfile = newProfile + this.state = 'INITIALIZED' + + if (this.cachedCodewhispererService) { + this.cachedCodewhispererService.profileArn = newProfile.arn + } + + if (this.cachedStreamingClient) { + this.cachedStreamingClient.profileArn = newProfile.arn + } + + return + } + + this.log(`Switching service client region from ${oldRegion} to ${newRegion}`) + this.notifyRegionChangeListeners(newRegion) + + this.handleTokenCancellationRequest(token) + + // Selected new profile is in different region. Re-initialize service + this.resetCodewhispererService() + + this.activeIdcProfile = newProfile + + this.createCodewhispererServiceInstances('identityCenter', newProfile.identityDetails.region) + this.state = 'INITIALIZED' + + return + } + + public getCodewhispererService(): CodeWhispererService { + // Prevent initiating requests while profile change is in progress. + if (this.state === 'PENDING_Q_PROFILE_UPDATE') { + throw new AmazonQServicePendingProfileUpdateError() + } + + this.handleConnectionChange() + + if (this.state === 'INITIALIZED' && this.cachedCodewhispererService) { + return this.cachedCodewhispererService + } + + if (this.state === 'PENDING_CONNECTION') { + throw new AmazonQServicePendingSigninError() + } + + if (this.state === 'PENDING_Q_PROFILE') { + throw new AmazonQServicePendingProfileError() + } + + throw new AmazonQServiceNotInitializedError() + } + + public getStreamingClient() { + this.log('Getting instance of CodeWhispererStreaming client') + + // Trigger checks in token service + const tokenService = this.getCodewhispererService() + + if (!tokenService || !this.region || !this.endpoint) { + throw new AmazonQServiceNotInitializedError() + } + + if (!this.cachedStreamingClient) { + this.cachedStreamingClient = this.streamingClientFactory(this.region, this.endpoint) + } + + return this.cachedStreamingClient + } + + private resetCodewhispererService() { + this.cachedCodewhispererService?.abortInflightRequests() + this.cachedCodewhispererService = undefined + this.cachedStreamingClient?.abortInflightRequests() + this.cachedStreamingClient = undefined + this.activeIdcProfile = undefined + this.region = undefined + this.endpoint = undefined + } + + private createCodewhispererServiceInstances( + connectionType: 'builderId' | 'identityCenter', + clientOrProfileRegion?: string + ) { + this.logServiceState('Initializing CodewhispererService') + + const { region, endpoint } = getAmazonQRegionAndEndpoint( + this.features.runtime, + this.features.logging, + clientOrProfileRegion + ) + + // Cache active region and endpoint selection + this.connectionType = connectionType + this.region = region + this.endpoint = endpoint + + this.cachedCodewhispererService = this.serviceFactory(region, endpoint) + this.log(`CodeWhispererToken service for connection type ${connectionType} was initialized, region=${region}`) + + this.cachedStreamingClient = this.streamingClientFactory(region, endpoint) + this.log(`StreamingClient service for connection type ${connectionType} was initialized, region=${region}`) + + this.logServiceState('CodewhispererService and StreamingClient Initialization finished') + } + + private getCustomUserAgent() { + const initializeParams = this.features.lsp.getClientInitializeParams() || {} + + return getUserAgent(initializeParams as InitializeParams, this.features.runtime.serverInfo) + } + + private serviceFactory(region: string, endpoint: string): CodeWhispererService { + const service = new CodeWhispererService( + this.features.credentialsProvider, + this.features.workspace, + this.features.logging, + region, + endpoint, + this.features.sdkInitializator + ) + + const customUserAgent = this.getCustomUserAgent() + service.updateClientConfig({ + customUserAgent: customUserAgent, + }) + service.customizationArn = this.configurationCache.getProperty('customizationArn') + service.profileArn = this.activeIdcProfile?.arn + service.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty( + 'shareCodeWhispererContentWithAWS' + ) + + this.log('Configured CodeWhispererService instance settings:') + this.log( + `customUserAgent=${customUserAgent}, customizationArn=${service.customizationArn}, shareCodeWhispererContentWithAWS=${service.shareCodeWhispererContentWithAWS}` + ) + + return service + } + + private streamingClientFactory(region: string, endpoint: string): StreamingClientService { + const streamingClient = new StreamingClientService( + this.features.credentialsProvider, + this.features.sdkInitializator, + this.features.logging, + region, + endpoint, + this.getCustomUserAgent() + ) + + if (this.features.credentialsProvider.hasCredentials('bearer')) { + streamingClient.profileArn = this.activeIdcProfile?.arn + } + + this.logging.debug(`Created streaming client instance region=${region}, endpoint=${endpoint}`) + return streamingClient + } + + private log(message: string): void { + const prefix = 'Amazon Q Token Service Manager' + this.logging?.log(`${prefix}: ${message}`) + } + + private logServiceState(context: string): void { + this.logging?.debug( + JSON.stringify({ + context, + state: { + serviceStatus: this.state, + connectionType: this.connectionType, + activeIdcProfile: this.activeIdcProfile, + }, + }) + ) + } + + // For Unit Tests + public static resetInstance(): void { + AmazonQServiceManager.instance = null + } + + public getState() { + return this.state + } + + public getConnectionType() { + return this.connectionType + } + + public override getActiveProfileArn() { + return this.activeIdcProfile?.arn + } + + public setServiceFactory(factory: (region: string, endpoint: string) => CodeWhispererService) { + this.serviceFactory = factory.bind(this) + } + + public getServiceFactory() { + return this.serviceFactory + } + + public getEnableDeveloperProfileSupport(): boolean { + return this.enableDeveloperProfileSupport === undefined ? false : this.enableDeveloperProfileSupport + } + + /** + * Registers a listener that will be called when the region changes + * @param listener Function that will be called with the new region + * @returns Function to unregister the listener + */ + public override onRegionChange(listener: (region: string) => void): () => void { + this.regionChangeListeners.push(listener) + // If we already have a region, notify the listener immediately + if (this.region) { + try { + listener(this.region) + } catch (error) { + this.logging.error(`Error in region change listener: ${error}`) + } + } + return () => { + this.regionChangeListeners = this.regionChangeListeners.filter(l => l !== listener) + } + } + + private notifyRegionChangeListeners(region: string): void { + this.logging.debug( + `Notifying ${this.regionChangeListeners.length} region change listeners of region: ${region}` + ) + this.regionChangeListeners.forEach(listener => { + try { + listener(region) + } catch (error) { + this.logging.error(`Error in region change listener: ${error}`) + } + }) + } + + public getRegion(): string | undefined { + return this.region + } +} + +export const initBaseServiceManager = (features: QServiceManagerFeatures) => + AmazonQServiceManager.initInstance(features) + +export const getOrThrowBaseServiceManager = (): AmazonQBaseServiceManager => AmazonQServiceManager.getInstance() diff --git a/server/aws-lsp-codewhisperer/src/shared/testUtils.ts b/server/aws-lsp-codewhisperer/src/shared/testUtils.ts index 79596d8cc6..e444809cc5 100644 --- a/server/aws-lsp-codewhisperer/src/shared/testUtils.ts +++ b/server/aws-lsp-codewhisperer/src/shared/testUtils.ts @@ -319,10 +319,10 @@ export function shuffleList(list: T[]): T[] { return shuffledList } -export const setCredentialsForAmazonQTokenServiceManagerFactory = (getFeatures: () => TestFeatures) => { +export const setTokenCredentialsForAmazonQServiceManagerFactory = (getFeatures: () => TestFeatures) => { return (connectionType: SsoConnectionType) => { const features = getFeatures() - features.credentialsProvider.hasCredentials.returns(true) + features.credentialsProvider.hasCredentials.withArgs('bearer').returns(true) features.credentialsProvider.getConnectionType.returns(connectionType) features.credentialsProvider.getCredentials.returns({ token: 'test-token', @@ -330,6 +330,22 @@ export const setCredentialsForAmazonQTokenServiceManagerFactory = (getFeatures: } } +// TODO: remove this when changing references +export const setCredentialsForAmazonQTokenServiceManagerFactory = setTokenCredentialsForAmazonQServiceManagerFactory + +export const setIamCredentialsForAmazonQServiceManagerFactory = (getFeatures: () => TestFeatures) => { + return () => { + const features = getFeatures() + features.credentialsProvider.hasCredentials.withArgs('iam').returns(true) + features.credentialsProvider.getConnectionType.returns('none') + features.credentialsProvider.getCredentials.returns({ + accessKeyId: 'test-access-key', + secretAccessKey: 'test-secret-key', + sessionToken: 'test-session-token', + }) + } +} + export const stubCodeWhispererService = () => { const service = stubInterface() From fc5c9a802c23db9dcd74e29ce56be836e0274190 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 25 Jul 2025 10:44:08 -0400 Subject: [PATCH 5/8] fix: update client references inside unified service manager --- .../AmazonQServiceManager.test.ts | 53 ++++----- .../AmazonQServiceManager.ts | 102 ++++++++++-------- 2 files changed, 79 insertions(+), 76 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts index fcec94ac7d..f8da4e8c32 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.test.ts @@ -29,7 +29,7 @@ import { setTokenCredentialsForAmazonQServiceManagerFactory, setIamCredentialsForAmazonQServiceManagerFactory, } from '../testUtils' -import { StreamingClientService } from '../streamingClientService' +import { StreamingClientServiceToken, StreamingClientServiceIAM } from '../streamingClientService' import { generateSingletonInitializationTests } from './testUtils' import * as utils from '../utils' @@ -61,8 +61,8 @@ const TEST_ENDPOINT_US_EAST_1 = 'http://amazon-q-in-us-east-1-endpoint' const TEST_ENDPOINT_EU_CENTRAL_1 = 'http://amazon-q-in-eu-central-1-endpoint' describe('Token', () => { - let codewhispererServiceStub: StubbedInstance - let codewhispererStubFactory: sinon.SinonStub> + let codewhispererServiceStub: StubbedInstance + let codewhispererStubFactory: sinon.SinonStub> let sdkInitializatorSpy: sinon.SinonSpy let getListAllAvailableProfilesHandlerStub: sinon.SinonStub @@ -94,7 +94,7 @@ describe('Token', () => { v2: sinon.spy(features.sdkInitializator.v2), }) - codewhispererServiceStub = stubInterface() + codewhispererServiceStub = stubInterface() // @ts-ignore codewhispererServiceStub.client = sinon.stub() codewhispererServiceStub.customizationArn = undefined @@ -141,7 +141,7 @@ describe('Token', () => { const setupServiceManagerWithProfile = async ( profileArn = 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' - ): Promise => { + ): Promise => { setupServiceManager(true) assert.strictEqual(amazonQServiceManager.getState(), 'PENDING_CONNECTION') @@ -161,7 +161,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') - return service + return service as CodeWhispererServiceToken } describe('Initialization process', () => { @@ -245,7 +245,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getState(), 'INITIALIZED') assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert(codewhispererServiceStub.generateSuggestions.calledOnce) }) @@ -318,7 +318,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') assert(codewhispererServiceStub.generateSuggestions.calledOnce) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) }) }) @@ -379,7 +379,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') }) @@ -456,7 +456,7 @@ describe('Token', () => { ) assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - assert(streamingClient1 instanceof StreamingClientService) + assert(streamingClient1 instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1') // Profile change @@ -483,7 +483,7 @@ describe('Token', () => { // CodeWhisperer Service was not recreated assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - assert(streamingClient2 instanceof StreamingClientService) + assert(streamingClient2 instanceof StreamingClientServiceToken) assert.strictEqual(streamingClient1, streamingClient2) assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1') }) @@ -516,7 +516,7 @@ describe('Token', () => { ) assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - assert(streamingClient1 instanceof StreamingClientService) + assert(streamingClient1 instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1') // Profile change @@ -548,7 +548,7 @@ describe('Token', () => { ]) // Streaming Client was recreated - assert(streamingClient2 instanceof StreamingClientService) + assert(streamingClient2 instanceof StreamingClientServiceToken) assert.notStrictEqual(streamingClient1, streamingClient2) assert.strictEqual(await streamingClient2.client.config.region(), 'eu-central-1') }) @@ -582,7 +582,7 @@ describe('Token', () => { ) assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') // Profile change to invalid profile @@ -691,7 +691,7 @@ describe('Token', () => { TEST_ENDPOINT_EU_CENTRAL_1, ]) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'eu-central-1') }) @@ -726,7 +726,7 @@ describe('Token', () => { ) assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1]) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') // Updaing profile @@ -907,7 +907,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') setCredentials('identityCenter') @@ -921,7 +921,7 @@ describe('Token', () => { assert(codewhispererStubFactory.calledTwice) assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) - assert(streamingClient2 instanceof StreamingClientService) + assert(streamingClient2 instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient2.client.config.region(), DEFAULT_AWS_Q_REGION) }) @@ -937,7 +937,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getConnectionType(), 'builderId') assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') setCredentials('identityCenter') @@ -967,7 +967,7 @@ describe('Token', () => { assert.strictEqual(amazonQServiceManager.getConnectionType(), 'identityCenter') assert.strictEqual(amazonQServiceManager.getActiveProfileArn(), undefined) - assert(streamingClient instanceof StreamingClientService) + assert(streamingClient instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') setCredentials('builderId') @@ -981,7 +981,7 @@ describe('Token', () => { assert(codewhispererStubFactory.calledTwice) assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) - assert(streamingClient2 instanceof StreamingClientService) + assert(streamingClient2 instanceof StreamingClientServiceToken) assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1') }) }) @@ -1081,19 +1081,12 @@ describe('IAM', () => { }) it('should initialize the streaming client only once', () => { + // Mock the credentials provider to return credentials when requested setCredentials() - // Mock getIAMCredentialsFromProvider to return dummy credentials - const getIAMCredentialsStub = sinon.stub(utils, 'getIAMCredentialsFromProvider').returns({ - accessKeyId: 'dummy-access-key', - secretAccessKey: 'dummy-secret-key', - sessionToken: 'dummy-session-token', - }) - const streamingClient = serviceManager.getStreamingClient() + // Verify that getting the client again returns the same instance assert.deepStrictEqual(serviceManager.getStreamingClient(), streamingClient) - - getIAMCredentialsStub.restore() }) }) }) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts index e9c1afc2df..089b30d4b9 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts @@ -32,8 +32,13 @@ import { AmazonQDeveloperProfile, signalsAWSQDeveloperProfilesEnabled } from './ import { isStringOrNull } from '../utils' import { getAmazonQRegionAndEndpoint } from './configurationUtils' import { getUserAgent } from '../telemetryUtils' -import { StreamingClientServiceToken, StreamingClientServiceIAM } from '../streamingClientService' +import { + StreamingClientServiceToken, + StreamingClientServiceIAM, + StreamingClientServiceBase, +} from '../streamingClientService' import { parse } from '@aws-sdk/util-arn-parser' +import { CodeWhispererServiceBase } from '../codeWhispererService/codeWhispererServiceBase' /** * AmazonQServiceManager manages state and provides centralized access to @@ -60,7 +65,10 @@ import { parse } from '@aws-sdk/util-arn-parser' * the LSP hand shake. * */ -export class AmazonQServiceManager extends BaseAmazonQServiceManager { +export class AmazonQServiceManager extends BaseAmazonQServiceManager< + CodeWhispererServiceBase, + StreamingClientServiceBase +> { private static instance: AmazonQServiceManager | null = null private enableDeveloperProfileSupport?: boolean private activeIdcProfile?: AmazonQDeveloperProfile @@ -185,33 +193,6 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager CodeWhispererService) { + public setServiceFactory(factory: (region: string, endpoint: string) => CodeWhispererServiceToken) { this.serviceFactory = factory.bind(this) } From eaf49320b56e2cf51ad035f715e0e12b9b8e51e9 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 25 Jul 2025 14:34:09 -0400 Subject: [PATCH 6/8] chore: sync AmazonQServiceManager with AmazonQTokenServiceManager --- .../AmazonQServiceManager.ts | 40 +++++++++++++++---- 1 file changed, 33 insertions(+), 7 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts index 089b30d4b9..e3dccf8cb2 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts @@ -97,6 +97,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< this.state = state } + endpointOverride(): string | undefined { + return this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities + ?.textDocument?.inlineCompletionWithReferences?.endpointOverride + } + public static initInstance(features: QServiceManagerFeatures): AmazonQServiceManager { if (!AmazonQServiceManager.instance) { AmazonQServiceManager.instance = new AmazonQServiceManager(features) @@ -224,6 +229,10 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< return } + const endpointOverride = + this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities + ?.textDocument?.inlineCompletionWithReferences?.endpointOverride + // Connection type changed to 'builderId' if (newConnectionType === 'builderId') { @@ -234,7 +243,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< // region set by client -> runtime region -> default region const clientParams = this.features.lsp.getClientInitializeParams() - this.createCodewhispererServiceInstances('builderId', clientParams?.initializationOptions?.aws?.region) + this.createCodewhispererServiceInstances( + 'builderId', + clientParams?.initializationOptions?.aws?.region, + endpointOverride + ) this.state = 'INITIALIZED' this.log('Initialized Amazon Q service with builderId connection') @@ -256,7 +269,7 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< return } - this.createCodewhispererServiceInstances('identityCenter') + this.createCodewhispererServiceInstances('identityCenter', undefined, endpointOverride) this.state = 'INITIALIZED' this.log('Initialized Amazon Q service with identityCenter connection') @@ -358,7 +371,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< if (!this.activeIdcProfile) { this.activeIdcProfile = newProfile - this.createCodewhispererServiceInstances('identityCenter', newProfile.identityDetails.region) + this.createCodewhispererServiceInstances( + 'identityCenter', + newProfile.identityDetails.region, + this.endpointOverride() + ) this.state = 'INITIALIZED' this.log( `Initialized identityCenter connection to region ${newProfile.identityDetails.region} for profile ${newProfile.arn}` @@ -410,7 +427,11 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< this.activeIdcProfile = newProfile - this.createCodewhispererServiceInstances('identityCenter', newProfile.identityDetails.region) + this.createCodewhispererServiceInstances( + 'identityCenter', + newProfile.identityDetails.region, + this.endpointOverride() + ) this.state = 'INITIALIZED' return @@ -489,7 +510,8 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< private createCodewhispererServiceInstances( connectionType: 'builderId' | 'identityCenter', - clientOrProfileRegion?: string + clientOrProfileRegion: string | undefined, + endpointOverride: string | undefined ) { this.logServiceState('Initializing CodewhispererService') @@ -504,10 +526,14 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< this.region = region this.endpoint = endpoint - this.cachedCodewhispererService = this.serviceFactory(region, endpoint) + if (endpointOverride) { + this.endpoint = endpointOverride + } + + this.cachedCodewhispererService = this.serviceFactory(region, this.endpoint) this.log(`CodeWhispererToken service for connection type ${connectionType} was initialized, region=${region}`) - this.cachedStreamingClient = this.streamingClientFactory(region, endpoint) + this.cachedStreamingClient = this.streamingClientFactory(region, this.endpoint) this.log(`StreamingClient service for connection type ${connectionType} was initialized, region=${region}`) this.logServiceState('CodewhispererService and StreamingClient Initialization finished') From 42bde1efb5bf4ebf53dc7ad4162a378402ec38a0 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Wed, 30 Jul 2025 13:24:33 -0400 Subject: [PATCH 7/8] refactor: move CodeWhispererServiceIAM construction to factory --- .../AmazonQServiceManager.ts | 69 ++++++++++--------- 1 file changed, 37 insertions(+), 32 deletions(-) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts index e3dccf8cb2..6d7ad4d76c 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQServiceManager.ts @@ -451,14 +451,7 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< ) this.region = amazonQRegionAndEndpoint.region this.endpoint = amazonQRegionAndEndpoint.endpoint - this.cachedCodewhispererService = new CodeWhispererServiceIAM( - this.features.credentialsProvider, - this.features.workspace, - this.features.logging, - this.region, - this.endpoint, - this.features.sdkInitializator - ) + this.cachedCodewhispererService = this.serviceFactory(this.region, this.endpoint) this.updateCachedServiceConfig() } this.state = 'INITIALIZED' @@ -485,9 +478,9 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< this.log('Getting instance of CodeWhispererStreaming client') // Trigger checks in token service - const tokenService = this.getCodewhispererService() + const service = this.getCodewhispererService() - if (!tokenService || !this.region || !this.endpoint) { + if (!service || !this.region || !this.endpoint) { throw new AmazonQServiceNotInitializedError() } @@ -545,30 +538,42 @@ export class AmazonQServiceManager extends BaseAmazonQServiceManager< return getUserAgent(initializeParams as InitializeParams, this.features.runtime.serverInfo) } - private serviceFactory(region: string, endpoint: string): CodeWhispererServiceToken { - const service = new CodeWhispererServiceToken( - this.features.credentialsProvider, - this.features.workspace, - this.features.logging, - region, - endpoint, - this.features.sdkInitializator - ) + private serviceFactory(region: string, endpoint: string): CodeWhispererServiceBase { + let service: CodeWhispererServiceBase + if (this.features.credentialsProvider.hasCredentials('iam')) { + service = new CodeWhispererServiceIAM( + this.features.credentialsProvider, + this.features.workspace, + this.features.logging, + region, + endpoint, + this.features.sdkInitializator + ) + } else { + service = new CodeWhispererServiceToken( + this.features.credentialsProvider, + this.features.workspace, + this.features.logging, + region, + endpoint, + this.features.sdkInitializator + ) - const customUserAgent = this.getCustomUserAgent() - service.updateClientConfig({ - customUserAgent: customUserAgent, - }) - service.customizationArn = this.configurationCache.getProperty('customizationArn') - service.profileArn = this.activeIdcProfile?.arn - service.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty( - 'shareCodeWhispererContentWithAWS' - ) + const customUserAgent = this.getCustomUserAgent() + service.updateClientConfig({ + customUserAgent: customUserAgent, + }) + service.customizationArn = this.configurationCache.getProperty('customizationArn') + service.profileArn = this.activeIdcProfile?.arn + service.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty( + 'shareCodeWhispererContentWithAWS' + ) - this.log('Configured CodeWhispererService instance settings:') - this.log( - `customUserAgent=${customUserAgent}, customizationArn=${service.customizationArn}, shareCodeWhispererContentWithAWS=${service.shareCodeWhispererContentWithAWS}` - ) + this.log('Configured CodeWhispererService instance settings:') + this.log( + `customUserAgent=${customUserAgent}, customizationArn=${service.customizationArn}, shareCodeWhispererContentWithAWS=${service.shareCodeWhispererContentWithAWS}` + ) + } return service } From 049c5f912e2d0c516b0250c9c17a60e8db76e6d7 Mon Sep 17 00:00:00 2001 From: Ramon Li Date: Fri, 25 Jul 2025 11:52:44 -0400 Subject: [PATCH 8/8] refactor: change references to AmazonQServiceManager --- .../src/agent-standalone.ts | 12 +- .../src/iam-standalone.ts | 6 +- .../src/iam-webworker.ts | 8 +- .../src/standalone-common.ts | 6 +- .../src/token-standalone.ts | 12 +- server/aws-lsp-codewhisperer/src/index.ts | 2 +- .../agenticChat/agenticChatController.test.ts | 43 +- .../agenticChat/agenticChatController.ts | 19 +- .../agenticChat/qAgenticChatServer.test.ts | 12 +- .../agenticChat/qAgenticChatServer.ts | 14 +- .../chat/chatController.test.ts | 14 +- .../chat/chatSessionService.test.ts | 73 +- .../language-server/chat/qChatServer.test.ts | 2 +- .../src/language-server/chat/qChatServer.ts | 8 +- .../qConfigurationServer.test.ts | 22 +- .../configuration/qConfigurationServer.ts | 8 +- .../inline-completion/codeWhispererServer.ts | 19 +- .../localProjectContextServer.ts | 8 +- .../netTransform/netTransformServer.ts | 6 +- .../tests/transformHandler.test.ts | 4 +- .../netTransform/transformHandler.ts | 6 +- .../codeWhispererSecurityScanServer.ts | 6 +- .../securityScan/securityScanHandler.test.ts | 4 +- .../securityScan/securityScanHandler.ts | 6 +- .../workspaceContextServer.ts | 6 +- .../workspaceFolderManager.test.ts | 6 +- .../workspaceFolderManager.ts | 8 +- .../src/shared/amazonQServer.ts | 6 +- .../AmazonQIAMServiceManager.test.ts | 59 - .../AmazonQIAMServiceManager.ts | 101 -- .../AmazonQTokenServiceManager.test.ts | 1060 ----------------- .../AmazonQTokenServiceManager.ts | 653 ---------- .../qDeveloperProfiles.ts | 6 +- .../shared/amazonQServiceManager/testUtils.ts | 2 +- .../codeWhispererServiceToken.ts | 2 +- .../src/shared/proxy-server.ts | 13 +- .../src/shared/testUtils.ts | 3 - .../src/shared/utils.test.ts | 18 +- tests/aws-lsp-codewhisperer.js | 12 +- 39 files changed, 210 insertions(+), 2065 deletions(-) delete mode 100644 server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.test.ts delete mode 100644 server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts delete mode 100644 server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts delete mode 100644 server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts diff --git a/app/aws-lsp-codewhisperer-runtimes/src/agent-standalone.ts b/app/aws-lsp-codewhisperer-runtimes/src/agent-standalone.ts index 7996472ada..40c03babad 100644 --- a/app/aws-lsp-codewhisperer-runtimes/src/agent-standalone.ts +++ b/app/aws-lsp-codewhisperer-runtimes/src/agent-standalone.ts @@ -1,11 +1,10 @@ import { standalone } from '@aws/language-server-runtimes/runtimes' import { - AmazonQServiceServerIAM, - AmazonQServiceServerToken, + AmazonQServiceServer, CodeWhispererSecurityScanServerTokenProxy, - CodeWhispererServerTokenProxy, - QAgenticChatServerProxy, + CodeWhispererServerProxy, QConfigurationServerTokenProxy, + QAgenticChatServerProxy, QLocalProjectContextServerProxy, QNetTransformServerTokenProxy, WorkspaceContextServerTokenProxy, @@ -25,7 +24,7 @@ const version = versionJson.agenticChat const props = { version: version, servers: [ - CodeWhispererServerTokenProxy, + CodeWhispererServerProxy, CodeWhispererSecurityScanServerTokenProxy, QConfigurationServerTokenProxy, QNetTransformServerTokenProxy, @@ -38,8 +37,7 @@ const props = { WorkspaceContextServerTokenProxy, McpToolsServer, // LspToolsServer, - AmazonQServiceServerIAM, - AmazonQServiceServerToken, + AmazonQServiceServer, ], name: 'AWS CodeWhisperer', } as RuntimeProps diff --git a/app/aws-lsp-codewhisperer-runtimes/src/iam-standalone.ts b/app/aws-lsp-codewhisperer-runtimes/src/iam-standalone.ts index d6c2449e4d..a975499ee8 100644 --- a/app/aws-lsp-codewhisperer-runtimes/src/iam-standalone.ts +++ b/app/aws-lsp-codewhisperer-runtimes/src/iam-standalone.ts @@ -1,7 +1,7 @@ import { standalone } from '@aws/language-server-runtimes/runtimes' -import { CodeWhispererServerIAM, QChatServerIAMProxy } from '@aws/lsp-codewhisperer' -import { createIAMRuntimeProps } from './standalone-common' +import { CodeWhispererServer, QChatServerProxy } from '@aws/lsp-codewhisperer' +import { createRuntimeProps } from './standalone-common' -const props = createIAMRuntimeProps('0.1.0', [CodeWhispererServerIAM, QChatServerIAMProxy]) +const props = createRuntimeProps('0.1.0', [CodeWhispererServer, QChatServerProxy]) standalone(props) diff --git a/app/aws-lsp-codewhisperer-runtimes/src/iam-webworker.ts b/app/aws-lsp-codewhisperer-runtimes/src/iam-webworker.ts index 95cb9bf9bd..8f87cda4c5 100644 --- a/app/aws-lsp-codewhisperer-runtimes/src/iam-webworker.ts +++ b/app/aws-lsp-codewhisperer-runtimes/src/iam-webworker.ts @@ -1,14 +1,14 @@ import { webworker } from '@aws/language-server-runtimes/runtimes/webworker' -import { CodeWhispererServerIAM } from '@aws/lsp-codewhisperer/out/language-server/inline-completion/codeWhispererServer' -import { QChatServerIAM } from '@aws/lsp-codewhisperer/out/language-server/chat/qChatServer' +import { CodeWhispererServer } from '@aws/lsp-codewhisperer/out/language-server/inline-completion/codeWhispererServer' +import { QChatServer } from '@aws/lsp-codewhisperer/out/language-server/chat/qChatServer' import { RuntimeProps } from '@aws/language-server-runtimes/runtimes/runtime' -import { AmazonQServiceServerIAM } from '@aws/lsp-codewhisperer/out/shared/amazonQServer' +import { AmazonQServiceServer } from '@aws/lsp-codewhisperer/out/shared/amazonQServer' // all bundles depend on AmazonQServiceServer, make sure to always include it. The standalone helper // to inject the AmazonQServiceServer does not work for webworker as it triggers missing polyfill errors const props: RuntimeProps = { version: '1.0.0', - servers: [AmazonQServiceServerIAM, CodeWhispererServerIAM, QChatServerIAM], + servers: [AmazonQServiceServer, CodeWhispererServer, QChatServer], name: 'AWS CodeWhisperer', } diff --git a/app/aws-lsp-codewhisperer-runtimes/src/standalone-common.ts b/app/aws-lsp-codewhisperer-runtimes/src/standalone-common.ts index 1867aa47c5..716d5a0ae1 100644 --- a/app/aws-lsp-codewhisperer-runtimes/src/standalone-common.ts +++ b/app/aws-lsp-codewhisperer-runtimes/src/standalone-common.ts @@ -1,7 +1,6 @@ import { RuntimeProps } from '@aws/language-server-runtimes/runtimes/runtime' import { Server } from '@aws/language-server-runtimes/server-interface' -import { AmazonQServiceServerToken } from '@aws/lsp-codewhisperer' -import { AmazonQServiceServerIAM } from '@aws/lsp-codewhisperer' +import { AmazonQServiceServer } from '@aws/lsp-codewhisperer' const createRuntimePropsFactory = (AmazonQServiceServer: Server) => @@ -13,5 +12,4 @@ const createRuntimePropsFactory = } } -export const createIAMRuntimeProps = createRuntimePropsFactory(AmazonQServiceServerIAM) -export const createTokenRuntimeProps = createRuntimePropsFactory(AmazonQServiceServerToken) +export const createRuntimeProps = createRuntimePropsFactory(AmazonQServiceServer) diff --git a/app/aws-lsp-codewhisperer-runtimes/src/token-standalone.ts b/app/aws-lsp-codewhisperer-runtimes/src/token-standalone.ts index 266dd06535..a95010026a 100644 --- a/app/aws-lsp-codewhisperer-runtimes/src/token-standalone.ts +++ b/app/aws-lsp-codewhisperer-runtimes/src/token-standalone.ts @@ -1,27 +1,27 @@ import { standalone } from '@aws/language-server-runtimes/runtimes' import { CodeWhispererSecurityScanServerTokenProxy, - CodeWhispererServerTokenProxy, - QChatServerTokenProxy, + CodeWhispererServerProxy, + QChatServerProxy, QConfigurationServerTokenProxy, QNetTransformServerTokenProxy, QLocalProjectContextServerProxy, WorkspaceContextServerTokenProxy, } from '@aws/lsp-codewhisperer' import { IdentityServer } from '@aws/lsp-identity' -import { createTokenRuntimeProps } from './standalone-common' +import { createRuntimeProps } from './standalone-common' const MAJOR = 0 const MINOR = 1 const PATCH = 0 const VERSION = `${MAJOR}.${MINOR}.${PATCH}` -const props = createTokenRuntimeProps(VERSION, [ - CodeWhispererServerTokenProxy, +const props = createRuntimeProps(VERSION, [ + CodeWhispererServerProxy, CodeWhispererSecurityScanServerTokenProxy, QConfigurationServerTokenProxy, QNetTransformServerTokenProxy, - QChatServerTokenProxy, + QChatServerProxy, IdentityServer.create, QLocalProjectContextServerProxy, WorkspaceContextServerTokenProxy, diff --git a/server/aws-lsp-codewhisperer/src/index.ts b/server/aws-lsp-codewhisperer/src/index.ts index bc74a8472b..8878d8733e 100644 --- a/server/aws-lsp-codewhisperer/src/index.ts +++ b/server/aws-lsp-codewhisperer/src/index.ts @@ -4,4 +4,4 @@ export * from './language-server/chat/qChatServer' export * from './language-server/agenticChat/qAgenticChatServer' export * from './shared/proxy-server' export * from './language-server/netTransform/netTransformServer' -export { AmazonQServiceServerIAM, AmazonQServiceServerToken } from './shared/amazonQServer' +export { AmazonQServiceServer } from './shared/amazonQServer' diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts index bb9ced1528..32d622a76f 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.test.ts @@ -37,7 +37,7 @@ import { TestFeatures } from '@aws/language-server-runtimes/testing' import * as assert from 'assert' import { createIterableResponse, - setCredentialsForAmazonQTokenServiceManagerFactory, + setTokenCredentialsForAmazonQServiceManagerFactory, setIamCredentialsForAmazonQServiceManagerFactory, } from '../../shared/testUtils' import sinon from 'ts-sinon' @@ -49,8 +49,7 @@ import { DocumentContextExtractor } from '../chat/contexts/documentContext' import * as utils from '../chat/utils' import { DEFAULT_HELP_FOLLOW_UP_PROMPT, HELP_MESSAGE } from '../chat/constants' import { TelemetryService } from '../../shared/telemetry/telemetryService' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' -import { AmazonQIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { TabBarController } from './tabBarController' import { getUserPromptsDirectory, promptFileExtension } from './context/contextUtils' import { AdditionalContextProvider } from './context/additionalContextProvider' @@ -174,7 +173,7 @@ describe('AgenticChatController', () => { let emitConversationMetricStub: sinon.SinonStub let testFeatures: TestFeatures - let serviceManager: AmazonQTokenServiceManager + let serviceManager: AmazonQServiceManager let chatSessionManagementService: ChatSessionManagementService let chatController: AgenticChatController let telemetryService: TelemetryService @@ -183,7 +182,7 @@ describe('AgenticChatController', () => { let getMessagesStub: sinon.SinonStub let addMessageStub: sinon.SinonStub - const setSsoCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures) + const setSsoCredentials = setTokenCredentialsForAmazonQServiceManagerFactory(() => testFeatures) const setIamCredentials = setIamCredentialsForAmazonQServiceManagerFactory(() => testFeatures) beforeEach(() => { @@ -286,14 +285,14 @@ describe('AgenticChatController', () => { disposeStub = sinon.stub(ChatSessionService.prototype, 'dispose') sinon.stub(ContextCommandsProvider.prototype, 'maybeUpdateCodeSymbols').resolves() - AmazonQTokenServiceManager.resetInstance() + AmazonQServiceManager.resetInstance() - serviceManager = AmazonQTokenServiceManager.initInstance(testFeatures) + serviceManager = AmazonQServiceManager.initInstance(testFeatures) chatSessionManagementService = ChatSessionManagementService.getInstance() chatSessionManagementService.withAmazonQServiceManager(serviceManager) const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: { @@ -2998,7 +2997,7 @@ ${' '.repeat(8)}} }) describe('onListAvailableModels', () => { - let tokenServiceManagerStub: sinon.SinonStub + let serviceManagerStub: sinon.SinonStub beforeEach(() => { // Create a session with a model ID @@ -3007,16 +3006,16 @@ ${' '.repeat(8)}} session.modelId = 'CLAUDE_3_7_SONNET_20250219_V1_0' // Stub the getRegion method - tokenServiceManagerStub = sinon.stub(AmazonQTokenServiceManager.prototype, 'getRegion') + serviceManagerStub = sinon.stub(AmazonQServiceManager.prototype, 'getRegion') }) afterEach(() => { - tokenServiceManagerStub.restore() + serviceManagerStub.restore() }) it('should return all available models for us-east-1 region', async () => { // Set up the region to be us-east-1 - tokenServiceManagerStub.returns('us-east-1') + serviceManagerStub.returns('us-east-1') // Call the method const params = { tabId: mockTabId } @@ -3035,7 +3034,7 @@ ${' '.repeat(8)}} it('should return limited models for eu-central-1 region', async () => { // Set up the region to be eu-central-1 - tokenServiceManagerStub.returns('eu-central-1') + serviceManagerStub.returns('eu-central-1') // Call the method const params = { tabId: mockTabId } @@ -3054,7 +3053,7 @@ ${' '.repeat(8)}} it('should return all models when region is unknown', async () => { // Set up the region to be unknown - tokenServiceManagerStub.returns('unknown-region') + serviceManagerStub.returns('unknown-region') // Call the method const params = { tabId: mockTabId } @@ -3089,7 +3088,7 @@ ${' '.repeat(8)}} it('should fallback to latest available model when saved model is not available in current region', async () => { // Set up the region to be eu-central-1 (which only has Claude 3.7) - tokenServiceManagerStub.returns('eu-central-1') + serviceManagerStub.returns('eu-central-1') // Mock database to return Claude Sonnet 4 (not available in eu-central-1) const getModelIdStub = sinon.stub(ChatDatabase.prototype, 'getModelId') @@ -3109,7 +3108,7 @@ ${' '.repeat(8)}} it('should use saved model when it is available in current region', async () => { // Set up the region to be us-east-1 (which has both models) - tokenServiceManagerStub.returns('us-east-1') + serviceManagerStub.returns('us-east-1') // Mock database to return Claude 3.7 (available in us-east-1) const getModelIdStub = sinon.stub(ChatDatabase.prototype, 'getModelId') @@ -3129,7 +3128,7 @@ ${' '.repeat(8)}} }) describe('IAM Authentication', () => { - let iamServiceManager: AmazonQIAMServiceManager + let serviceManager: AmazonQServiceManager let iamChatController: AgenticChatController let iamChatSessionManagementService: ChatSessionManagementService @@ -3153,25 +3152,25 @@ ${' '.repeat(8)}} setIamCredentials() // Create IAM service manager - AmazonQIAMServiceManager.resetInstance() - iamServiceManager = AmazonQIAMServiceManager.initInstance(testFeatures) + AmazonQServiceManager.resetInstance() + serviceManager = AmazonQServiceManager.initInstance(testFeatures) // Create chat session management service with IAM service manager iamChatSessionManagementService = ChatSessionManagementService.getInstance() - iamChatSessionManagementService.withAmazonQServiceManager(iamServiceManager) + iamChatSessionManagementService.withAmazonQServiceManager(serviceManager) // Create controller with IAM service manager iamChatController = new AgenticChatController( iamChatSessionManagementService, testFeatures, telemetryService, - iamServiceManager + serviceManager ) }) afterEach(() => { iamChatController.dispose() ChatSessionManagementService.reset() - AmazonQIAMServiceManager.resetInstance() + AmazonQServiceManager.resetInstance() }) it('creates a session with IAM service manager', () => { diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts index 535d17ce56..656f168d99 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts @@ -133,8 +133,7 @@ import { AmazonQServicePendingProfileError, AmazonQServicePendingSigninError, } from '../../shared/amazonQServiceManager/errors' -import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils' import { TabBarController } from './tabBarController' import { ChatDatabase, MaxOverallCharacters, ToolResultValidationError } from './tools/chatDb/chatDb' @@ -223,7 +222,7 @@ import { sanitize } from '@aws/lsp-core/out/util/path' import { getLatestAvailableModel } from './utils/agenticChatControllerHelper' import { ActiveUserTracker } from '../../shared/activeUserTracker' import { UserContext } from '../../client/token/codewhispererbearertokenclient' -import { CodeWhispererServiceToken } from '../../shared/codeWhispererService/codeWhispererServiceToken' +import { CodeWhispererServiceBase } from '../../shared/codeWhispererService/codeWhispererServiceBase' type ChatHandlers = Omit< LspHandlers, @@ -255,7 +254,7 @@ export class AgenticChatController implements ChatHandlers { #triggerContext: AgenticChatTriggerContext #customizationArn?: string #telemetryService: TelemetryService - #serviceManager?: AmazonQBaseServiceManager + #serviceManager?: AmazonQServiceManager #tabBarController: TabBarController #chatHistoryDb: ChatDatabase #additionalContextProvider: AdditionalContextProvider @@ -333,7 +332,7 @@ export class AgenticChatController implements ChatHandlers { chatSessionManagementService: ChatSessionManagementService, features: Features, telemetryService: TelemetryService, - serviceManager?: AmazonQBaseServiceManager + serviceManager?: AmazonQServiceManager ) { this.#features = features this.#chatSessionManagementService = chatSessionManagementService @@ -675,7 +674,7 @@ export class AgenticChatController implements ChatHandlers { } async onListAvailableModels(params: ListAvailableModelsParams): Promise { - const region = AmazonQTokenServiceManager.getInstance().getRegion() + const region = AmazonQServiceManager.getInstance().getRegion() const models = region && MODEL_OPTIONS_FOR_REGION[region] ? MODEL_OPTIONS_FOR_REGION[region] : MODEL_OPTIONS const sessionResult = this.#chatSessionManagementService.getSession(params.tabId) @@ -3473,7 +3472,7 @@ export class AgenticChatController implements ChatHandlers { // In that case, we use the default modelId. let modelId = this.#chatHistoryDb.getModelId() ?? DEFAULT_MODEL_ID - const region = AmazonQTokenServiceManager.getInstance().getRegion() + const region = AmazonQServiceManager.getInstance().getRegion() if (region === 'eu-central-1') { // Only 3.7 Sonnet is available in eu-central-1 for now modelId = 'CLAUDE_3_7_SONNET_20250219_V1_0' @@ -4389,9 +4388,9 @@ export class AgenticChatController implements ChatHandlers { } this.#abTestingFetchingTimeout = setInterval(() => { - let codeWhispererServiceToken: CodeWhispererServiceToken + let codeWhispererService: CodeWhispererServiceBase try { - codeWhispererServiceToken = AmazonQTokenServiceManager.getInstance().getCodewhispererService() + codeWhispererService = AmazonQServiceManager.getInstance().getCodewhispererService() } catch (error) { // getCodewhispererService only returns the cwspr client if the service manager was initialized // i.e. profile was selected otherwise it throws an error @@ -4403,7 +4402,7 @@ export class AgenticChatController implements ChatHandlers { clearInterval(this.#abTestingFetchingTimeout) this.#abTestingFetchingTimeout = undefined - codeWhispererServiceToken + codeWhispererService .listFeatureEvaluations({ userContext }) .then(result => { const feature = result.featureEvaluations?.find( diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.test.ts index a4d1fa2ae9..a476a9bbf3 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.test.ts @@ -9,7 +9,7 @@ import sinon from 'ts-sinon' import { AgenticChatController } from './agenticChatController' import { ChatSessionManagementService } from '../chat/chatSessionManagementService' import { QAgenticChatServer } from './qAgenticChatServer' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' import { Features } from '../types' @@ -21,7 +21,7 @@ describe('QAgenticChatServer', () => { ChatSessionManagementService > let testFeatures: TestFeatures - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let disposeServer: () => void let chatSessionManagementService: ChatSessionManagementService @@ -52,10 +52,10 @@ describe('QAgenticChatServer', () => { } testFeatures.setClientParams(cachedInitializeParams) - AmazonQTokenServiceManager.resetInstance() + AmazonQServiceManager.resetInstance() - AmazonQTokenServiceManager.initInstance(testFeatures) - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + AmazonQServiceManager.initInstance(testFeatures) + amazonQServiceManager = AmazonQServiceManager.getInstance() disposeStub = sinon.stub(ChatSessionManagementService.prototype, 'dispose') chatSessionManagementService = ChatSessionManagementService.getInstance() @@ -75,7 +75,7 @@ describe('QAgenticChatServer', () => { testFeatures.dispose() }) - it('should initialize ChatSessionManagementService with AmazonQTokenServiceManager instance', () => { + it('should initialize ChatSessionManagementService with AmazonQServiceManager instance', () => { sinon.assert.calledOnceWithExactly(withAmazonQServiceSpy, amazonQServiceManager, testFeatures.lsp) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts index 205e5aff69..616a32de1a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts @@ -9,9 +9,10 @@ import { ChatSessionManagementService } from '../chat/chatSessionManagementServi import { CLEAR_QUICK_ACTION, COMPACT_QUICK_ACTION, HELP_QUICK_ACTION } from '../chat/quickActions' import { TelemetryService } from '../../shared/telemetry/telemetryService' import { makeUserContextObject } from '../../shared/telemetryUtils' -import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' -import { getOrThrowBaseTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' -import { getOrThrowBaseIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' +import { + AmazonQServiceManager, + getOrThrowBaseServiceManager, +} from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils' import { TabBarController } from './tabBarController' import { AmazonQServiceInitializationError } from '../../shared/amazonQServiceManager/errors' @@ -45,8 +46,8 @@ export const QAgenticChatServer = (): Server => features => { const { chat, credentialsProvider, telemetry, logging, lsp, runtime, agent } = features - // AmazonQTokenServiceManager and TelemetryService are initialized in `onInitialized` handler to make sure Language Server connection is started - let amazonQServiceManager: AmazonQBaseServiceManager + // AmazonQServiceManager and TelemetryService are initialized in `onInitialized` handler to make sure Language Server connection is started + let amazonQServiceManager: AmazonQServiceManager let telemetryService: TelemetryService let chatController: AgenticChatController @@ -96,8 +97,7 @@ export const QAgenticChatServer = lsp.onInitialized(async () => { // Get initialized service manager and inject it to chatSessionManagementService to pass it down - logging.info(`In IAM Auth mode: ${isUsingIAMAuth()}`) - amazonQServiceManager = isUsingIAMAuth() ? getOrThrowBaseIAMServiceManager() : getOrThrowBaseTokenServiceManager() + amazonQServiceManager = AmazonQServiceManager.getInstance() chatSessionManagementService = ChatSessionManagementService.getInstance().withAmazonQServiceManager(amazonQServiceManager, features.lsp) diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatController.test.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatController.test.ts index 8b78a09e78..fe967df31e 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/chat/chatController.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/chat/chatController.test.ts @@ -15,7 +15,7 @@ import { import { TestFeatures } from '@aws/language-server-runtimes/testing' import * as assert from 'assert' -import { createIterableResponse, setCredentialsForAmazonQTokenServiceManagerFactory } from '../../shared/testUtils' +import { createIterableResponse, setTokenCredentialsForAmazonQServiceManagerFactory } from '../../shared/testUtils' import sinon from 'ts-sinon' import { ChatController } from './chatController' @@ -26,7 +26,7 @@ import { DocumentContextExtractor } from './contexts/documentContext' import * as utils from './utils' import { DEFAULT_HELP_FOLLOW_UP_PROMPT, HELP_MESSAGE } from './constants' import { TelemetryService } from '../../shared/telemetry/telemetryService' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AmazonQError, AmazonQServicePendingProfileError, @@ -102,13 +102,13 @@ describe('ChatController', () => { let emitConversationMetricStub: sinon.SinonStub let testFeatures: TestFeatures - let serviceManager: AmazonQTokenServiceManager + let serviceManager: AmazonQServiceManager let chatSessionManagementService: ChatSessionManagementService let chatController: ChatController let telemetryService: TelemetryService let telemetry: Telemetry - const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures) + const setCredentials = setTokenCredentialsForAmazonQServiceManagerFactory(() => testFeatures) beforeEach(() => { sendMessageStub = sinon.stub(CodeWhispererStreaming.prototype, 'sendMessage').callsFake(() => { @@ -147,15 +147,15 @@ describe('ChatController', () => { disposeStub = sinon.stub(ChatSessionService.prototype, 'dispose') - AmazonQTokenServiceManager.resetInstance() + AmazonQServiceManager.resetInstance() - serviceManager = AmazonQTokenServiceManager.initInstance(testFeatures) + serviceManager = AmazonQServiceManager.initInstance(testFeatures) chatSessionManagementService = ChatSessionManagementService.getInstance() chatSessionManagementService.withAmazonQServiceManager(serviceManager) const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: { diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts index bc776c2f85..e5a27c32de 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts @@ -2,10 +2,9 @@ import { SendMessageCommandInput, SendMessageCommandOutput, ChatTriggerType } fr import * as assert from 'assert' import sinon, { StubbedInstance, stubInterface } from 'ts-sinon' import { ChatSessionService } from './chatSessionService' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' import { StreamingClientServiceToken, StreamingClientServiceIAM } from '../../shared/streamingClientService' import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' -import { AmazonQIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import * as sharedUtils from '../../shared/utils' import { Utils } from 'vscode-uri' import { wrapErrorWithCode } from '../agenticChat/errors' @@ -49,7 +48,7 @@ describe('Chat Session Service', () => { chatSessionService = new ChatSessionService(amazonQServiceManager, mockLsp as any) // needed to identify the stubs as the actual class when checking 'instanceof' in generateAssistantResponse - Object.setPrototypeOf(amazonQServiceManager, AmazonQTokenServiceManager.prototype) + Object.setPrototypeOf(amazonQServiceManager, AmazonQServiceManager.prototype) Object.setPrototypeOf(codeWhispererStreamingClient, StreamingClientServiceToken.prototype) }) @@ -58,7 +57,7 @@ describe('Chat Session Service', () => { }) describe('calling SendMessage', () => { - it('throws error is AmazonQTokenServiceManager is not initialized', async () => { + it('throws error is AmazonQServiceManager is not initialized', async () => { chatSessionService = new ChatSessionService(undefined) await assert.rejects( @@ -119,7 +118,7 @@ describe('Chat Session Service', () => { }) describe('calling GenerateAssistantResponse', () => { - it('throws error is AmazonQTokenServiceManager is not initialized', async () => { + it('throws error is AmazonQServiceManager is not initialized', async () => { chatSessionService = new ChatSessionService(undefined) await assert.rejects( @@ -157,7 +156,7 @@ describe('Chat Session Service', () => { const codeWhispererStreamingClientIAM = stubInterface() codeWhispererStreamingClientIAM.sendMessage.callsFake(() => Promise.resolve(mockRequestResponse)) - const amazonQServiceManagerIAM = stubInterface() + const amazonQServiceManagerIAM = stubInterface() amazonQServiceManagerIAM.getStreamingClient.returns(codeWhispererStreamingClientIAM) const chatSessionServiceIAM = new ChatSessionService(amazonQServiceManagerIAM) @@ -172,7 +171,7 @@ describe('Chat Session Service', () => { const codeWhispererStreamingClientIAM = stubInterface() codeWhispererStreamingClientIAM.sendMessage.callsFake(() => Promise.resolve(mockRequestResponse)) - const amazonQServiceManagerIAM = stubInterface() + const amazonQServiceManagerIAM = stubInterface() amazonQServiceManagerIAM.getStreamingClient.returns(codeWhispererStreamingClientIAM) const chatSessionServiceIAM = new ChatSessionService(amazonQServiceManagerIAM) @@ -216,7 +215,7 @@ describe('Chat Session Service', () => { const codeWhispererStreamingClientIAM = stubInterface() codeWhispererStreamingClientIAM.sendMessage.callsFake(() => Promise.resolve(mockRequestResponse)) - const amazonQServiceManagerIAM = stubInterface() + const amazonQServiceManagerIAM = stubInterface() amazonQServiceManagerIAM.getStreamingClient.returns(codeWhispererStreamingClientIAM) const chatSessionServiceIAM = new ChatSessionService(amazonQServiceManagerIAM) @@ -328,18 +327,62 @@ describe('Chat Session Service', () => { assert.ok(approvedPaths.has(unixPath)) }) }) +}) + +describe('IAM Chat Session Service', () => { + let abortStub: sinon.SinonStub + let chatSessionService: ChatSessionService + let amazonQServiceManager: StubbedInstance + let codeWhispererStreamingClient: StubbedInstance + const mockConversationId = 'mockConversationId' + + const mockRequestParams: SendMessageCommandInput = { + conversationState: { + chatTriggerType: 'MANUAL', + currentMessage: { + userInputMessage: { + content: 'hello', + }, + }, + }, + } + + const mockRequestResponse: SendMessageCommandOutput = { + $metadata: {}, + sendMessageResponse: undefined, + } + + beforeEach(() => { + codeWhispererStreamingClient = stubInterface() + codeWhispererStreamingClient.sendMessage.callsFake(() => Promise.resolve(mockRequestResponse)) + + amazonQServiceManager = stubInterface() + amazonQServiceManager.getStreamingClient.returns(codeWhispererStreamingClient) + + abortStub = sinon.stub(AbortController.prototype, 'abort') + + chatSessionService = new ChatSessionService(amazonQServiceManager) + + // needed to identify the stubs as the actual class when checking 'instanceof' in generateAssistantResponse + Object.setPrototypeOf(amazonQServiceManager, AmazonQServiceManager.prototype) + Object.setPrototypeOf(codeWhispererStreamingClient, StreamingClientServiceIAM.prototype) + }) + + afterEach(() => { + abortStub.restore() + }) describe('IAM client source property', () => { it('sets source to Origin.IDE when using StreamingClientServiceIAM', async () => { const codeWhispererStreamingClientIAM = stubInterface() codeWhispererStreamingClientIAM.sendMessage.callsFake(() => Promise.resolve(mockRequestResponse)) - const amazonQServiceManagerIAM = stubInterface() + const amazonQServiceManagerIAM = stubInterface() amazonQServiceManagerIAM.getStreamingClient.returns(codeWhispererStreamingClientIAM) // Set prototype to make instanceof check work Object.setPrototypeOf(codeWhispererStreamingClientIAM, StreamingClientServiceIAM.prototype) - Object.setPrototypeOf(amazonQServiceManagerIAM, AmazonQIAMServiceManager.prototype) + Object.setPrototypeOf(amazonQServiceManagerIAM, AmazonQServiceManager.prototype) const chatSessionServiceIAM = new ChatSessionService(amazonQServiceManagerIAM) @@ -369,12 +412,12 @@ describe('Chat Session Service', () => { const codeWhispererStreamingClientIAM = stubInterface() codeWhispererStreamingClientIAM.sendMessage.callsFake(() => Promise.resolve(mockRequestResponse)) - const amazonQServiceManagerIAM = stubInterface() + const amazonQServiceManagerIAM = stubInterface() amazonQServiceManagerIAM.getStreamingClient.returns(codeWhispererStreamingClientIAM) // Set prototype to make instanceof check work Object.setPrototypeOf(codeWhispererStreamingClientIAM, StreamingClientServiceIAM.prototype) - Object.setPrototypeOf(amazonQServiceManagerIAM, AmazonQIAMServiceManager.prototype) + Object.setPrototypeOf(amazonQServiceManagerIAM, AmazonQServiceManager.prototype) const chatSessionServiceIAM = new ChatSessionService(amazonQServiceManagerIAM) @@ -541,16 +584,16 @@ describe('Chat Session Service', () => { describe('IAM client error handling', () => { let codeWhispererStreamingClientIAM: StubbedInstance - let amazonQServiceManagerIAM: StubbedInstance + let amazonQServiceManagerIAM: StubbedInstance let chatSessionServiceIAM: ChatSessionService beforeEach(() => { codeWhispererStreamingClientIAM = stubInterface() - amazonQServiceManagerIAM = stubInterface() + amazonQServiceManagerIAM = stubInterface() amazonQServiceManagerIAM.getStreamingClient.returns(codeWhispererStreamingClientIAM) Object.setPrototypeOf(codeWhispererStreamingClientIAM, StreamingClientServiceIAM.prototype) - Object.setPrototypeOf(amazonQServiceManagerIAM, AmazonQIAMServiceManager.prototype) + Object.setPrototypeOf(amazonQServiceManagerIAM, AmazonQServiceManager.prototype) const mockLsp = { getClientInitializeParams: () => ({}), diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.test.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.test.ts index 36b005193b..e2dda88992 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.test.ts @@ -51,7 +51,7 @@ describe('QChatServer', () => { testFeatures.dispose() }) - it('should initialize ChatSessionManagementService with AmazonQTokenServiceManager instance', () => { + it('should initialize ChatSessionManagementService with AmazonQServiceManager instance', () => { sinon.assert.calledOnceWithExactly(withAmazonQServiceSpy, amazonQServiceManager) }) diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.ts index 2452426b1b..504877b588 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/chat/qChatServer.ts @@ -5,8 +5,7 @@ import { CLEAR_QUICK_ACTION, HELP_QUICK_ACTION } from './quickActions' import { TelemetryService } from '../../shared/telemetry/telemetryService' import { makeUserContextObject } from '../../shared/telemetryUtils' import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' -import { getOrThrowBaseTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' -import { getOrThrowBaseIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' +import { getOrThrowBaseServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils' import { AmazonQServiceInitializationError } from '../../shared/amazonQServiceManager/errors' @@ -17,7 +16,7 @@ export const QChatServerFactory = features => { const { chat, credentialsProvider, lsp, telemetry, logging, runtime } = features - // AmazonQTokenServiceManager and TelemetryService are initialized in `onInitialized` handler to make sure Language Server connection is started + // AmazonQServiceManager and TelemetryService are initialized in `onInitialized` handler to make sure Language Server connection is started let amazonQServiceManager: AmazonQBaseServiceManager let telemetryService: TelemetryService @@ -127,5 +126,4 @@ export const QChatServerFactory = } } -export const QChatServerIAM = QChatServerFactory(getOrThrowBaseIAMServiceManager) -export const QChatServerToken = QChatServerFactory(getOrThrowBaseTokenServiceManager) +export const QChatServer = QChatServerFactory(getOrThrowBaseServiceManager) diff --git a/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.test.ts b/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.test.ts index 72decf69c1..3133e6d09a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.test.ts @@ -16,8 +16,8 @@ import { ResponseError, Server, } from '@aws/language-server-runtimes/server-interface' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' -import { setCredentialsForAmazonQTokenServiceManagerFactory } from '../../shared/testUtils' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' +import { setTokenCredentialsForAmazonQServiceManagerFactory } from '../../shared/testUtils' import { Q_CONFIGURATION_SECTION, AWS_Q_ENDPOINTS } from '../../shared/constants' import { AmazonQDeveloperProfile } from '../../shared/amazonQServiceManager/qDeveloperProfiles' @@ -72,7 +72,7 @@ const mockCustomizations = [ describe('QConfigurationServerToken', () => { let testFeatures: TestFeatures - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let listAvailableProfilesStub: sinon.SinonStub let listAvailableCustomizationsStub: sinon.SinonStub let listAllAvailableCustomizationsWithMetadataStub: sinon.SinonStub @@ -82,9 +82,9 @@ describe('QConfigurationServerToken', () => { testFeatures = new TestFeatures() testFeatures.setClientParams(getInitializeParams(customizationsWithMetadata, developerProfiles)) - AmazonQTokenServiceManager.resetInstance() - AmazonQTokenServiceManager.initInstance(testFeatures) - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + AmazonQServiceManager.resetInstance() + AmazonQServiceManager.initInstance(testFeatures) + amazonQServiceManager = AmazonQServiceManager.getInstance() const codeWhispererService = stubInterface() const configurationServer: Server = QConfigurationServerToken() @@ -215,21 +215,21 @@ describe('QConfigurationServerToken', () => { describe('ServerConfigurationProvider', () => { let serverConfigurationProvider: ServerConfigurationProvider - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let codeWhispererService: StubbedInstance let testFeatures: TestFeatures let listAvailableProfilesHandlerSpy: sinon.SinonSpy let tokenSource: CancellationTokenSource let serviceFactoryStub: sinon.SinonStub - const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures) + const setCredentials = setTokenCredentialsForAmazonQServiceManagerFactory(() => testFeatures) const setupServerConfigurationProvider = (developerProfiles = true) => { testFeatures.setClientParams(getInitializeParams(false, developerProfiles)) - AmazonQTokenServiceManager.resetInstance() - AmazonQTokenServiceManager.initInstance(testFeatures) - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + AmazonQServiceManager.resetInstance() + AmazonQServiceManager.initInstance(testFeatures) + amazonQServiceManager = AmazonQServiceManager.getInstance() serviceFactoryStub = sinon.stub().returns(codeWhispererService) amazonQServiceManager.setServiceFactory(serviceFactoryStub) diff --git a/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts b/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts index a2a4de7447..fe078a7e63 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/configuration/qConfigurationServer.ts @@ -14,7 +14,7 @@ import { ListAllAvailableProfilesHandler, } from '../../shared/amazonQServiceManager/qDeveloperProfiles' import { Customization, Customizations } from '../../client/token/codewhispererbearertokenclient' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AWS_Q_ENDPOINTS, Q_CONFIGURATION_SECTION } from '../../shared/constants' import { AmazonQError } from '../../shared/amazonQServiceManager/errors' @@ -54,7 +54,7 @@ type QConfigurationResponse = export const QConfigurationServerToken = (): Server => ({ credentialsProvider, lsp, logging }) => { - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let serverConfigurationProvider: ServerConfigurationProvider let enableCustomizationsWithMetadata = false @@ -103,7 +103,7 @@ export const QConfigurationServerToken = }) lsp.onInitialized(async () => { - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + amazonQServiceManager = AmazonQServiceManager.getInstance() serverConfigurationProvider = new ServerConfigurationProvider( amazonQServiceManager, @@ -186,7 +186,7 @@ export class ServerConfigurationProvider { private listAllAvailableProfilesHandler: ListAllAvailableProfilesHandler constructor( - private serviceManager: AmazonQTokenServiceManager, + private serviceManager: AmazonQServiceManager, private credentialsProvider: CredentialsProvider, private logging: Logging ) { diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts index 5a95a6d424..e99d28cb62 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts @@ -37,11 +37,13 @@ import { AmazonQServiceConnectionExpiredError, AmazonQServiceInitializationError, } from '../../shared/amazonQServiceManager/errors' -import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' -import { getOrThrowBaseTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { + AmazonQBaseServiceManager, + QServiceManagerFeatures, +} from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' +import { getOrThrowBaseServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils' import { hasConnectionExpired } from '../../shared/utils' -import { getOrThrowBaseIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' import { WorkspaceFolderManager } from '../workspaceContext/workspaceFolderManager' import path = require('path') import { getRelativePath } from '../workspaceContext/util' @@ -161,7 +163,7 @@ export const CodewhispererServerFactory = const sessionManager = SessionManager.getInstance() - // AmazonQTokenServiceManager and TelemetryService are initialized in `onInitialized` handler to make sure Language Server connection is started + // AmazonQServiceManager and TelemetryService are initialized in `onInitialized` handler to make sure Language Server connection is started let amazonQServiceManager: AmazonQBaseServiceManager let telemetryService: TelemetryService @@ -895,9 +897,9 @@ export const CodewhispererServerFactory = } logging.debug(`CodePercentageTracker customizationArn updated to ${customizationArn}`) /* - The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination - configuration post all events migration to STE. It'll be replaced by qConfig['enableTelemetryEventsToDestination'] === true - */ + The flag enableTelemetryEventsToDestination is set to true temporarily. It's value will be determined through destination + configuration post all events migration to STE. It'll be replaced by qConfig['enableTelemetryEventsToDestination'] === true + */ // const enableTelemetryEventsToDestination = true // telemetryService.updateEnableTelemetryEventsToDestination(enableTelemetryEventsToDestination) telemetryService.updateOptOutPreference(optOutTelemetryPreference) @@ -1051,5 +1053,4 @@ export const CodewhispererServerFactory = } } -export const CodeWhispererServerIAM = CodewhispererServerFactory(getOrThrowBaseIAMServiceManager) -export const CodeWhispererServerToken = CodewhispererServerFactory(getOrThrowBaseTokenServiceManager) +export const CodeWhispererServer = CodewhispererServerFactory(getOrThrowBaseServiceManager) diff --git a/server/aws-lsp-codewhisperer/src/language-server/localProjectContext/localProjectContextServer.ts b/server/aws-lsp-codewhisperer/src/language-server/localProjectContext/localProjectContextServer.ts index 1fa2b8301b..e1cc669a6a 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/localProjectContext/localProjectContextServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/localProjectContext/localProjectContextServer.ts @@ -1,13 +1,11 @@ import { InitializeParams, Server, TextDocumentSyncKind } from '@aws/language-server-runtimes/server-interface' -import { getOrThrowBaseTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { getOrThrowBaseServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { TelemetryService } from '../../shared/telemetry/telemetryService' import { LocalProjectContextController } from '../../shared/localProjectContextController' import { languageByExtension } from '../../shared/languageDetection' import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils' import { URI } from 'vscode-uri' -import { isUsingIAMAuth } from '../../shared/utils' import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager' -import { getOrThrowBaseIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' export const LocalProjectContextServer = (): Server => @@ -69,9 +67,7 @@ export const LocalProjectContextServer = lsp.onInitialized(async () => { try { - amazonQServiceManager = isUsingIAMAuth() - ? getOrThrowBaseIAMServiceManager() - : getOrThrowBaseTokenServiceManager() + amazonQServiceManager = getOrThrowBaseServiceManager() telemetryService = new TelemetryService(amazonQServiceManager, credentialsProvider, telemetry, logging) await amazonQServiceManager.addDidChangeConfigurationListener(updateConfigurationHandler) diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/netTransformServer.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/netTransformServer.ts index 2b3571576c..64f2e5b35b 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/netTransform/netTransformServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/netTransform/netTransformServer.ts @@ -43,7 +43,7 @@ const CancelTransformCommand = 'aws/qNetTransform/cancelTransform' const DownloadArtifactsCommand = 'aws/qNetTransform/downloadArtifacts' const CancelPollingCommand = 'aws/qNetTransform/cancelPolling' import { SDKInitializator } from '@aws/language-server-runtimes/server-interface' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' /** * @@ -53,7 +53,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A export const QNetTransformServerToken = (): Server => ({ workspace, logging, lsp, telemetry, runtime }) => { - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let transformHandler: TransformHandler const runTransformCommand = async (params: ExecuteCommandParams, _token: CancellationToken) => { @@ -200,7 +200,7 @@ export const QNetTransformServerToken = } const onInitializedHandler = () => { - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + amazonQServiceManager = AmazonQServiceManager.getInstance() transformHandler = new TransformHandler(amazonQServiceManager, workspace, logging, runtime) } diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/tests/transformHandler.test.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/tests/transformHandler.test.ts index f42f351278..25dc27e419 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/netTransform/tests/transformHandler.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/netTransform/tests/transformHandler.test.ts @@ -32,7 +32,7 @@ import { Readable } from 'stream' import { ArtifactManager } from '../artifactManager' import path = require('path') import { IZipEntry } from 'adm-zip' -import { AmazonQTokenServiceManager } from '../../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../../shared/amazonQServiceManager/AmazonQServiceManager' const mocked$Response = { $response: { @@ -64,7 +64,7 @@ describe('Test Transform handler ', () => { workspace = stubInterface() runtime = stubInterface() - const serviceManager = stubInterface() + const serviceManager = stubInterface() client = stubInterface() serviceManager.getCodewhispererService.returns(client) diff --git a/server/aws-lsp-codewhisperer/src/language-server/netTransform/transformHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/netTransform/transformHandler.ts index b0afaa06f1..c5caffa7ac 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/netTransform/transformHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/netTransform/transformHandler.ts @@ -28,18 +28,18 @@ import { import * as validation from './validation' import path = require('path') import AdmZip = require('adm-zip') -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' const workspaceFolderName = 'artifactWorkspace' export class TransformHandler { - private serviceManager: AmazonQTokenServiceManager + private serviceManager: AmazonQServiceManager private workspace: Workspace private logging: Logging private runtime: Runtime private cancelPollingEnabled: Boolean = false - constructor(serviceManager: AmazonQTokenServiceManager, workspace: Workspace, logging: Logging, runtime: Runtime) { + constructor(serviceManager: AmazonQServiceManager, workspace: Workspace, logging: Logging, runtime: Runtime) { this.serviceManager = serviceManager this.workspace = workspace this.logging = logging diff --git a/server/aws-lsp-codewhisperer/src/language-server/securityScan/codeWhispererSecurityScanServer.ts b/server/aws-lsp-codewhisperer/src/language-server/securityScan/codeWhispererSecurityScanServer.ts index 3eca7e77c8..dcd4e0aa75 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/securityScan/codeWhispererSecurityScanServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/securityScan/codeWhispererSecurityScanServer.ts @@ -17,7 +17,7 @@ import { SecurityScanRequestParams, SecurityScanResponse } from './types' import { SecurityScanEvent } from '../../shared/telemetry/types' import { getErrorMessage, parseJson } from '../../shared/utils' import { v4 as uuidv4 } from 'uuid' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { hasConnectionExpired } from '../../shared/utils' import { AmazonQServiceConnectionExpiredError } from '../../shared/amazonQServiceManager/errors' @@ -27,7 +27,7 @@ const CancelSecurityScanCommand = 'aws/codewhisperer/cancelSecurityScan' export const SecurityScanServerToken = (): Server => ({ credentialsProvider, workspace, logging, lsp, telemetry, runtime, sdkInitializator }) => { - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let scanHandler: SecurityScanHandler const diagnosticsProvider = new SecurityScanDiagnosticsProvider(lsp, logging) @@ -241,7 +241,7 @@ export const SecurityScanServerToken = } const onInitializedHandler = async () => { - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + amazonQServiceManager = AmazonQServiceManager.getInstance() scanHandler = new SecurityScanHandler(amazonQServiceManager, workspace, logging) } diff --git a/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.test.ts b/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.test.ts index 7715541476..f933658418 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.test.ts @@ -9,7 +9,7 @@ import { CodeWhispererServiceToken } from '../../shared/codeWhispererService/cod import { SecurityScanHandler } from './securityScanHandler' import { RawCodeScanIssue } from './types' import * as ScanConstants from './constants' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' const mockCodeScanFindings = JSON.stringify([ { @@ -56,7 +56,7 @@ describe('securityScanHandler', () => { const mockedLogging = stubInterface() beforeEach(async () => { // Set up the server with a mock service - const serviceManager = stubInterface() + const serviceManager = stubInterface() client = stubInterface() serviceManager.getCodewhispererService.returns(client) workspace = stubInterface() diff --git a/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.ts b/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.ts index cba587b706..bd10ae44b0 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/securityScan/securityScanHandler.ts @@ -19,15 +19,15 @@ import { } from '../../client/token/codewhispererbearertokenclient' import { sleep } from './dependencyGraph/commonUtil' import { AggregatedCodeScanIssue, RawCodeScanIssue } from './types' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' export class SecurityScanHandler { - private serviceManager: AmazonQTokenServiceManager + private serviceManager: AmazonQServiceManager private workspace: Workspace private logging: Logging public tokenSource: CancellationTokenSource - constructor(serviceManager: AmazonQTokenServiceManager, workspace: Workspace, logging: Logging) { + constructor(serviceManager: AmazonQServiceManager, workspace: Workspace, logging: Logging) { this.serviceManager = serviceManager this.workspace = workspace this.logging = logging diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts index 09336ffb2f..79c4c0e8ba 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceContextServer.ts @@ -19,7 +19,7 @@ import { DependencyDiscoverer } from './dependency/dependencyDiscoverer' import { getCodeWhispererLanguageIdFromPath } from '../../shared/languageDetection' import { makeUserContextObject } from '../../shared/telemetryUtils' import { safeGet } from '../../shared/utils' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { FileUploadJobManager, FileUploadJobType } from './fileUploadJobManager' import { DependencyEvent, DependencyEventBundler } from './dependency/dependencyEventBundler' import ignore = require('ignore') @@ -50,7 +50,7 @@ export const WorkspaceContextServer = (): Server => features => { let isOptedIn: boolean = false let abTestingEvaluated = false let abTestingEnabled = false - let amazonQServiceManager: AmazonQTokenServiceManager + let amazonQServiceManager: AmazonQServiceManager let allowedExtension: string[] = ['AmazonQ-For-VSCode', 'Amazon Q For JetBrains'] let isSupportedExtension = false @@ -231,7 +231,7 @@ export const WorkspaceContextServer = (): Server => features => { if (!isSupportedExtension) { return {} } - amazonQServiceManager = AmazonQTokenServiceManager.getInstance() + amazonQServiceManager = AmazonQServiceManager.getInstance() artifactManager = new ArtifactManager(workspace, logging, workspaceFolders) dependencyDiscoverer = new DependencyDiscoverer(workspace, logging, workspaceFolders, artifactManager) diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts index 9764734254..9c6c895ed7 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.test.ts @@ -1,6 +1,5 @@ import { WorkspaceFolderManager } from './workspaceFolderManager' import sinon, { stubInterface, StubbedInstance } from 'ts-sinon' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' import { CredentialsProvider, Logging } from '@aws/language-server-runtimes/server-interface' import { DependencyDiscoverer } from './dependency/dependencyDiscoverer' import { WorkspaceFolder } from 'vscode-languageserver-protocol' @@ -8,9 +7,10 @@ import { ArtifactManager } from './artifactManager' import { CodeWhispererServiceToken } from '../../shared/codeWhispererService/codeWhispererServiceToken' import { CreateWorkspaceResponse } from '../../client/token/codewhispererbearertokenclient' import { AWSError } from 'aws-sdk' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' describe('WorkspaceFolderManager', () => { - let mockServiceManager: StubbedInstance + let mockServiceManager: StubbedInstance let mockLogging: StubbedInstance let mockCredentialsProvider: StubbedInstance let mockDependencyDiscoverer: StubbedInstance @@ -19,7 +19,7 @@ describe('WorkspaceFolderManager', () => { let workspaceFolderManager: WorkspaceFolderManager beforeEach(() => { - mockServiceManager = stubInterface() + mockServiceManager = stubInterface() mockLogging = stubInterface() mockCredentialsProvider = stubInterface() mockDependencyDiscoverer = stubInterface() diff --git a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts index 00048cd62f..8102f593f7 100644 --- a/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts +++ b/server/aws-lsp-codewhisperer/src/language-server/workspaceContext/workspaceFolderManager.ts @@ -16,7 +16,7 @@ import { uploadArtifactToS3, } from './util' import { DependencyDiscoverer } from './dependency/dependencyDiscoverer' -import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager' +import { AmazonQServiceManager } from '../../shared/amazonQServiceManager/AmazonQServiceManager' import { URI } from 'vscode-uri' import path = require('path') import { isAwsError } from '../../shared/utils' @@ -31,7 +31,7 @@ interface WorkspaceState { type WorkspaceRoot = string export class WorkspaceFolderManager { - private serviceManager: AmazonQTokenServiceManager + private serviceManager: AmazonQServiceManager private logging: Logging private artifactManager: ArtifactManager private dependencyDiscoverer: DependencyDiscoverer @@ -58,7 +58,7 @@ export class WorkspaceFolderManager { private isServiceQuotaExceeded: boolean = false static createInstance( - serviceManager: AmazonQTokenServiceManager, + serviceManager: AmazonQServiceManager, logging: Logging, artifactManager: ArtifactManager, dependencyDiscoverer: DependencyDiscoverer, @@ -85,7 +85,7 @@ export class WorkspaceFolderManager { } private constructor( - serviceManager: AmazonQTokenServiceManager, + serviceManager: AmazonQServiceManager, logging: Logging, artifactManager: ArtifactManager, dependencyDiscoverer: DependencyDiscoverer, diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts index ce00760319..7a6a4b72b6 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServer.ts @@ -6,8 +6,7 @@ import { UpdateConfigurationParams, } from '@aws/language-server-runtimes/server-interface' import { AmazonQBaseServiceManager, QServiceManagerFeatures } from './amazonQServiceManager/BaseAmazonQServiceManager' -import { initBaseIAMServiceManager } from './amazonQServiceManager/AmazonQIAMServiceManager' -import { initBaseTokenServiceManager } from './amazonQServiceManager/AmazonQTokenServiceManager' +import { initBaseServiceManager } from './amazonQServiceManager/AmazonQServiceManager' const LOGGING_PREFIX = '[AMAZON Q SERVER]: ' @@ -65,5 +64,4 @@ export const AmazonQServiceServerFactory = return () => {} } -export const AmazonQServiceServerIAM = AmazonQServiceServerFactory(initBaseIAMServiceManager) -export const AmazonQServiceServerToken = AmazonQServiceServerFactory(initBaseTokenServiceManager) +export const AmazonQServiceServer = AmazonQServiceServerFactory(initBaseServiceManager) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.test.ts deleted file mode 100644 index 02dd270e12..0000000000 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.test.ts +++ /dev/null @@ -1,59 +0,0 @@ -import { TestFeatures } from '@aws/language-server-runtimes/testing' -import { deepStrictEqual } from 'assert' -import sinon from 'ts-sinon' -import { AmazonQIAMServiceManager } from './AmazonQIAMServiceManager' -import { generateSingletonInitializationTests } from './testUtils' -import * as utils from '../utils' - -describe('AmazonQIAMServiceManager', () => { - describe('Initialization process', () => { - generateSingletonInitializationTests(AmazonQIAMServiceManager) - }) - - describe('Service caching', () => { - let serviceManager: AmazonQIAMServiceManager - let features: TestFeatures - let updateCachedServiceConfigSpy: sinon.SinonSpy - - beforeEach(() => { - features = new TestFeatures() - - updateCachedServiceConfigSpy = sinon.spy( - AmazonQIAMServiceManager.prototype, - 'updateCachedServiceConfig' as keyof AmazonQIAMServiceManager - ) - - AmazonQIAMServiceManager.resetInstance() - serviceManager = AmazonQIAMServiceManager.initInstance(features) - }) - - afterEach(() => { - AmazonQIAMServiceManager.resetInstance() - features.dispose() - sinon.restore() - }) - - it('should initialize the CodeWhisperer service only once', () => { - const service = serviceManager.getCodewhispererService() - sinon.assert.calledOnce(updateCachedServiceConfigSpy) - - deepStrictEqual(serviceManager.getCodewhispererService(), service) - sinon.assert.calledOnce(updateCachedServiceConfigSpy) - }) - - it('should initialize the streaming client only once', () => { - // Mock the credentials provider to return credentials when requested - features.credentialsProvider.hasCredentials.withArgs('iam').returns(true) - features.credentialsProvider.getCredentials.withArgs('iam').returns({ - accessKeyId: 'dummy-access-key', - secretAccessKey: 'dummy-secret-key', - sessionToken: 'dummy-session-token', - }) - - const streamingClient = serviceManager.getStreamingClient() - - // Verify that getting the client again returns the same instance - deepStrictEqual(serviceManager.getStreamingClient(), streamingClient) - }) - }) -}) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts deleted file mode 100644 index a7298bb1b5..0000000000 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts +++ /dev/null @@ -1,101 +0,0 @@ -import { CodeWhispererServiceIAM } from '../codeWhispererService/codeWhispererServiceIAM' -import { - AmazonQBaseServiceManager, - BaseAmazonQServiceManager, - QServiceManagerFeatures, -} from './BaseAmazonQServiceManager' -import { getAmazonQRegionAndEndpoint } from './configurationUtils' -import { StreamingClientServiceIAM } from '../streamingClientService' -import { AmazonQServiceAlreadyInitializedError, AmazonQServiceInitializationError } from './errors' -import { - CancellationToken, - CredentialsType, - UpdateConfigurationParams, -} from '@aws/language-server-runtimes/server-interface' - -export class AmazonQIAMServiceManager extends BaseAmazonQServiceManager< - CodeWhispererServiceIAM, - StreamingClientServiceIAM -> { - private static instance: AmazonQIAMServiceManager | null = null - private region: string - private endpoint: string - - private constructor(features: QServiceManagerFeatures) { - super(features) - const amazonQRegionAndEndpoint = getAmazonQRegionAndEndpoint(features.runtime, features.logging) - this.region = amazonQRegionAndEndpoint.region - this.endpoint = amazonQRegionAndEndpoint.endpoint - } - - public static initInstance(features: QServiceManagerFeatures): AmazonQIAMServiceManager { - if (!AmazonQIAMServiceManager.instance) { - AmazonQIAMServiceManager.instance = new AmazonQIAMServiceManager(features) - - return AmazonQIAMServiceManager.instance - } - - throw new AmazonQServiceAlreadyInitializedError() - } - - public static getInstance(): AmazonQIAMServiceManager { - if (!AmazonQIAMServiceManager.instance) { - throw new AmazonQServiceInitializationError( - 'Amazon Q service has not been initialized yet. Make sure the Amazon Q service server is present and properly initialized.' - ) - } - - return AmazonQIAMServiceManager.instance - } - - public getCodewhispererService() { - if (!this.cachedCodewhispererService) { - this.cachedCodewhispererService = new CodeWhispererServiceIAM( - this.features.credentialsProvider, - this.features.workspace, - this.features.logging, - this.region, - this.endpoint, - this.features.sdkInitializator - ) - - this.updateCachedServiceConfig() - } - - return this.cachedCodewhispererService - } - - public getStreamingClient() { - if (!this.cachedStreamingClient) { - this.cachedStreamingClient = new StreamingClientServiceIAM( - this.features.credentialsProvider, - this.features.sdkInitializator, - this.features.logging, - this.region, - this.endpoint - ) - } - return this.cachedStreamingClient - } - - public handleOnCredentialsDeleted(_type: CredentialsType): void { - return - } - - public override handleOnUpdateConfiguration( - _params: UpdateConfigurationParams, - _token: CancellationToken - ): Promise { - return Promise.resolve() - } - - // For Unit Tests - public static resetInstance(): void { - AmazonQIAMServiceManager.instance = null - } -} - -export const initBaseIAMServiceManager = (features: QServiceManagerFeatures) => - AmazonQIAMServiceManager.initInstance(features) - -export const getOrThrowBaseIAMServiceManager = (): AmazonQBaseServiceManager => AmazonQIAMServiceManager.getInstance() diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts deleted file mode 100644 index 9cf10fc623..0000000000 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.test.ts +++ /dev/null @@ -1,1060 +0,0 @@ -import * as assert from 'assert' -import sinon, { StubbedInstance, stubInterface } from 'ts-sinon' -import { AmazonQTokenServiceManager } from './AmazonQTokenServiceManager' -import { TestFeatures } from '@aws/language-server-runtimes/testing' -import { GenerateSuggestionsRequest } from '../codeWhispererService/codeWhispererServiceBase' -import { CodeWhispererServiceToken } from '../codeWhispererService/codeWhispererServiceToken' -import { - AmazonQServiceInitializationError, - AmazonQServicePendingProfileError, - AmazonQServicePendingProfileUpdateError, - AmazonQServicePendingSigninError, -} from './errors' -import { - CancellationToken, - InitializeParams, - LSPErrorCodes, - ResponseError, -} from '@aws/language-server-runtimes/protocol' -import { - AWS_Q_ENDPOINT_URL_ENV_VAR, - AWS_Q_ENDPOINTS, - AWS_Q_REGION_ENV_VAR, - DEFAULT_AWS_Q_ENDPOINT_URL, - DEFAULT_AWS_Q_REGION, -} from '../constants' -import * as qDeveloperProfilesFetcherModule from './qDeveloperProfiles' -import { setCredentialsForAmazonQTokenServiceManagerFactory } from '../testUtils' -import { StreamingClientServiceToken } from '../streamingClientService' -import { generateSingletonInitializationTests } from './testUtils' - -export const mockedProfiles: qDeveloperProfilesFetcherModule.AmazonQDeveloperProfile[] = [ - { - arn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - name: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - identityDetails: { - region: 'us-east-1', - }, - }, - { - arn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2', - name: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2', - identityDetails: { - region: 'us-east-1', - }, - }, - { - arn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', - name: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', - identityDetails: { - region: 'eu-central-1', - }, - }, -] - -const TEST_ENDPOINT_US_EAST_1 = 'http://amazon-q-in-us-east-1-endpoint' -const TEST_ENDPOINT_EU_CENTRAL_1 = 'http://amazon-q-in-eu-central-1-endpoint' - -describe('AmazonQTokenServiceManager', () => { - let codewhispererServiceStub: StubbedInstance - let codewhispererStubFactory: sinon.SinonStub> - let sdkInitializatorSpy: sinon.SinonSpy - let getListAllAvailableProfilesHandlerStub: sinon.SinonStub - - let amazonQTokenServiceManager: AmazonQTokenServiceManager - let features: TestFeatures - - beforeEach(() => { - // Override endpoints for testing - AWS_Q_ENDPOINTS.set('us-east-1', TEST_ENDPOINT_US_EAST_1) - AWS_Q_ENDPOINTS.set('eu-central-1', TEST_ENDPOINT_EU_CENTRAL_1) - - getListAllAvailableProfilesHandlerStub = sinon - .stub() - .resolves( - Promise.resolve(mockedProfiles).then(() => - new Promise(resolve => setTimeout(resolve, 1)).then(() => mockedProfiles) - ) - ) - - sinon - .stub(qDeveloperProfilesFetcherModule, 'getListAllAvailableProfilesHandler') - .returns(getListAllAvailableProfilesHandlerStub) - - AmazonQTokenServiceManager.resetInstance() - - features = new TestFeatures() - - sdkInitializatorSpy = Object.assign(sinon.spy(features.sdkInitializator), { - v2: sinon.spy(features.sdkInitializator.v2), - }) - - codewhispererServiceStub = stubInterface() - // @ts-ignore - codewhispererServiceStub.client = sinon.stub() - codewhispererServiceStub.customizationArn = undefined - codewhispererServiceStub.shareCodeWhispererContentWithAWS = false - codewhispererServiceStub.profileArn = undefined - - // Initialize the class with mocked dependencies - codewhispererStubFactory = sinon.stub().returns(codewhispererServiceStub) - }) - - afterEach(() => { - AmazonQTokenServiceManager.resetInstance() - features.dispose() - sinon.restore() - }) - - const setupServiceManager = (enableProfiles = false) => { - // @ts-ignore - const cachedInitializeParams: InitializeParams = { - initializationOptions: { - aws: { - awsClientCapabilities: { - q: { - developerProfiles: enableProfiles, - }, - }, - }, - }, - } - features.setClientParams(cachedInitializeParams) - - AmazonQTokenServiceManager.initInstance(features) - amazonQTokenServiceManager = AmazonQTokenServiceManager.getInstance() - amazonQTokenServiceManager.setServiceFactory(codewhispererStubFactory) - } - - const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => features) - - const clearCredentials = () => { - features.credentialsProvider.hasCredentials.returns(false) - features.credentialsProvider.getCredentials.returns(undefined) - features.credentialsProvider.getConnectionType.returns('none') - } - - const setupServiceManagerWithProfile = async ( - profileArn = 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' - ): Promise => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: profileArn, - }, - }, - {} as CancellationToken - ) - - const service = amazonQTokenServiceManager.getCodewhispererService() - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - - return service - } - - describe('Initialization process', () => { - generateSingletonInitializationTests(AmazonQTokenServiceManager) - }) - - describe('Client is not connected', () => { - it('should be in PENDING_CONNECTION state when bearer token is not set', () => { - setupServiceManager() - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - clearCredentials() - - assert.throws(() => amazonQTokenServiceManager.getCodewhispererService(), AmazonQServicePendingSigninError) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none') - }) - }) - - describe('Clear state upon bearer token deletion', () => { - let cancelActiveProfileChangeTokenSpy: sinon.SinonSpy - - beforeEach(() => { - setupServiceManager() - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - cancelActiveProfileChangeTokenSpy = sinon.spy( - amazonQTokenServiceManager as any, - 'cancelActiveProfileChangeToken' - ) - - setCredentials('builderId') - }) - - it('should clear local state variables on receiving bearer token deletion event', () => { - amazonQTokenServiceManager.getCodewhispererService() - - amazonQTokenServiceManager.handleOnCredentialsDeleted('bearer') - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none') - assert.strictEqual((amazonQTokenServiceManager as any)['cachedCodewhispererService'], undefined) - assert.strictEqual((amazonQTokenServiceManager as any)['cachedStreamingClient'], undefined) - assert.strictEqual((amazonQTokenServiceManager as any)['activeIdcProfile'], undefined) - sinon.assert.calledOnce(cancelActiveProfileChangeTokenSpy) - }) - - it('should not clear local state variables on receiving iam token deletion event', () => { - amazonQTokenServiceManager.getCodewhispererService() - - amazonQTokenServiceManager.handleOnCredentialsDeleted('iam') - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'builderId') - assert(!(amazonQTokenServiceManager['cachedCodewhispererService'] === undefined)) - assert.strictEqual((amazonQTokenServiceManager as any)['activeIdcProfile'], undefined) - sinon.assert.notCalled(cancelActiveProfileChangeTokenSpy) - }) - }) - - describe('BuilderId support', () => { - const testRegion = 'some-region' - const testEndpoint = 'http://some-endpoint-in-some-region' - - beforeEach(() => { - setupServiceManager() - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('builderId') - - AWS_Q_ENDPOINTS.set(testRegion, testEndpoint) - - features.lsp.getClientInitializeParams.reset() - }) - - it('should be INITIALIZED with BuilderId Connection', async () => { - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'builderId') - - assert(streamingClient instanceof StreamingClientServiceToken) - assert(codewhispererServiceStub.generateSuggestions.calledOnce) - }) - - it('should initialize service with region set by client', async () => { - features.setClientParams({ - processId: 0, - rootUri: 'some-root-uri', - capabilities: {}, - initializationOptions: { - aws: { - region: testRegion, - }, - }, - }) - - amazonQTokenServiceManager.getCodewhispererService() - assert(codewhispererStubFactory.calledOnceWithExactly(testRegion, testEndpoint)) - - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - assert.strictEqual(await streamingClient.client.config.region(), testRegion) - assert.strictEqual( - (await streamingClient.client.config.endpoint()).hostname, - 'some-endpoint-in-some-region' - ) - }) - - it('should initialize service with region set by runtime if not set by client', async () => { - features.runtime.getConfiguration.withArgs(AWS_Q_REGION_ENV_VAR).returns('eu-central-1') - features.runtime.getConfiguration.withArgs(AWS_Q_ENDPOINT_URL_ENV_VAR).returns(TEST_ENDPOINT_EU_CENTRAL_1) - - amazonQTokenServiceManager.getCodewhispererService() - assert(codewhispererStubFactory.calledOnceWithExactly('eu-central-1', TEST_ENDPOINT_EU_CENTRAL_1)) - - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - assert.strictEqual(await streamingClient.client.config.region(), 'eu-central-1') - assert.strictEqual( - (await streamingClient.client.config.endpoint()).hostname, - 'amazon-q-in-eu-central-1-endpoint' - ) - }) - - it('should initialize service with default region if not set by client and runtime', async () => { - amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - - assert(codewhispererStubFactory.calledOnceWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) - - assert.strictEqual(await streamingClient.client.config.region(), DEFAULT_AWS_Q_REGION) - assert.strictEqual( - (await streamingClient.client.config.endpoint()).hostname, - 'codewhisperer.us-east-1.amazonaws.com' - ) - }) - }) - - describe('IdentityCenter support', () => { - describe('Developer Profiles Support is disabled', () => { - it('should be INITIALIZED with IdentityCenter Connection', async () => { - setupServiceManager() - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert(codewhispererServiceStub.generateSuggestions.calledOnce) - - assert(streamingClient instanceof StreamingClientServiceToken) - }) - }) - - describe('Developer Profiles Support is enabled', () => { - it('should not throw when receiving null profile arn in PENDING_CONNECTION state', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - await assert.doesNotReject( - amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: null, - }, - }, - {} as CancellationToken - ) - ) - - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - }) - - it('should initialize to PENDING_Q_PROFILE state when IdentityCenter Connection is set', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileError - ) - assert.throws(() => amazonQTokenServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - }) - - it('handles Profile configuration request for valid profile and initializes to INITIALIZED state', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') - }) - - it('handles Profile configuration request for valid profile & cancels the old in-flight update request', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - assert.strictEqual((amazonQTokenServiceManager as any)['profileChangeTokenSource'], undefined) - - let firstRequestStarted = false - const originalHandleProfileChange = amazonQTokenServiceManager['handleProfileChange'] - amazonQTokenServiceManager['handleProfileChange'] = async (...args) => { - firstRequestStarted = true - return originalHandleProfileChange.apply(amazonQTokenServiceManager, args) - } - const firstUpdate = amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - while (!firstRequestStarted) { - await new Promise(resolve => setTimeout(resolve, 1)) - } - const secondUpdate = amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - const results = await Promise.allSettled([firstUpdate, secondUpdate]) - - assert.strictEqual((amazonQTokenServiceManager as any)['profileChangeTokenSource'], undefined) - const service = amazonQTokenServiceManager.getCodewhispererService() - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - - assert.strictEqual(results[0].status, 'fulfilled') - assert.strictEqual(results[1].status, 'fulfilled') - }) - - it('handles Profile configuration change to valid profile in same region', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient1 = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' - ) - - assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - assert(streamingClient1 instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1') - - // Profile change - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2', - }, - }, - {} as CancellationToken - ) - await service.generateSuggestions({} as GenerateSuggestionsRequest) - const streamingClient2 = amazonQTokenServiceManager.getStreamingClient() - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ-2' - ) - - // CodeWhisperer Service was not recreated - assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - - assert(streamingClient2 instanceof StreamingClientServiceToken) - assert.strictEqual(streamingClient1, streamingClient2) - assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1') - }) - - it('handles Profile configuration change to valid profile in different region', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient1 = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' - ) - assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - - assert(streamingClient1 instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient1.client.config.region(), 'us-east-1') - - // Profile change - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - await service.generateSuggestions({} as GenerateSuggestionsRequest) - const streamingClient2 = amazonQTokenServiceManager.getStreamingClient() - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ' - ) - - // CodeWhisperer Service was recreated - assert(codewhispererStubFactory.calledTwice) - assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, [ - 'eu-central-1', - TEST_ENDPOINT_EU_CENTRAL_1, - ]) - - // Streaming Client was recreated - assert(streamingClient2 instanceof StreamingClientServiceToken) - assert.notStrictEqual(streamingClient1, streamingClient2) - assert.strictEqual(await streamingClient2.client.config.region(), 'eu-central-1') - }) - - // As we're not validating profile at this moment, there is no "invalid" profile - it.skip('handles Profile configuration change from valid to invalid profile', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - let service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' - ) - assert(codewhispererStubFactory.calledOnceWithExactly('us-east-1', TEST_ENDPOINT_US_EAST_1)) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') - - // Profile change to invalid profile - - await assert.rejects( - amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/invalid-profile-arn', - }, - }, - {} as CancellationToken - ), - new ResponseError(LSPErrorCodes.RequestFailed, 'Requested Amazon Q Profile does not exist', { - awsErrorCode: 'E_AMAZON_Q_INVALID_PROFILE', - }) - ) - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileError - ) - assert.throws(() => amazonQTokenServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - // CodeWhisperer Service was not recreated - assert(codewhispererStubFactory.calledOnce) - assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1]) - }) - - // As we're not validating profile at this moment, there is no "non-existing" profile - it.skip('handles non-existing profile selection', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await assert.rejects( - amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/invalid-profile-arn', - }, - }, - {} as CancellationToken - ), - new ResponseError(LSPErrorCodes.RequestFailed, 'Requested Amazon Q Profile does not exist', { - awsErrorCode: 'E_AMAZON_Q_INVALID_PROFILE', - }) - ) - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileError - ) - assert.throws(() => amazonQTokenServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(codewhispererStubFactory.notCalled) - }) - - it('prevents service usage while profile change is inflight when profile was not set', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileError - ) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - - amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE') - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileUpdateError - ) - assert.throws( - () => amazonQTokenServiceManager.getStreamingClient(), - AmazonQServicePendingProfileUpdateError - ) - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:eu-central-1:11111111111111:profile/QQQQQQQQQQQQ' - ) - assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, [ - 'eu-central-1', - TEST_ENDPOINT_EU_CENTRAL_1, - ]) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'eu-central-1') - }) - - it('prevents service usage while profile change is inflight when profile was set before', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileError - ) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - const service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual( - amazonQTokenServiceManager.getActiveProfileArn(), - 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ' - ) - assert.deepStrictEqual(codewhispererStubFactory.lastCall.args, ['us-east-1', TEST_ENDPOINT_US_EAST_1]) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') - - // Updaing profile - amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE') - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileUpdateError - ) - assert.throws( - () => amazonQTokenServiceManager.getStreamingClient(), - AmazonQServicePendingProfileUpdateError - ) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') - }) - - it('resets to PENDING_PROFILE from INITIALIZED when receiving null profileArn', async () => { - await setupServiceManagerWithProfile() - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: null, - }, - }, - {} as CancellationToken - ) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests) - }) - - it('resets to PENDING_Q_PROFILE from PENDING_Q_PROFILE_UPDATE when receiving null profileArn', async () => { - await setupServiceManagerWithProfile() - - amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE') - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') - - // Null profile arn - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: null, - }, - }, - {} as CancellationToken - ) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests) - assert.throws(() => amazonQTokenServiceManager.getCodewhispererService()) - }) - - it('cancels on-going profile update when credentials are deleted', async () => { - await setupServiceManagerWithProfile() - - amazonQTokenServiceManager.setState('PENDING_Q_PROFILE_UPDATE') - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE_UPDATE') - - amazonQTokenServiceManager.handleOnCredentialsDeleted('bearer') - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - sinon.assert.calledOnce(codewhispererServiceStub.abortInflightRequests) - assert.throws(() => amazonQTokenServiceManager.getCodewhispererService()) - }) - - // Due to service limitation, validation was removed for the sake of recovering API availability - // When service is ready to take more tps, revert https://github.com/aws/language-servers/pull/1329 to add profile validation - it('should not call service to validate profile and always assume its validness', async () => { - setupServiceManager(true) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - sinon.assert.notCalled(getListAllAvailableProfilesHandlerStub) - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - }) - }) - }) - - describe('Connection types with no Developer Profiles support', () => { - it('handles reauthentication scenario when connection type is none but profile ARN is provided', async () => { - setupServiceManager(true) - clearCredentials() - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ) - - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - }) - - it('ignores null profile when connection type is none', async () => { - setupServiceManager(true) - clearCredentials() - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none') - - await amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: null, - }, - }, - {} as CancellationToken - ) - - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'none') - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_CONNECTION') - }) - - it('returns error when profile update is requested and connection type is builderId', async () => { - setupServiceManager(true) - setCredentials('builderId') - - await assert.rejects( - amazonQTokenServiceManager.handleOnUpdateConfiguration( - { - section: 'aws.q', - settings: { - profileArn: 'arn:aws:testprofilearn:us-east-1:11111111111111:profile/QQQQQQQQQQQQ', - }, - }, - {} as CancellationToken - ), - new ResponseError( - LSPErrorCodes.RequestFailed, - 'Connection type builderId does not support Developer Profiles feature.', - { - awsErrorCode: 'E_AMAZON_Q_CONNECTION_NO_PROFILE_SUPPORT', - } - ) - ) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'builderId') - }) - }) - - describe('Handle connection type changes', () => { - describe('connection changes from BuilderId to IdentityCenter', () => { - it('should initialize service with default region when profile support is disabled', async () => { - setupServiceManager(false) - setCredentials('builderId') - - let service1 = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service1.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'builderId') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') - - setCredentials('identityCenter') - let service2 = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient2 = amazonQTokenServiceManager.getStreamingClient() - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(codewhispererStubFactory.calledTwice) - assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) - - assert(streamingClient2 instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient2.client.config.region(), DEFAULT_AWS_Q_REGION) - }) - - it('should initialize service to PENDING_Q_PROFILE state when profile support is enabled', async () => { - setupServiceManager(true) - setCredentials('builderId') - - let service = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'builderId') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') - - setCredentials('identityCenter') - - assert.throws( - () => amazonQTokenServiceManager.getCodewhispererService(), - AmazonQServicePendingProfileError - ) - assert.throws(() => amazonQTokenServiceManager.getStreamingClient(), AmazonQServicePendingProfileError) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'PENDING_Q_PROFILE') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(codewhispererStubFactory.calledOnce) - assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) - }) - }) - - describe('connection changes from IdentityCenter to BuilderId', () => { - it('should initialize service in default IAD region', async () => { - setupServiceManager(false) - setCredentials('identityCenter') - - let service1 = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient = amazonQTokenServiceManager.getStreamingClient() - await service1.generateSuggestions({} as GenerateSuggestionsRequest) - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'identityCenter') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(streamingClient instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient.client.config.region(), 'us-east-1') - - setCredentials('builderId') - let service2 = amazonQTokenServiceManager.getCodewhispererService() - const streamingClient2 = amazonQTokenServiceManager.getStreamingClient() - - assert.strictEqual(amazonQTokenServiceManager.getState(), 'INITIALIZED') - assert.strictEqual(amazonQTokenServiceManager.getConnectionType(), 'builderId') - assert.strictEqual(amazonQTokenServiceManager.getActiveProfileArn(), undefined) - - assert(codewhispererStubFactory.calledTwice) - assert(codewhispererStubFactory.calledWithExactly(DEFAULT_AWS_Q_REGION, DEFAULT_AWS_Q_ENDPOINT_URL)) - - assert(streamingClient2 instanceof StreamingClientServiceToken) - assert.strictEqual(await streamingClient2.client.config.region(), 'us-east-1') - }) - }) - }) - - describe('handle LSP Configuration settings', () => { - it('should initialize codewhisperer service with default configurations when not set by client', async () => { - setupServiceManager() - setCredentials('identityCenter') - - await amazonQTokenServiceManager.handleDidChangeConfiguration() - - const service = amazonQTokenServiceManager.getCodewhispererService() - - assert.strictEqual(service.customizationArn, undefined) - assert.strictEqual(service.shareCodeWhispererContentWithAWS, false) - }) - - it('should returned configured codewhispererService with expected configuration values', async () => { - const getConfigStub = features.lsp.workspace.getConfiguration - getConfigStub.withArgs('aws.q').resolves({ - customization: 'test-customization-arn', - optOutTelemetryPreference: true, - }) - getConfigStub.withArgs('aws.codeWhisperer').resolves({ - includeSuggestionsWithCodeReferences: true, - shareCodeWhispererContentWithAWS: true, - }) - - // Initialize mock server - setupServiceManager() - setCredentials('identityCenter') - - amazonQTokenServiceManager = AmazonQTokenServiceManager.getInstance() - const service = amazonQTokenServiceManager.getCodewhispererService() - - assert.strictEqual(service.customizationArn, undefined) - assert.strictEqual(service.shareCodeWhispererContentWithAWS, false) - - await amazonQTokenServiceManager.handleDidChangeConfiguration() - - // Force next tick to allow async work inside handleDidChangeConfiguration to complete - await Promise.resolve() - - assert.strictEqual(service.customizationArn, 'test-customization-arn') - assert.strictEqual(service.shareCodeWhispererContentWithAWS, true) - }) - }) - - describe('Initialize', () => { - it('should throw when initialize is called before LSP has been initialized with InitializeParams', () => { - features.resetClientParams() - - assert.throws(() => AmazonQTokenServiceManager.initInstance(features), AmazonQServiceInitializationError) - }) - }) -}) diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts deleted file mode 100644 index 7b182f85dd..0000000000 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQTokenServiceManager.ts +++ /dev/null @@ -1,653 +0,0 @@ -import { - UpdateConfigurationParams, - ResponseError, - LSPErrorCodes, - SsoConnectionType, - CancellationToken, - CredentialsType, - InitializeParams, - CancellationTokenSource, -} from '@aws/language-server-runtimes/server-interface' -import { CodeWhispererServiceToken } from '../codeWhispererService/codeWhispererServiceToken' -import { - AmazonQError, - AmazonQServiceAlreadyInitializedError, - AmazonQServiceInitializationError, - AmazonQServiceInvalidProfileError, - AmazonQServiceNoProfileSupportError, - AmazonQServiceNotInitializedError, - AmazonQServicePendingProfileError, - AmazonQServicePendingProfileUpdateError, - AmazonQServicePendingSigninError, - AmazonQServiceProfileUpdateCancelled, -} from './errors' -import { - AmazonQBaseServiceManager, - BaseAmazonQServiceManager, - QServiceManagerFeatures, -} from './BaseAmazonQServiceManager' -import { AWS_Q_ENDPOINTS, Q_CONFIGURATION_SECTION } from '../constants' -import { AmazonQDeveloperProfile, signalsAWSQDeveloperProfilesEnabled } from './qDeveloperProfiles' -import { isStringOrNull } from '../utils' -import { getAmazonQRegionAndEndpoint } from './configurationUtils' -import { getUserAgent } from '../telemetryUtils' -import { StreamingClientServiceToken } from '../streamingClientService' -import { parse } from '@aws-sdk/util-arn-parser' - -/** - * AmazonQTokenServiceManager manages state and provides centralized access to - * instance of CodeWhispererServiceToken SDK client to any consuming code. - * It ensures that CodeWhispererServiceToken is configured to always access correct regionalized Amazon Q Developer API endpoint. - * Regional endppoint is selected based on: - * 1) current SSO auth connection type (BuilderId or IDC). - * 2) selected Amazon Q Developer profile (only for IDC connection type). - * - * @states - * - PENDING_CONNECTION: Initial state when no bearer token is set - * - PENDING_Q_PROFILE: When using Identity Center and waiting for profile selection - * - PENDING_Q_PROFILE_UPDATE: During profile update operation - * - INITIALIZED: Service is ready to handle requests - * - * @connectionTypes - * - none: No active connection - * - builderId: Connected via Builder ID - * - identityCenter: Connected via Identity Center - * - * AmazonQTokenServiceManager is a singleton class, which must be instantiated with Language Server runtimes [Features](https://github.com/aws/language-server-runtimes/blob/21d5d1dc7c73499475b7c88c98d2ce760e5d26c8/runtimes/server-interface/server.ts#L31-L42) - * in the `AmazonQServiceServer` via the `initBaseTokenServiceManager` factory. Dependencies of this class can access the singleton via - * the `getOrThrowBaseTokenServiceManager` factory or `getInstance()` method after the initialized notification has been received during - * the LSP hand shake. - * - */ -export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager< - CodeWhispererServiceToken, - StreamingClientServiceToken -> { - private static instance: AmazonQTokenServiceManager | null = null - private enableDeveloperProfileSupport?: boolean - private activeIdcProfile?: AmazonQDeveloperProfile - private connectionType?: SsoConnectionType - private profileChangeTokenSource: CancellationTokenSource | undefined - private region?: string - private endpoint?: string - private regionChangeListeners: Array<(region: string) => void> = [] - /** - * Internal state of Service connection, based on status of bearer token and Amazon Q Developer profile selection. - * Supported states: - * PENDING_CONNECTION - Waiting for Bearer Token and StartURL to be passed - * PENDING_Q_PROFILE - (only for identityCenter connection) waiting for setting Developer Profile - * PENDING_Q_PROFILE_UPDATE (only for identityCenter connection) waiting for Developer Profile to complete - * INITIALIZED - Service is initialized - */ - private state: 'PENDING_CONNECTION' | 'PENDING_Q_PROFILE' | 'PENDING_Q_PROFILE_UPDATE' | 'INITIALIZED' = - 'PENDING_CONNECTION' - - private constructor(features: QServiceManagerFeatures) { - super(features) - } - - // @VisibleForTesting, please DO NOT use in production - setState(state: 'PENDING_CONNECTION' | 'PENDING_Q_PROFILE' | 'PENDING_Q_PROFILE_UPDATE' | 'INITIALIZED') { - this.state = state - } - - endpointOverride(): string | undefined { - return this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities - ?.textDocument?.inlineCompletionWithReferences?.endpointOverride - } - - public static initInstance(features: QServiceManagerFeatures): AmazonQTokenServiceManager { - if (!AmazonQTokenServiceManager.instance) { - AmazonQTokenServiceManager.instance = new AmazonQTokenServiceManager(features) - AmazonQTokenServiceManager.instance.initialize() - - return AmazonQTokenServiceManager.instance - } - - throw new AmazonQServiceAlreadyInitializedError() - } - - public static getInstance(): AmazonQTokenServiceManager { - if (!AmazonQTokenServiceManager.instance) { - throw new AmazonQServiceInitializationError( - 'Amazon Q service has not been initialized yet. Make sure the Amazon Q server is present and properly initialized.' - ) - } - - return AmazonQTokenServiceManager.instance - } - - private initialize(): void { - if (!this.features.lsp.getClientInitializeParams()) { - this.log('AmazonQTokenServiceManager initialized before LSP connection was initialized.') - throw new AmazonQServiceInitializationError( - 'AmazonQTokenServiceManager initialized before LSP connection was initialized.' - ) - } - - // Bind methods that are passed by reference to some handlers to maintain proper scope. - this.serviceFactory = this.serviceFactory.bind(this) - - this.log('Reading enableDeveloperProfileSupport setting from AWSInitializationOptions') - if (this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws) { - const awsOptions = this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws || {} - this.enableDeveloperProfileSupport = signalsAWSQDeveloperProfilesEnabled(awsOptions) - - this.log(`Enabled Q Developer Profile support: ${this.enableDeveloperProfileSupport}`) - } - - this.connectionType = 'none' - this.state = 'PENDING_CONNECTION' - - this.log('Manager instance is initialize') - } - - public handleOnCredentialsDeleted(type: CredentialsType): void { - this.log(`Received credentials delete event for type: ${type}`) - if (type === 'iam') { - return - } - this.cancelActiveProfileChangeToken() - - this.resetCodewhispererService() - this.connectionType = 'none' - this.state = 'PENDING_CONNECTION' - } - - public async handleOnUpdateConfiguration(params: UpdateConfigurationParams, _token: CancellationToken) { - try { - if (params.section === Q_CONFIGURATION_SECTION && params.settings.profileArn !== undefined) { - const profileArn = params.settings.profileArn - const region = params.settings.region - - if (!isStringOrNull(profileArn)) { - throw new Error('Expected params.settings.profileArn to be of either type string or null') - } - - this.log(`Profile update is requested for profile ${profileArn}`) - this.cancelActiveProfileChangeToken() - this.profileChangeTokenSource = new CancellationTokenSource() - - await this.handleProfileChange(profileArn, this.profileChangeTokenSource.token) - } - } catch (error) { - this.log('Error updating profiles: ' + error) - if (error instanceof AmazonQServiceProfileUpdateCancelled) { - throw new ResponseError(LSPErrorCodes.ServerCancelled, error.message, { - awsErrorCode: error.code, - }) - } - if (error instanceof AmazonQError) { - throw new ResponseError(LSPErrorCodes.RequestFailed, error.message, { - awsErrorCode: error.code, - }) - } - - throw new ResponseError(LSPErrorCodes.RequestFailed, 'Failed to update configuration') - } finally { - if (this.profileChangeTokenSource) { - this.profileChangeTokenSource.dispose() - this.profileChangeTokenSource = undefined - } - } - } - - /** - * Validate if Bearer Token Connection type has changed mid-session. - * When connection type change is detected: reinitialize CodeWhispererService class with current connection type. - */ - private handleSsoConnectionChange() { - const newConnectionType = this.features.credentialsProvider.getConnectionType() - - this.logServiceState('Validate State of SSO Connection') - - const noCreds = !this.features.credentialsProvider.hasCredentials('bearer') - const noConnectionType = newConnectionType === 'none' - if (noCreds || noConnectionType) { - // Connection was reset, wait for SSO connection token from client - this.log( - `No active SSO connection is detected: no ${noCreds ? 'credentials' : 'connection type'} provided. Resetting the client` - ) - this.resetCodewhispererService() - this.connectionType = 'none' - this.state = 'PENDING_CONNECTION' - - return - } - - // Connection type hasn't change. - - if (newConnectionType === this.connectionType) { - this.logging.debug(`Connection type did not change: ${this.connectionType}`) - - return - } - - const endpointOverride = - this.features.lsp.getClientInitializeParams()?.initializationOptions?.aws?.awsClientCapabilities - ?.textDocument?.inlineCompletionWithReferences?.endpointOverride - - // Connection type changed to 'builderId' - - if (newConnectionType === 'builderId') { - this.log('Detected New connection type: builderId') - this.resetCodewhispererService() - - // For the builderId connection type regional endpoint discovery chain is: - // region set by client -> runtime region -> default region - const clientParams = this.features.lsp.getClientInitializeParams() - - this.createCodewhispererServiceInstances( - 'builderId', - clientParams?.initializationOptions?.aws?.region, - endpointOverride - ) - this.state = 'INITIALIZED' - this.log('Initialized Amazon Q service with builderId connection') - - return - } - - // Connection type changed to 'identityCenter' - - if (newConnectionType === 'identityCenter') { - this.log('Detected New connection type: identityCenter') - - this.resetCodewhispererService() - - if (this.enableDeveloperProfileSupport) { - this.connectionType = 'identityCenter' - this.state = 'PENDING_Q_PROFILE' - this.logServiceState('Pending profile selection for IDC connection') - - return - } - - this.createCodewhispererServiceInstances('identityCenter', undefined, endpointOverride) - this.state = 'INITIALIZED' - this.log('Initialized Amazon Q service with identityCenter connection') - - return - } - - this.logServiceState('Unknown Connection state') - } - - private cancelActiveProfileChangeToken() { - this.profileChangeTokenSource?.cancel() - this.profileChangeTokenSource?.dispose() - this.profileChangeTokenSource = undefined - } - - private handleTokenCancellationRequest(token: CancellationToken) { - if (token.isCancellationRequested) { - this.logServiceState('Handling CancellationToken cancellation request') - throw new AmazonQServiceProfileUpdateCancelled('Requested profile update got cancelled') - } - } - - private async handleProfileChange(newProfileArn: string | null, token: CancellationToken): Promise { - if (!this.enableDeveloperProfileSupport) { - this.log('Developer Profiles Support is not enabled') - return - } - - if (typeof newProfileArn === 'string' && newProfileArn.length === 0) { - throw new Error('Received invalid Profile ARN (empty string)') - } - - this.logServiceState('UpdateProfile is requested') - - // Test if connection type changed - this.handleSsoConnectionChange() - - if (this.connectionType === 'none') { - if (newProfileArn !== null) { - // During reauthentication, connection might be temporarily 'none' but user is providing a profile - // Set connection type to identityCenter to proceed with profile setting - this.connectionType = 'identityCenter' - this.state = 'PENDING_Q_PROFILE_UPDATE' - } else { - this.logServiceState('Received null profile while not connected, ignoring request') - return - } - } - - if (this.connectionType !== 'identityCenter') { - this.logServiceState('Q Profile can not be set') - throw new AmazonQServiceNoProfileSupportError( - `Connection type ${this.connectionType} does not support Developer Profiles feature.` - ) - } - - if ((this.state === 'INITIALIZED' && this.activeIdcProfile) || this.state === 'PENDING_Q_PROFILE') { - // Change status to pending to prevent API calls until profile is updated. - // Because `listAvailableProfiles` below can take few seconds to complete, - // there is possibility that client could send requests while profile is changing. - this.state = 'PENDING_Q_PROFILE_UPDATE' - } - - // Client sent an explicit null, indicating they want to reset the assigned profile (if any) - if (newProfileArn === null) { - this.logServiceState('Received null profile, resetting to PENDING_Q_PROFILE state') - this.resetCodewhispererService() - this.state = 'PENDING_Q_PROFILE' - - return - } - - const parsedArn = parse(newProfileArn) - const region = parsedArn.region - const endpoint = AWS_Q_ENDPOINTS.get(region) - if (!endpoint) { - throw new Error('Requested profileArn region is not supported') - } - - // Hack to inject a dummy profile name as it's not used by client IDE for now, if client IDE starts consuming name field then we should also pass both profile name and arn from the IDE - // When service is ready to take more tps, revert https://github.com/aws/language-servers/pull/1329 to add profile validation - const newProfile: AmazonQDeveloperProfile = { - arn: newProfileArn, - name: 'Client provided profile', - identityDetails: { - region: parsedArn.region, - }, - } - - if (!newProfile || !newProfile.identityDetails?.region) { - this.log(`Amazon Q Profile ${newProfileArn} is not valid`) - this.resetCodewhispererService() - this.state = 'PENDING_Q_PROFILE' - - throw new AmazonQServiceInvalidProfileError('Requested Amazon Q Profile does not exist') - } - - this.handleTokenCancellationRequest(token) - - if (!this.activeIdcProfile) { - this.activeIdcProfile = newProfile - this.createCodewhispererServiceInstances( - 'identityCenter', - newProfile.identityDetails.region, - this.endpointOverride() - ) - this.state = 'INITIALIZED' - this.log( - `Initialized identityCenter connection to region ${newProfile.identityDetails.region} for profile ${newProfile.arn}` - ) - - return - } - - // Profile didn't change - if (this.activeIdcProfile && this.activeIdcProfile.arn === newProfile.arn) { - // Update cached profile fields, keep existing client - this.log(`Profile selection did not change, active profile is ${this.activeIdcProfile.arn}`) - this.activeIdcProfile = newProfile - this.state = 'INITIALIZED' - - return - } - - this.handleTokenCancellationRequest(token) - - // At this point new valid profile is selected. - - const oldRegion = this.activeIdcProfile.identityDetails?.region - const newRegion = newProfile.identityDetails.region - if (oldRegion === newRegion) { - this.log(`New profile is in the same region as old one, keeping exising service.`) - this.log(`New active profile is ${this.activeIdcProfile.arn}, region ${oldRegion}`) - this.activeIdcProfile = newProfile - this.state = 'INITIALIZED' - - if (this.cachedCodewhispererService) { - this.cachedCodewhispererService.profileArn = newProfile.arn - } - - if (this.cachedStreamingClient) { - this.cachedStreamingClient.profileArn = newProfile.arn - } - - return - } - - this.log(`Switching service client region from ${oldRegion} to ${newRegion}`) - this.notifyRegionChangeListeners(newRegion) - - this.handleTokenCancellationRequest(token) - - // Selected new profile is in different region. Re-initialize service - this.resetCodewhispererService() - - this.activeIdcProfile = newProfile - - this.createCodewhispererServiceInstances( - 'identityCenter', - newProfile.identityDetails.region, - this.endpointOverride() - ) - this.state = 'INITIALIZED' - - return - } - - public getCodewhispererService(): CodeWhispererServiceToken { - // Prevent initiating requests while profile change is in progress. - if (this.state === 'PENDING_Q_PROFILE_UPDATE') { - throw new AmazonQServicePendingProfileUpdateError() - } - - this.handleSsoConnectionChange() - - if (this.state === 'INITIALIZED' && this.cachedCodewhispererService) { - return this.cachedCodewhispererService - } - - if (this.state === 'PENDING_CONNECTION') { - throw new AmazonQServicePendingSigninError() - } - - if (this.state === 'PENDING_Q_PROFILE') { - throw new AmazonQServicePendingProfileError() - } - - throw new AmazonQServiceNotInitializedError() - } - - public getStreamingClient() { - this.log('Getting instance of CodeWhispererStreaming client') - - // Trigger checks in token service - const tokenService = this.getCodewhispererService() - - if (!tokenService || !this.region || !this.endpoint) { - throw new AmazonQServiceNotInitializedError() - } - - if (!this.cachedStreamingClient) { - this.cachedStreamingClient = this.streamingClientFactory(this.region, this.endpoint) - } - - return this.cachedStreamingClient - } - - private resetCodewhispererService() { - this.cachedCodewhispererService?.abortInflightRequests() - this.cachedCodewhispererService = undefined - this.cachedStreamingClient?.abortInflightRequests() - this.cachedStreamingClient = undefined - this.activeIdcProfile = undefined - this.region = undefined - this.endpoint = undefined - } - - private createCodewhispererServiceInstances( - connectionType: 'builderId' | 'identityCenter', - clientOrProfileRegion: string | undefined, - endpointOverride: string | undefined - ) { - this.logServiceState('Initializing CodewhispererService') - - const { region, endpoint } = getAmazonQRegionAndEndpoint( - this.features.runtime, - this.features.logging, - clientOrProfileRegion - ) - - // Cache active region and endpoint selection - this.connectionType = connectionType - this.region = region - this.endpoint = endpoint - - if (endpointOverride) { - this.endpoint = endpointOverride - } - - this.cachedCodewhispererService = this.serviceFactory(region, this.endpoint) - this.log(`CodeWhispererToken service for connection type ${connectionType} was initialized, region=${region}`) - - this.cachedStreamingClient = this.streamingClientFactory(region, this.endpoint) - this.log(`StreamingClient service for connection type ${connectionType} was initialized, region=${region}`) - - this.logServiceState('CodewhispererService and StreamingClient Initialization finished') - } - - private getCustomUserAgent() { - const initializeParams = this.features.lsp.getClientInitializeParams() || {} - - return getUserAgent(initializeParams as InitializeParams, this.features.runtime.serverInfo) - } - - private serviceFactory(region: string, endpoint: string): CodeWhispererServiceToken { - const service = new CodeWhispererServiceToken( - this.features.credentialsProvider, - this.features.workspace, - this.features.logging, - region, - endpoint, - this.features.sdkInitializator - ) - - const customUserAgent = this.getCustomUserAgent() - service.updateClientConfig({ - customUserAgent: customUserAgent, - }) - service.customizationArn = this.configurationCache.getProperty('customizationArn') - service.profileArn = this.activeIdcProfile?.arn - service.shareCodeWhispererContentWithAWS = this.configurationCache.getProperty( - 'shareCodeWhispererContentWithAWS' - ) - - this.log('Configured CodeWhispererServiceToken instance settings:') - this.log( - `customUserAgent=${customUserAgent}, customizationArn=${service.customizationArn}, shareCodeWhispererContentWithAWS=${service.shareCodeWhispererContentWithAWS}` - ) - - return service - } - - private streamingClientFactory(region: string, endpoint: string): StreamingClientServiceToken { - const streamingClient = new StreamingClientServiceToken( - this.features.credentialsProvider, - this.features.sdkInitializator, - this.features.logging, - region, - endpoint, - this.getCustomUserAgent() - ) - streamingClient.profileArn = this.activeIdcProfile?.arn - - this.logging.debug(`Created streaming client instance region=${region}, endpoint=${endpoint}`) - return streamingClient - } - - private log(message: string): void { - const prefix = 'Amazon Q Token Service Manager' - this.logging?.log(`${prefix}: ${message}`) - } - - private logServiceState(context: string): void { - this.logging?.debug( - JSON.stringify({ - context, - state: { - serviceStatus: this.state, - connectionType: this.connectionType, - activeIdcProfile: this.activeIdcProfile, - }, - }) - ) - } - - // For Unit Tests - public static resetInstance(): void { - AmazonQTokenServiceManager.instance = null - } - - public getState() { - return this.state - } - - public getConnectionType() { - return this.connectionType - } - - public override getActiveProfileArn() { - return this.activeIdcProfile?.arn - } - - public setServiceFactory(factory: (region: string, endpoint: string) => CodeWhispererServiceToken) { - this.serviceFactory = factory.bind(this) - } - - public getServiceFactory() { - return this.serviceFactory - } - - public getEnableDeveloperProfileSupport(): boolean { - return this.enableDeveloperProfileSupport === undefined ? false : this.enableDeveloperProfileSupport - } - - /** - * Registers a listener that will be called when the region changes - * @param listener Function that will be called with the new region - * @returns Function to unregister the listener - */ - public override onRegionChange(listener: (region: string) => void): () => void { - this.regionChangeListeners.push(listener) - // If we already have a region, notify the listener immediately - if (this.region) { - try { - listener(this.region) - } catch (error) { - this.logging.error(`Error in region change listener: ${error}`) - } - } - return () => { - this.regionChangeListeners = this.regionChangeListeners.filter(l => l !== listener) - } - } - - private notifyRegionChangeListeners(region: string): void { - this.logging.debug( - `Notifying ${this.regionChangeListeners.length} region change listeners of region: ${region}` - ) - this.regionChangeListeners.forEach(listener => { - try { - listener(region) - } catch (error) { - this.logging.error(`Error in region change listener: ${error}`) - } - }) - } - - public getRegion(): string | undefined { - return this.region - } -} - -export const initBaseTokenServiceManager = (features: QServiceManagerFeatures) => - AmazonQTokenServiceManager.initInstance(features) - -export const getOrThrowBaseTokenServiceManager = (): AmazonQBaseServiceManager => - AmazonQTokenServiceManager.getInstance() diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.ts index e7048afe52..d34e71c4aa 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/qDeveloperProfiles.ts @@ -7,7 +7,7 @@ import { } from '@aws/language-server-runtimes/server-interface' import { isBool, isObject, SsoConnectionType } from '../utils' import { AWS_Q_ENDPOINTS } from '../../shared/constants' -import { CodeWhispererServiceToken } from '../codeWhispererService/codeWhispererServiceToken' +import { CodeWhispererServiceBase } from '../codeWhispererService/codeWhispererServiceBase' import { AmazonQServiceProfileThrottlingError } from './errors' export interface AmazonQDeveloperProfile { @@ -35,7 +35,7 @@ export const MAX_Q_DEVELOPER_PROFILE_PAGES = 10 const MAX_Q_DEVELOPER_PROFILES_PER_PAGE = 10 export const getListAllAvailableProfilesHandler = - (service: (region: string, endpoint: string) => CodeWhispererServiceToken): ListAllAvailableProfilesHandler => + (service: (region: string, endpoint: string) => CodeWhispererServiceBase): ListAllAvailableProfilesHandler => async ({ connectionType, logging, endpoints, token }): Promise => { if (!connectionType || connectionType !== 'identityCenter') { logging.debug('Connection type is not set or not identityCenter - returning empty response.') @@ -117,7 +117,7 @@ export const getListAllAvailableProfilesHandler = } async function fetchProfilesFromRegion( - service: CodeWhispererServiceToken, + service: CodeWhispererServiceBase, region: string, logging: Logging, token: CancellationToken diff --git a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/testUtils.ts b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/testUtils.ts index 622e325780..eb60657dc9 100644 --- a/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/testUtils.ts +++ b/server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/testUtils.ts @@ -125,7 +125,7 @@ export const initBaseTestServiceManager = { - * generateSingletonInitializationTests(AmazonQTokenServiceManager) + * generateSingletonInitializationTests(AmazonQServiceManager) * }) * ``` */ diff --git a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService/codeWhispererServiceToken.ts b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService/codeWhispererServiceToken.ts index 648ac26b58..2cd8db0982 100644 --- a/server/aws-lsp-codewhisperer/src/shared/codeWhispererService/codeWhispererServiceToken.ts +++ b/server/aws-lsp-codewhisperer/src/shared/codeWhispererService/codeWhispererServiceToken.ts @@ -53,7 +53,7 @@ import { } from './codeWhispererServiceBase' /** - * Hint: to get an instance of this: `AmazonQTokenServiceManager.getInstance().getCodewhispererService()` + * Hint: to get an instance of this: `AmazonQServiceManager.getInstance().getCodewhispererService()` */ export class CodeWhispererServiceToken extends CodeWhispererServiceBase { client: CodeWhispererTokenClient diff --git a/server/aws-lsp-codewhisperer/src/shared/proxy-server.ts b/server/aws-lsp-codewhisperer/src/shared/proxy-server.ts index 7313aacaba..9a07fa0271 100644 --- a/server/aws-lsp-codewhisperer/src/shared/proxy-server.ts +++ b/server/aws-lsp-codewhisperer/src/shared/proxy-server.ts @@ -1,24 +1,19 @@ import { QAgenticChatServer } from '../language-server/agenticChat/qAgenticChatServer' import { SecurityScanServerToken } from '../language-server/securityScan/codeWhispererSecurityScanServer' -import { CodewhispererServerFactory } from '../language-server/inline-completion/codeWhispererServer' +import { CodeWhispererServer } from '../language-server/inline-completion/codeWhispererServer' import { QNetTransformServerToken } from '../language-server/netTransform/netTransformServer' -import { QChatServerFactory } from '../language-server/chat/qChatServer' +import { QChatServer } from '../language-server/chat/qChatServer' import { QConfigurationServerToken } from '../language-server/configuration/qConfigurationServer' -import { getOrThrowBaseTokenServiceManager } from './amazonQServiceManager/AmazonQTokenServiceManager' -import { getOrThrowBaseIAMServiceManager } from './amazonQServiceManager/AmazonQIAMServiceManager' import { LocalProjectContextServer } from '../language-server/localProjectContext/localProjectContextServer' import { WorkspaceContextServer } from '../language-server/workspaceContext/workspaceContextServer' -export const CodeWhispererServerTokenProxy = CodewhispererServerFactory(getOrThrowBaseTokenServiceManager) - -export const CodeWhispererServerIAMProxy = CodewhispererServerFactory(getOrThrowBaseIAMServiceManager) +export const CodeWhispererServerProxy = CodeWhispererServer export const CodeWhispererSecurityScanServerTokenProxy = SecurityScanServerToken() export const QNetTransformServerTokenProxy = QNetTransformServerToken() -export const QChatServerTokenProxy = QChatServerFactory(getOrThrowBaseTokenServiceManager) -export const QChatServerIAMProxy = QChatServerFactory(getOrThrowBaseIAMServiceManager) +export const QChatServerProxy = QChatServer export const QAgenticChatServerProxy = QAgenticChatServer() diff --git a/server/aws-lsp-codewhisperer/src/shared/testUtils.ts b/server/aws-lsp-codewhisperer/src/shared/testUtils.ts index e444809cc5..71856eeb6e 100644 --- a/server/aws-lsp-codewhisperer/src/shared/testUtils.ts +++ b/server/aws-lsp-codewhisperer/src/shared/testUtils.ts @@ -330,9 +330,6 @@ export const setTokenCredentialsForAmazonQServiceManagerFactory = (getFeatures: } } -// TODO: remove this when changing references -export const setCredentialsForAmazonQTokenServiceManagerFactory = setTokenCredentialsForAmazonQServiceManagerFactory - export const setIamCredentialsForAmazonQServiceManagerFactory = (getFeatures: () => TestFeatures) => { return () => { const features = getFeatures() diff --git a/server/aws-lsp-codewhisperer/src/shared/utils.test.ts b/server/aws-lsp-codewhisperer/src/shared/utils.test.ts index 7136b04053..0f572b9d12 100644 --- a/server/aws-lsp-codewhisperer/src/shared/utils.test.ts +++ b/server/aws-lsp-codewhisperer/src/shared/utils.test.ts @@ -32,7 +32,7 @@ describe('getBearerTokenFromProvider', () => { const mockToken = 'mockToken' it('returns the bearer token from the provider', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: mockToken }), getConnectionMetadata: sinon.stub(), getConnectionType: sinon.stub(), @@ -43,7 +43,7 @@ describe('getBearerTokenFromProvider', () => { it('throws an error if the credentials does not contain bearer credentials', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(false), + hasCredentials: sinon.stub().withArgs('bearer').returns(false), getCredentials: sinon.stub().returns({ token: mockToken }), getConnectionMetadata: sinon.stub(), getConnectionType: sinon.stub(), @@ -58,7 +58,7 @@ describe('getBearerTokenFromProvider', () => { it('throws an error if token is empty in bearer token', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: '' }), getConnectionMetadata: sinon.stub(), getConnectionType: sinon.stub(), @@ -97,7 +97,7 @@ describe('getOriginFromClientInfo', () => { describe('getSsoConnectionType', () => { const mockToken = 'mockToken' const mockCredsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: mockToken }), getConnectionMetadata: sinon.stub(), getConnectionType: sinon.stub(), @@ -105,7 +105,7 @@ describe('getSsoConnectionType', () => { } it('should return ssoConnectionType as builderId', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: { @@ -121,7 +121,7 @@ describe('getSsoConnectionType', () => { it('should return ssoConnectionType as identityCenter', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: { @@ -142,7 +142,7 @@ describe('getSsoConnectionType', () => { it('should return ssoConnectionType as none when getConnectionMetadata.sso returns undefined', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: undefined, @@ -156,7 +156,7 @@ describe('getSsoConnectionType', () => { it('should return ssoConnectionType as none when getConnectionMetadata.sso.startUrl is empty string', () => { const mockCredentialsProvider: CredentialsProvider = { - hasCredentials: sinon.stub().returns(true), + hasCredentials: sinon.stub().withArgs('bearer').returns(true), getCredentials: sinon.stub().returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: { @@ -173,7 +173,7 @@ describe('getSsoConnectionType', () => { it('should return ssoConnectionType as none when getConnectionMetadata.sso.startUrl returns undefined', () => { const mockCredentialsProvider: CredentialsProvider = { hasCredentials: sinon.stub().returns(true), - getCredentials: sinon.stub().returns({ token: 'token' }), + getCredentials: sinon.stub().withArgs('bearer').returns({ token: 'token' }), getConnectionMetadata: sinon.stub().returns({ sso: { startUrl: undefined, diff --git a/tests/aws-lsp-codewhisperer.js b/tests/aws-lsp-codewhisperer.js index 668426a455..0e9370390d 100644 --- a/tests/aws-lsp-codewhisperer.js +++ b/tests/aws-lsp-codewhisperer.js @@ -1,15 +1,11 @@ const assert = require('assert') const codewhispererPackage = require('@aws/lsp-codewhisperer') -assert.ok(codewhispererPackage.CodeWhispererServerIAM) -assert.ok(codewhispererPackage.CodeWhispererServerIAMProxy) -assert.ok(codewhispererPackage.CodeWhispererServerToken) -assert.ok(codewhispererPackage.CodeWhispererServerTokenProxy) +assert.ok(codewhispererPackage.CodeWhispererServer) +assert.ok(codewhispererPackage.CodeWhispererServerProxy) assert.ok(codewhispererPackage.SecurityScanServerToken) assert.ok(codewhispererPackage.CodeWhispererSecurityScanServerTokenProxy) -assert.ok(codewhispererPackage.QChatServerIAM) -assert.ok(codewhispererPackage.QChatServerIAMProxy) -assert.ok(codewhispererPackage.QChatServerToken) -assert.ok(codewhispererPackage.QChatServerTokenProxy) +assert.ok(codewhispererPackage.QChatServer) +assert.ok(codewhispererPackage.QChatServerProxy) console.info('AWS Codewhisperer LSP: all tests passed')