fix: add workspace folders as context for agentic-chat

core/aws-lsp-core/src/util/workspaceUtils.ts

@@ -85,6 +85,10 @@ export function getEntryPath(entry: Dirent) {
 }
 
 // TODO: port this to runtimes?
+export function getWorkspaceFolders(lsp: Features['lsp']): string[] {
+    return lsp.getClientInitializeParams()?.workspaceFolders?.map(({ uri }) => URI.parse(uri).fsPath) ?? []
+}
+
 export async function inWorkspace(workspace: Features['workspace'], filepath: string) {
     return (await workspace.getTextDocument(URI.file(filepath).toString())) !== undefined
 }

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

@@ -926,7 +926,7 @@ describe('AgenticChatController', () => {
                 extractDocumentContextStub.restore()
             })
 
-            it('leaves editor state as undefined if cursorState is not passed', async () => {
+            it('leaves cursorState as undefined if cursorState is not passed', async () => {
                 const documentContextObject = {
                     programmingLanguage: 'typescript',
                     cursorState: undefined,
@@ -949,12 +949,12 @@ describe('AgenticChatController', () => {
 
                 assert.strictEqual(
                     calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
-                        ?.editorState,
+                        ?.editorState?.cursorState,
                     undefined
                 )
             })
 
-            it('leaves editor state as undefined if relative file path is undefined', async () => {
+            it('leaves document as undefined if relative file path is undefined', async () => {
                 const documentContextObject = {
                     programmingLanguage: 'typescript',
                     cursorState: [],
@@ -976,7 +976,7 @@ describe('AgenticChatController', () => {
 
                 assert.strictEqual(
                     calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
-                        ?.editorState,
+                        ?.editorState?.document,
                     undefined
                 )
             })
@@ -1012,6 +1012,7 @@ describe('AgenticChatController', () => {
                             relativeFilePath: 'file:///test.ts',
                             text: undefined,
                         },
+                        workspaceFolders: [],
                     }
                 )
             })
@@ -1177,7 +1178,7 @@ describe('AgenticChatController', () => {
                 extractDocumentContextStub.restore()
             })
 
-            it('leaves editor state as undefined if cursorState is not passed', async () => {
+            it('leaves cursorState as undefined if cursorState is not passed', async () => {
                 const documentContextObject = {
                     programmingLanguage: 'typescript',
                     cursorState: undefined,
@@ -1198,12 +1199,12 @@ describe('AgenticChatController', () => {
 
                 assert.strictEqual(
                     calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
-                        ?.editorState,
+                        ?.editorState?.cursorState,
                     undefined
                 )
             })
 
-            it('leaves editor state as undefined if relative file path is undefined', async () => {
+            it('leaves document as undefined if relative file path is undefined', async () => {
                 const documentContextObject = {
                     programmingLanguage: 'typescript',
                     cursorState: [],
@@ -1223,7 +1224,7 @@ describe('AgenticChatController', () => {
 
                 assert.strictEqual(
                     calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
-                        ?.editorState,
+                        ?.editorState?.document,
                     undefined
                 )
             })
@@ -1257,6 +1258,7 @@ describe('AgenticChatController', () => {
                             relativeFilePath: 'file:///test.ts',
                             text: undefined,
                         },
+                        workspaceFolders: [],
                     }
                 )
             })

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

@@ -119,7 +119,7 @@ export class AgenticChatController implements ChatHandlers {
     ) {
         this.#features = features
         this.#chatSessionManagementService = chatSessionManagementService
-        this.#triggerContext = new AgenticChatTriggerContext(features.workspace, features.logging)
+        this.#triggerContext = new AgenticChatTriggerContext(features)
         this.#telemetryController = new ChatTelemetryController(features, telemetryService)
         this.#telemetryService = telemetryService
         this.#amazonQServiceManager = amazonQServiceManager

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContext.ts

@@ -22,6 +22,7 @@ import {
 } from '@aws/language-server-runtimes/server-interface'
 import { Features } from '../../types'
 import { DocumentContext, DocumentContextExtractor } from '../../chat/contexts/documentContext'
+import { workspaceUtils } from '@aws/lsp-core'
 
 export interface TriggerContext extends Partial<DocumentContext> {
     userIntent?: UserIntent
@@ -37,11 +38,13 @@ export class AgenticChatTriggerContext {
     private static readonly DEFAULT_CURSOR_STATE: CursorState = { position: { line: 0, character: 0 } }
 
     #workspace: Features['workspace']
+    #lsp: Features['lsp']
     #documentContextExtractor: DocumentContextExtractor
 
-    constructor(workspace: Features['workspace'], logger: Features['logging']) {
+    constructor({ workspace, lsp, logging }: Pick<Features, 'workspace' | 'lsp' | 'logging'> & Partial<Features>) {
         this.#workspace = workspace
-        this.#documentContextExtractor = new DocumentContextExtractor({ logger, workspace })
+        this.#lsp = lsp
+        this.#documentContextExtractor = new DocumentContextExtractor({ logger: logging, workspace })
     }
 
     async getNewTriggerContext(params: ChatParams | InlineChatParams): Promise<TriggerContext> {
@@ -63,7 +66,7 @@ export class AgenticChatTriggerContext {
         additionalContent?: AdditionalContentEntryAddition[]
     ): GenerateAssistantResponseCommandInput {
         const { prompt } = params
-
+        const defaultEditorState = { workspaceFolders: workspaceUtils.getWorkspaceFolders(this.#lsp) }
         const data: GenerateAssistantResponseCommandInput = {
             conversationState: {
                 chatTriggerType: chatTriggerType,
@@ -80,13 +83,17 @@ export class AgenticChatTriggerContext {
                                               programmingLanguage: triggerContext.programmingLanguage,
                                               relativeFilePath: triggerContext.relativeFilePath,
                                           },
+                                          ...defaultEditorState,
                                       },
                                       tools,
                                       additionalContext: additionalContent,
                                   }
                                 : {
                                       tools,
                                       additionalContext: additionalContent,
+                                      editorState: {
+                                          ...defaultEditorState,
+                                      },
                                   },
                         userIntent: triggerContext.userIntent,
                         origin: 'IDE',

server/aws-lsp-codewhisperer/src/language-server/agenticChat/context/agenticChatTriggerContexts.test.ts

@@ -9,11 +9,15 @@ import { TextDocument } from 'vscode-languageserver-textdocument'
 import sinon = require('sinon')
 import { AgenticChatTriggerContext } from './agenticChatTriggerContext'
 import { DocumentContext, DocumentContextExtractor } from '../../chat/contexts/documentContext'
+import { ChatTriggerType, CursorState } from '@amzn/codewhisperer-streaming'
+import { URI } from 'vscode-uri'
+import { InitializeParams } from '@aws/language-server-runtimes/protocol'
 
 describe('AgenticChatTriggerContext', () => {
     let testFeatures: TestFeatures
 
     const filePath = 'file://test.ts'
+    const mockWorkspaceFolders = [{ uri: URI.file('/path/to/my/workspace/').toString(), name: 'myWorkspace' }]
     const mockTSDocument = TextDocument.create(filePath, 'typescript', 1, '')
     const mockDocumentContext: DocumentContext = {
         text: '',
@@ -25,6 +29,9 @@ describe('AgenticChatTriggerContext', () => {
 
     beforeEach(() => {
         testFeatures = new TestFeatures()
+        testFeatures.lsp.getClientInitializeParams.returns({
+            workspaceFolders: mockWorkspaceFolders,
+        } as InitializeParams)
         sinon.stub(DocumentContextExtractor.prototype, 'extractDocumentContext').resolves(mockDocumentContext)
     })
 
@@ -33,7 +40,7 @@ describe('AgenticChatTriggerContext', () => {
     })
 
     it('returns null if text document is not defined in params', async () => {
-        const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new AgenticChatTriggerContext(testFeatures)
 
         const documentContext = await triggerContext.extractDocumentContext({
             cursorState: [
@@ -51,7 +58,7 @@ describe('AgenticChatTriggerContext', () => {
     })
 
     it('returns null if text document is not found', async () => {
-        const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new AgenticChatTriggerContext(testFeatures)
 
         const documentContext = await triggerContext.extractDocumentContext({
             cursorState: [
@@ -71,7 +78,7 @@ describe('AgenticChatTriggerContext', () => {
     })
 
     it('passes default cursor state if no cursor is found', async () => {
-        const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new AgenticChatTriggerContext(testFeatures)
 
         const documentContext = await triggerContext.extractDocumentContext({
             cursorState: [],
@@ -84,7 +91,7 @@ describe('AgenticChatTriggerContext', () => {
     })
 
     it('includes cursor state from the parameters and text document if found', async () => {
-        const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new AgenticChatTriggerContext(testFeatures)
 
         testFeatures.openDocument(mockTSDocument)
         const documentContext = await triggerContext.extractDocumentContext({
@@ -96,4 +103,29 @@ describe('AgenticChatTriggerContext', () => {
 
         assert.deepStrictEqual(documentContext, mockDocumentContext)
     })
+
+    it('includes workspace folders as part of editor state in chat params', async () => {
+        const triggerContext = new AgenticChatTriggerContext(testFeatures)
+        const chatParams = triggerContext.getChatParamsFromTrigger(
+            { tabId: 'tab', prompt: {} },
+            {},
+            ChatTriggerType.MANUAL
+        )
+        const chatParamsWithMore = 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)
+        )
+    })
 })