refactor(amazonq): change AmazonQServiceManager references

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

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)

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',
 }
 

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)

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,

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'

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,
+    setTokenCredentialsForAmazonQServiceManagerFactory,
+    setIamCredentialsForAmazonQServiceManagerFactory,
+} from '../../shared/testUtils'
 import sinon from 'ts-sinon'
 import { AgenticChatController } from './agenticChatController'
 import { ChatSessionManagementService } from '../chat/chatSessionManagementService'
@@ -45,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'
@@ -170,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
@@ -179,7 +182,8 @@ describe('AgenticChatController', () => {
     let getMessagesStub: sinon.SinonStub
     let addMessageStub: sinon.SinonStub
 
-    const setCredentials = setCredentialsForAmazonQTokenServiceManagerFactory(() => testFeatures)
+    const setSsoCredentials = setTokenCredentialsForAmazonQServiceManagerFactory(() => testFeatures)
+    const setIamCredentials = setIamCredentialsForAmazonQServiceManagerFactory(() => testFeatures)
 
     beforeEach(() => {
         // Override the response timeout for tests to avoid long waits
@@ -272,7 +276,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')
@@ -281,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: {
@@ -2993,7 +2997,7 @@ ${' '.repeat(8)}}
     })
 
     describe('onListAvailableModels', () => {
-        let tokenServiceManagerStub: sinon.SinonStub
+        let serviceManagerStub: sinon.SinonStub
 
         beforeEach(() => {
             // Create a session with a model ID
@@ -3002,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 }
@@ -3030,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 }
@@ -3049,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 }
@@ -3084,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')
@@ -3104,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')
@@ -3124,7 +3128,7 @@ ${' '.repeat(8)}}
     })
 
     describe('IAM Authentication', () => {
-        let iamServiceManager: AmazonQIAMServiceManager
+        let serviceManager: AmazonQServiceManager
         let iamChatController: AgenticChatController
         let iamChatSessionManagementService: ChatSessionManagementService
 
@@ -3144,26 +3148,29 @@ ${' '.repeat(8)}}
             // Reset the singleton instance
             ChatSessionManagementService.reset()
 
+            // Store IAM credentials
+            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', () => {

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<Chat>,
@@ -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<ListAvailableModelsResult> {
-        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(