feat: add iam support in q chat server and q agentic server (#945)

* feat: add StreamingClientServiceIAM, refactor StreamingClientService and QServiceManager

* refactor: qChat and qAgenticChat to support both iam and token auth

* fix: features object parameter in other servers after changes to qServiceManager

* fix: tests after refactor and renaming

* feat: export QChatServerIAMProxy and fix naming of other servers after refactor

* docs: testing and debugging with IAM auth

* feat: enable chat for CodeWhisperer Server IAM launch config

* fix: test/aws-lsp-codewhisperer.js assertions names

* refactor: restore runtimes-defined features type in controllers & original QServiceManagerFeatures

* refactor: remove unneeded functions and check in iam service manager

* chore: merge branch 'main' into frapicc/iam-support-q-chat-server

* chore: revert merge from main

* chore: remove profileArn from agenticChatController

* chore: remove unused parameters for qConfigurationServer

* feat: add QChatServerIAM to iam-webworker servers in app/aws-lsp-codewhisperer-runtimes

* test: add some tests for IAM q chat server

* fix: revert changes to agentic chat (remove iam support)

* fix: test

* test: avoid running seperate tests for q chat server iam / token

* chore: renam q chat server test class

---------

Co-authored-by: Francesco Piccoli <[email protected]>

Author: francescoopiccoli

.vscode/launch.json

@@ -128,7 +128,8 @@
             "env": {
                 "LSP_SERVER": "${workspaceFolder}/app/aws-lsp-codewhisperer-runtimes/out/iam-standalone.js",
                 "ENABLE_INLINE_COMPLETION": "true",
-                "ENABLE_IAM_PROVIDER": "true"
+                "ENABLE_IAM_PROVIDER": "true",
+                "ENABLE_CHAT": "true"
                 // "HTTPS_PROXY": "http://127.0.0.1:8888",
             }
             // "preLaunchTask": "compile"

CONTRIBUTING.md

@@ -211,15 +211,22 @@ and be able to debug it all.
    A new window will open.
 
 ### With CodeWhisperer Server in VSCode
-
-1. In the `Run & Debug` menu, run `"CodeWhisperer Server"`
+#### With token credentials
+1. In the `Run & Debug` menu, run `"CodeWhisperer Server Token"`
 2. Set breakpoints in `src` where needed
 3. Check the logs in `"AWS Documents Language Server"` output window.
 
 > **NOTE**: If you see "Recommendation failure: Error: Authorization failed, bearer token is not set" errors, make sure to authenticate using `"AWS LSP - Obtain bearer token and send to LSP server"` command.
 
 > **NOTE**: The lsp client is activated by one of the `activationEvents` defined [here](https://github.com/aws/language-servers/blob/06fd81d1e936648ef43243865039f89c7ac142a7/client/vscode/package.json#L18-L22), the lsp client then starts the LSP server.
 
+#### With IAM credentials
+1. In the `Run & Debug` menu, run `"CodeWhisperer Server IAM"`
+2. Set breakpoints in `src` where needed
+3. Check the logs in `"AWS Documents Language Server"` output window.
+
+> **NOTE**: To authenticate use ['AWS LSP - Choose credentials profile, resolve, and send to LSP Server'](https://github.com/aws/language-servers/blob/694bbb85580cc79313d65ad77b224875f74280c2/client/vscode/package.json#L32-L33) command, giving as input your iam credential profile name, for more info see [here](https://github.com/aws/language-servers/blob/694bbb85580cc79313d65ad77b224875f74280c2/client/vscode/README.md?plain=1#L11).
+
 ### With Other Clients
 Using other clients can also be done with the bundle created from this package.
 

app/aws-lsp-codewhisperer-runtimes/src/agent-standalone.ts

@@ -3,7 +3,7 @@ import { RuntimeProps } from '@aws/language-server-runtimes/runtimes/runtime'
 import {
     CodeWhispererSecurityScanServerTokenProxy,
     CodeWhispererServerTokenProxy,
-    QAgenticChatServerProxy,
+    QAgenticChatServerTokenProxy,
     QConfigurationServerTokenProxy,
     QNetTransformServerTokenProxy,
 } from '@aws/lsp-codewhisperer'
@@ -22,7 +22,7 @@ const props: RuntimeProps = {
         CodeWhispererSecurityScanServerTokenProxy,
         QConfigurationServerTokenProxy,
         QNetTransformServerTokenProxy,
-        QAgenticChatServerProxy,
+        QAgenticChatServerTokenProxy,
         IdentityServer.create,
         FsToolsServer,
     ],

app/aws-lsp-codewhisperer-runtimes/src/iam-standalone.ts

@@ -1,10 +1,10 @@
 import { standalone } from '@aws/language-server-runtimes/runtimes'
 import { RuntimeProps } from '@aws/language-server-runtimes/runtimes/runtime'
-import { CodeWhispererServerIAM } from '@aws/lsp-codewhisperer'
+import { CodeWhispererServerIAM, QChatServerIAMProxy } from '@aws/lsp-codewhisperer'
 
 const props: RuntimeProps = {
     version: '0.1.0',
-    servers: [CodeWhispererServerIAM],
+    servers: [CodeWhispererServerIAM, QChatServerIAMProxy],
     name: 'AWS CodeWhisperer',
 }
 standalone(props)

app/aws-lsp-codewhisperer-runtimes/src/iam-webworker.ts

@@ -1,10 +1,11 @@
 import { webworker } from '@aws/language-server-runtimes/runtimes/webworker'
 import { RuntimeProps } from '@aws/language-server-runtimes/runtimes/runtime'
 import { CodeWhispererServerIAM } from '@aws/lsp-codewhisperer/out/language-server/inline-completion/codeWhispererServer'
+import { QChatServerIAM } from '@aws/lsp-codewhisperer/out/language-server/chat/qChatServer'
 
 const props: RuntimeProps = {
     version: '1.0.0',
-    servers: [CodeWhispererServerIAM],
+    servers: [CodeWhispererServerIAM, QChatServerIAM],
     name: 'AWS CodeWhisperer',
 }
 

app/aws-lsp-codewhisperer-runtimes/src/token-standalone.ts

@@ -3,7 +3,7 @@ import { RuntimeProps } from '@aws/language-server-runtimes/runtimes/runtime'
 import {
     CodeWhispererSecurityScanServerTokenProxy,
     CodeWhispererServerTokenProxy,
-    QChatServerProxy,
+    QChatServerTokenProxy,
     QConfigurationServerTokenProxy,
     QNetTransformServerTokenProxy,
     QLocalProjectContextServerTokenProxy,
@@ -22,7 +22,7 @@ const props: RuntimeProps = {
         CodeWhispererSecurityScanServerTokenProxy,
         QConfigurationServerTokenProxy,
         QNetTransformServerTokenProxy,
-        QChatServerProxy,
+        QChatServerTokenProxy,
         IdentityServer.create,
         QLocalProjectContextServerTokenProxy,
     ],

client/vscode/README.md

@@ -8,7 +8,7 @@ This extension contains functioning code that shows an example of how extensions
 
 The intent behind this concept code is to show how an extension could "push" credentials to the server whenever they are resolved (eg: when a user configures credentials for use with the server, or when the credentials have been refreshed). To simulate these events, the following commands can be run from the extension:
 
--   awslsp.selectProfile (sigv4) - used to simulate when credentials are pushed to the server. When you run this command, you will be prompted to enter a profile name. This profile's credentials will be resolved from your shared credentials file, then pushed to the server. (Only basic access key - secret key credentials type is supported in this example).
+-   awslsp.selectProfile (sigv4) - used to simulate when credentials are pushed to the server. When you run this command, you will be prompted to enter a profile name. This profile's credentials will be resolved from your shared credentials file, then pushed to the server. (Only basic access key - secret key credentials type is supported in this example). From the minimal vscode client launched from "Run and Debug" with [CodeWhisperer Server IAM](https://github.com/aws/language-servers/blob/694bbb85580cc79313d65ad77b224875f74280c2/.vscode/launch.json#L100) configuration, the command is accessible via the command palette as ['AWS LSP - Choose credentials profile, resolve, and send to LSP Server'](https://github.com/aws/language-servers/blob/694bbb85580cc79313d65ad77b224875f74280c2/client/vscode/package.json#L32-L33).
 -   awslsp.clearProfile (sigv4), awslsp.clearBearerToken (token) - used to simulate when credentials are removed from the server
 -   awslsp.resolveBearerToken (token) - used to simulate when bearer tokens are pushed to the server. When you run this command, you will be taken through an SSO login flow, similar to what the AWS Toolkit for VS Code does. Once logged in, the bearer token will be pushed to the server.
 

server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts

@@ -7,8 +7,7 @@ import {
     ChatTriggerType,
     GenerateAssistantResponseCommandInput,
     GenerateAssistantResponseCommandOutput,
-    SendMessageCommandInput,
-    SendMessageCommandOutput,
+    SendMessageCommandInput as SendMessageCommandInputCodeWhispererStreaming,
     ToolResult,
     ToolResultContentBlock,
     ToolUse,
@@ -64,6 +63,7 @@ import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/A
 import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils'
 import { TabBarController } from './tabBarController'
 import { ChatDatabase } from './tools/chatDb/chatDb'
+import { SendMessageCommandInput, SendMessageCommandOutput } from '../../shared/streamingClientService'
 import {
     AgenticChatEventParser,
     ChatResultWithMetadata as AgenticChatResultWithMetadata,
@@ -512,21 +512,19 @@ export class AgenticChatController implements ChatHandlers {
         let requestInput: SendMessageCommandInput
 
         try {
-            const profileArn = AmazonQTokenServiceManager.getInstance(this.#features).getActiveProfileArn()
             requestInput = this.#triggerContext.getChatParamsFromTrigger(
                 params,
                 triggerContext,
                 ChatTriggerType.INLINE_CHAT,
-                this.#customizationArn,
-                profileArn
+                this.#customizationArn
             )
 
             if (!this.#amazonQServiceManager) {
                 throw new Error('amazonQServiceManager is not initialized')
             }
 
             const client = this.#amazonQServiceManager.getStreamingClient()
-            response = await client.sendMessage(requestInput)
+            response = await client.sendMessage(requestInput as SendMessageCommandInputCodeWhispererStreaming)
             this.#log('Response for inline chat', JSON.stringify(response.$metadata), JSON.stringify(response))
         } catch (err) {
             if (err instanceof AmazonQServicePendingSigninError || err instanceof AmazonQServicePendingProfileError) {

server/aws-lsp-codewhisperer/src/language-server/agenticChat/qAgenticChatServer.ts

@@ -20,10 +20,12 @@ 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: AmazonQTokenServiceManager
+        let telemetryService: TelemetryService
+
         let chatController: AgenticChatController
         let chatSessionManagementService: ChatSessionManagementService
-        let telemetryService: TelemetryService;
 
         lsp.addInitializer((params: InitializeParams) => {
             return {
@@ -52,29 +54,32 @@ export const QAgenticChatServer =
         lsp.onInitialized(async () => {
             // Initialize service manager and inject it to chatSessionManagementService to pass it down
             amazonQServiceManager = AmazonQTokenServiceManager.getInstance(features)
-            chatSessionManagementService = ChatSessionManagementService
-                .getInstance()
-                .withAmazonQServiceManager(amazonQServiceManager)
-
-            telemetryService = new TelemetryService(
-                amazonQServiceManager,
-                credentialsProvider,
-                telemetry,
-                logging,
-            )
+            chatSessionManagementService =
+                ChatSessionManagementService.getInstance().withAmazonQServiceManager(amazonQServiceManager)
 
-            const clientParams = safeGet(lsp.getClientInitializeParams(), new AmazonQServiceInitializationError(
-                'TelemetryService initialized before LSP connection was initialized.'))
+            telemetryService = new TelemetryService(amazonQServiceManager, credentialsProvider, telemetry, logging)
+
+            const clientParams = safeGet(
+                lsp.getClientInitializeParams(),
+                new AmazonQServiceInitializationError(
+                    'TelemetryService initialized before LSP connection was initialized.'
+                )
+            )
 
             telemetryService.updateUserContext(makeUserContextObject(clientParams, runtime.platform, 'CHAT'))
 
-            chatController = new AgenticChatController(chatSessionManagementService, features, telemetryService, amazonQServiceManager)
+            chatController = new AgenticChatController(
+                chatSessionManagementService,
+                features,
+                telemetryService,
+                amazonQServiceManager
+            )
 
             /* 
-                Calling handleDidChangeConfiguration once to ensure we get configuration atleast once at start up
-                
-                TODO: TODO: consider refactoring such responsibilities to common service manager config/initialisation server
-        */
+                            Calling handleDidChangeConfiguration once to ensure we get configuration atleast once at start up
+                            
+                            TODO: TODO: consider refactoring such responsibilities to common service manager config/initialisation server
+                    */
             await amazonQServiceManager.handleDidChangeConfiguration()
             await amazonQServiceManager.addDidChangeConfigurationListener(updateConfigurationHandler)
         })
@@ -134,7 +139,7 @@ export const QAgenticChatServer =
                     }
                 }
             }
-        } as const, async (input) => 
+        } as const, async (input) =>
             (await features.workspace.getAllTextDocuments()).map(td => td.uri).filter(uri => input.filter === undefined || uri.includes(input.filter))
         )
 
@@ -156,7 +161,7 @@ If a file is not open, use the \`fsRead\` tool to read from disk. Use this tool
                 },
                 required: ['paths']
             }
-        } as const, async (input) => 
+        } as const, async (input) =>
             input.paths.reduce(async (acc, path) => {
                 const doc = await features.workspace.getTextDocument(path)
                 if (doc) {

server/aws-lsp-codewhisperer/src/language-server/chat/chatController.ts

@@ -1,4 +1,4 @@
-import { ChatTriggerType, SendMessageCommandInput, SendMessageCommandOutput } from '@amzn/codewhisperer-streaming'
+import { ChatTriggerType } from '@amzn/codewhisperer-streaming'
 import {
     ApplyWorkspaceEditParams,
     ErrorCodes,
@@ -47,7 +47,8 @@ import {
 } from '../../shared/amazonQServiceManager/errors'
 import { TelemetryService } from '../../shared/telemetry/telemetryService'
 import { AmazonQWorkspaceConfig } from '../../shared/amazonQServiceManager/configurationUtils'
-import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
+import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
+import { SendMessageCommandInput, SendMessageCommandOutput } from '../../shared/streamingClientService'
 
 type ChatHandlers = Omit<
     LspHandlers<Chat>,
@@ -71,13 +72,13 @@ export class ChatController implements ChatHandlers {
     #triggerContext: QChatTriggerContext
     #customizationArn?: string
     #telemetryService: TelemetryService
-    #amazonQServiceManager: AmazonQTokenServiceManager
+    #amazonQServiceManager: AmazonQBaseServiceManager
 
     constructor(
         chatSessionManagementService: ChatSessionManagementService,
         features: Features,
         telemetryService: TelemetryService,
-        amazonQServiceManager: AmazonQTokenServiceManager
+        amazonQServiceManager: AmazonQBaseServiceManager
     ) {
         this.#features = features
         this.#chatSessionManagementService = chatSessionManagementService

server/aws-lsp-codewhisperer/src/language-server/chat/chatEventParser.ts

@@ -1,4 +1,9 @@
-import { ChatResponseStream, Reference, SupplementaryWebLink } from '@amzn/codewhisperer-streaming'
+import {
+    ChatResponseStream as ChatResponseStreamCodeWhispererStreaming,
+    Reference,
+    SupplementaryWebLink,
+} from '@amzn/codewhisperer-streaming'
+import { ChatResponseStream as ChatResponseStreamQDeveloperStreaming } from '@amzn/amazon-q-developer-streaming-client'
 import {
     ChatItemAction,
     ChatResult,
@@ -8,6 +13,7 @@ import {
 import { Result } from '../types'
 import { AddMessageEvent } from '../../shared/telemetry/types'
 import { Metric } from '../../shared/telemetry/metric'
+export type ChatResponseStream = ChatResponseStreamCodeWhispererStreaming | ChatResponseStreamQDeveloperStreaming
 
 export type ChatResultWithMetadata = {
     chatResult: ChatResult

server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionManagementService.ts

@@ -1,11 +1,10 @@
 import { Result } from '../types'
 import { ChatSessionService } from './chatSessionService'
-import { AmazonQTokenServiceManager } from '../../shared/amazonQServiceManager/AmazonQTokenServiceManager'
-
+import { AmazonQBaseServiceManager } from '../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
 export class ChatSessionManagementService {
     static #instance?: ChatSessionManagementService
     #sessionByTab: Map<string, ChatSessionService> = new Map<string, any>()
-    #amazonQServiceManager?: AmazonQTokenServiceManager
+    #amazonQServiceManager?: AmazonQBaseServiceManager
 
     public static getInstance() {
         if (!ChatSessionManagementService.#instance) {
@@ -21,7 +20,7 @@ export class ChatSessionManagementService {
 
     private constructor() {}
 
-    public withAmazonQServiceManager(amazonQServiceManager: AmazonQTokenServiceManager) {
+    public withAmazonQServiceManager(amazonQServiceManager: AmazonQBaseServiceManager) {
         this.#amazonQServiceManager = amazonQServiceManager
 
         return this