feat: integrate server side project context into agentic chat (#1405)

guvishl · Ege Ozcan · web-flow · commit e4d8f6144aef · 2025-05-21T18:06:46.000+02:00
* feat: pass remote workspaceId in agentic chat API calls

* feat: removed dev logs

* feat: add unit test

* feat: bumping up dependency version

* chore: updating package-lock.json

* fix: install new streaming client correctly

---------

Co-authored-by: Ege Ozcan &lt;eoozca@amazon.com&gt;
diff --git a/core/codewhisperer-streaming/amzn-codewhisperer-streaming-1.0.7.tgz b/core/codewhisperer-streaming/amzn-codewhisperer-streaming-1.0.7.tgz
diff --git a/package-lock.json b/package-lock.json
diff --git a/server/aws-lsp-codewhisperer/package.json b/server/aws-lsp-codewhisperer/package.json
@@ -28,7 +28,7 @@
     },
     "dependencies": {
         "@amzn/amazon-q-developer-streaming-client": "file:../../core/q-developer-streaming-client/amzn-amazon-q-developer-streaming-client-1.0.0.tgz",
-        "@amzn/codewhisperer-streaming": "file:../../core/codewhisperer-streaming/amzn-codewhisperer-streaming-1.0.6.tgz",
+        "@amzn/codewhisperer-streaming": "file:../../core/codewhisperer-streaming/amzn-codewhisperer-streaming-1.0.7.tgz",
         "@aws-sdk/util-arn-parser": "^3.723.0",
         "@aws-sdk/util-retry": "^3.374.0",
         "@aws/chat-client-ui-types": "^0.1.33",
diff --git a/server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json b/server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json
@@ -1506,6 +1506,10 @@
                     "shape": "ConversationId",
                     "documentation": "<p>Unique identifier for the chat conversation stream</p>"
                 },
+                "workspaceId": {
+                    "shape": "UUID",
+                    "documentation": "<p>Unique identifier for remote workspace</p>"
+                },
                 "history": {
                     "shape": "ChatHistory",
                     "documentation": "<p>Holds the history of chat messages.</p>"
diff --git a/server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts b/server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts
@@ -82,11 +82,11 @@ declare class CodeWhispererBearerTokenClient extends Service {
    */
   getCodeAnalysis(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetCodeAnalysisResponse) => void): Request<CodeWhispererBearerTokenClient.Types.GetCodeAnalysisResponse, AWSError>;
   /**
-   * 
+   *
    */
   getCodeFixJob(params: CodeWhispererBearerTokenClient.Types.GetCodeFixJobRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse, AWSError>;
   /**
-   * 
+   *
    */
   getCodeFixJob(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse, AWSError>;
   /**
@@ -130,11 +130,11 @@ declare class CodeWhispererBearerTokenClient extends Service {
    */
   listAvailableCustomizations(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.ListAvailableCustomizationsResponse) => void): Request<CodeWhispererBearerTokenClient.Types.ListAvailableCustomizationsResponse, AWSError>;
   /**
-   * 
+   *
    */
   listAvailableProfiles(params: CodeWhispererBearerTokenClient.Types.ListAvailableProfilesRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse) => void): Request<CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse, AWSError>;
   /**
-   * 
+   *
    */
   listAvailableProfiles(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse) => void): Request<CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse, AWSError>;
   /**
@@ -186,11 +186,11 @@ declare class CodeWhispererBearerTokenClient extends Service {
    */
   startCodeAnalysis(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.StartCodeAnalysisResponse) => void): Request<CodeWhispererBearerTokenClient.Types.StartCodeAnalysisResponse, AWSError>;
   /**
-   * 
+   *
    */
   startCodeFixJob(params: CodeWhispererBearerTokenClient.Types.StartCodeFixJobRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse, AWSError>;
   /**
-   * 
+   *
    */
   startCodeFixJob(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse, AWSError>;
   /**
@@ -463,6 +463,10 @@ declare namespace CodeWhispererBearerTokenClient {
      * Unique identifier for the chat conversation stream
      */
     conversationId?: ConversationId;
+    /**
+     * Unique identifier for remote workspace
+     */
+    workspaceId?: UUID;
     /**
      * Holds the history of chat messages.
      */
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContext.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContext.ts
@@ -33,6 +33,7 @@ import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
 import { languageByExtension } from '../../../shared/languageDetection'
 import { AgenticChatResultStream } from '../agenticChatResultStream'
 import { ContextInfo, mergeFileLists, mergeRelevantTextDocuments } from './contextUtils'
+import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
 
 export interface TriggerContext extends Partial<DocumentContext> {
     userIntent?: UserIntent
@@ -127,6 +128,16 @@ export class AgenticChatTriggerContext {
             promptContent = promptContent?.replace(/\*\*@workspace\*\*/, '')
         }
 
+        // Append remote workspaceId if it exists
+        // Only append workspaceId to GenerateCompletions when WebSocket client is connected
+        const remoteWsFolderManager = WorkspaceFolderManager.getInstance()
+        const workspaceId =
+            (remoteWsFolderManager &&
+                remoteWsFolderManager.getWorkspaceState().webSocketClient?.isConnected &&
+                remoteWsFolderManager.getWorkspaceState().workspaceId) ||
+            undefined
+        this.#logging.info(`remote workspaceId: ${workspaceId}`)
+
         // Get workspace documents if @workspace is used
         let relevantDocuments = hasWorkspace
             ? await this.#getRelevantDocuments(promptContent ?? '', chatResultStream)
@@ -176,6 +187,7 @@ export class AgenticChatTriggerContext {
 
         const data: GenerateAssistantResponseCommandInput = {
             conversationState: {
+                workspaceId: workspaceId,
                 chatTriggerType: chatTriggerType,
                 currentMessage: {
                     userInputMessage: {
diff --git a/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContexts.test.ts b/server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContexts.test.ts
@@ -14,6 +14,7 @@ import { ChatTriggerType, CursorState } from '@amzn/codewhisperer-streaming'
 import { URI } from 'vscode-uri'
 import { InitializeParams } from '@aws/language-server-runtimes/protocol'
 import { TestFolder } from '@aws/lsp-core/out/test/testFolder'
+import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
 
 describe('AgenticChatTriggerContext', () => {
     let testFeatures: TestFeatures
@@ -28,6 +29,7 @@ describe('AgenticChatTriggerContext', () => {
         hasCodeSnippet: false,
         totalEditorCharacters: 0,
     }
+    let mockWorkspaceFolderManager: any
 
     beforeEach(() => {
         testFeatures = new TestFeatures()
@@ -127,6 +129,39 @@ describe('AgenticChatTriggerContext', () => {
             mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
         )
     })
+    it('includes remote workspaceId if it exists and is connected', async () => {
+        mockWorkspaceFolderManager = {
+            getWorkspaceState: sinon.stub(),
+        }
+        sinon.stub(WorkspaceFolderManager, 'getInstance').returns(mockWorkspaceFolderManager)
+        mockWorkspaceFolderManager.getWorkspaceState.returns({
+            webSocketClient: { isConnected: true },
+            workspaceId: 'test-workspace-123',
+        })
+        const triggerContext = new AgenticChatTriggerContext(testFeatures)
+        const chatParams = await triggerContext.getChatParamsFromTrigger(
+            { tabId: 'tab', prompt: {} },
+            {},
+            ChatTriggerType.MANUAL
+        )
+        const chatParamsWithMore = await triggerContext.getChatParamsFromTrigger(
+            { tabId: 'tab', prompt: {} },
+            { cursorState: {} as CursorState, relativeFilePath: '' },
+            ChatTriggerType.MANUAL
+        )
+
+        assert.deepStrictEqual(
+            chatParams.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
+                ?.workspaceFolders,
+            mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
+        )
+        assert.deepStrictEqual(
+            chatParamsWithMore.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
+                ?.workspaceFolders,
+            mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
+        )
+        assert.deepStrictEqual(chatParamsWithMore.conversationState?.workspaceId, 'test-workspace-123')
+    })
     describe('getTextDocument', function () {
         let tempFolder: TestFolder
 
