feat: add iam support in q chat server #945
Conversation
…and QServiceManager
server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/BaseAmazonQServiceManager.ts
Outdated
Show resolved
Hide resolved
server/aws-lsp-codewhisperer/src/shared/amazonQServiceManager/AmazonQIAMServiceManager.ts
Outdated
Show resolved
Hide resolved
server/aws-lsp-codewhisperer/src/language-server/agenticChat/agenticChatController.ts
Outdated
Show resolved
Hide resolved
server/aws-lsp-codewhisperer/src/language-server/chat/chatController.ts
Outdated
Show resolved
Hide resolved
server/aws-lsp-codewhisperer/src/language-server/chat/chatSessionService.test.ts
Show resolved
Hide resolved
…ginal QServiceManagerFeatures
…whisperer-runtimes
| @@ -26,8 +26,10 @@ export const QNetTransformServerTokenProxy = QNetTransformServerToken( | |||
| } | |||
| ) | |||
|
|
|||
| export const QChatServerProxy = QChatServer() | |||
| export const QChatServerTokenProxy = QChatServerFactory(initBaseTokenServiceManager) | |||
| export const QChatServerIAMProxy = QChatServerFactory(initBaseIAMServiceManager) | |||
There was a problem hiding this comment.
We should stop using Proxy suffix for servers, it is meaningless and will be deprected
There was a problem hiding this comment.
would you suggest removing it all for all servers and eventually remove or rename the proxy-server.ts file itself? I understand your point, but for consistency either I would keep the same proxy naming pattern for all or for none, if we decide to remove the proxy suffix for all i we should do it in a followup PR imo.
There was a problem hiding this comment.
We will need to do both: remove Proxy suffix from all these functions, and remove proxy-server file altogether. I'm ok with keeping current change for consistency.
| @@ -551,6 +557,6 @@ export class AmazonQTokenServiceManager extends BaseAmazonQServiceManager<CodeWh | |||
| } | |||
| } | |||
|
|
|||
| export const initBaseTokenServiceManager = (features: Features): AmazonQBaseServiceManager => { | |||
| export const initBaseTokenServiceManager = (features: QServiceManagerFeatures): AmazonQBaseServiceManager => { | |||
There was a problem hiding this comment.
nit: why do we need to change the type from Features? I think it's cleaner to keep using previous one
There was a problem hiding this comment.
With @patrickvbeurden we agreed (see comment) that it was confusing having two types having the same name but different properties (i.e. the Features defined in runtime and Features type defined in BaseAmazonQServiceManager), one being a subset of the other. Imo QServiceManagerFeatures gives a better idea of the fact that it is a subset of the general Features type. Wdyt?
| this.inflightRequests.clear() | ||
| } | ||
|
|
||
| public async sendMessage( |
There was a problem hiding this comment.
Looks like common sendMessage logic could also be consolidated in base class, it's only 1 parameter that's different, which could be passed to base method.
There was a problem hiding this comment.
Initially I agreed, but i tried modifying but it ends up making the code harder to read, we need to pass a child class specific param (profileArn) as optional to base class, and the call this.client.sendMessage has to be done via if / else check on the streaming client to make sure params of the call are compatible. I will keep current implementation
| const response = await client.generateAssistantResponse(request, this.#abortController) | ||
|
|
||
| return response | ||
| if (client instanceof StreamingClientServiceToken) { |
There was a problem hiding this comment.
Can we avoid relying of Class instance here? You can check availability of generateAssistantResponse method on the StreamingClient class instead, and keep ChatSessionService independent from type of streaming service used.
There was a problem hiding this comment.
checking availability of generateAssistantResponse method still requires casting with specific streaming service (StreamingClientServiceToken), as client is an instance of the base class which does not have generateAssistantResponse method
if ('generateAssistantResponse' in client) {
const tokenClient = client as StreamingClientServiceToken
const response = await tokenClient.generateAssistantResponse(request, this.#abortController)
without casting it throws a typescript compile error: 'client.generateAssistantResponse' is of type 'unknown'.
| import { QChatServerFactory } from './qChatServer' | ||
| import { AmazonQIAMServiceManager } from '../../shared/amazonQServiceManager/AmazonQIAMServiceManager' | ||
|
|
||
| describe('QChatServerIAM', () => { |
There was a problem hiding this comment.
Cna we optimize tests setup and avoid having 2 identical test suite files? Given QChatServer does not have any specific implementation based on type of service used, we should be able to run same test suite for both service manager variants.
francescoopiccoli
changed the title
* 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]>
Problem
In this PR we add IAM auth support to q chat server
and q agentic server. It required some refactoring of the chat servers and the StreamingClientService, as previously chat servers were coupled with token auth only. To do this, for consistency and clarity I followed the same patterns as we do for inline-completion servers (CodewhispererServerFactory) and for inline-completion services (CodeWhispererServiceBase).To manually test follow the updated instructions in readme and contributing.md, this pr does not include any new tests, some basic tests have been added in #953.
The PR was initially meant to be reviewed commit by commit, after requested changes this no longer holds. I would still suggest to review based on which files were modified in each initial commit.
License
By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.