Skip to content

Commit f300ca5

Browse files
HweinstockHweinstock
authored
fix: add workspace folders as context for agentic-chat (#995)
* fix: add workspace folders as context * test: add test for fix and update existing tests * fix: parse URI result to be platform agnostic * fix: properly merge to avoid compile error
1 parent 3010da0 commit f300ca5

File tree

5 files changed

+61
-16
lines changed

5 files changed

+61
-16
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,10 @@ export function getEntryPath(entry: Dirent) {
8585
}
8686

8787
// TODO: port this to runtimes?
88+
export function getWorkspaceFolders(lsp: Features['lsp']): string[] {
89+
return lsp.getClientInitializeParams()?.workspaceFolders?.map(({ uri }) => URI.parse(uri).fsPath) ?? []
90+
}
91+
8892
export async function inWorkspace(workspace: Features['workspace'], filepath: string) {
8993
return (await workspace.getTextDocument(URI.file(filepath).toString())) !== undefined
9094
}

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

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -926,7 +926,7 @@ describe('AgenticChatController', () => {
926926
extractDocumentContextStub.restore()
927927
})
928928

929-
it('leaves editor state as undefined if cursorState is not passed', async () => {
929+
it('leaves cursorState as undefined if cursorState is not passed', async () => {
930930
const documentContextObject = {
931931
programmingLanguage: 'typescript',
932932
cursorState: undefined,
@@ -949,12 +949,12 @@ describe('AgenticChatController', () => {
949949

950950
assert.strictEqual(
951951
calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
952-
?.editorState,
952+
?.editorState?.cursorState,
953953
undefined
954954
)
955955
})
956956

957-
it('leaves editor state as undefined if relative file path is undefined', async () => {
957+
it('leaves document as undefined if relative file path is undefined', async () => {
958958
const documentContextObject = {
959959
programmingLanguage: 'typescript',
960960
cursorState: [],
@@ -976,7 +976,7 @@ describe('AgenticChatController', () => {
976976

977977
assert.strictEqual(
978978
calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
979-
?.editorState,
979+
?.editorState?.document,
980980
undefined
981981
)
982982
})
@@ -1012,6 +1012,7 @@ describe('AgenticChatController', () => {
10121012
relativeFilePath: 'file:///test.ts',
10131013
text: undefined,
10141014
},
1015+
workspaceFolders: [],
10151016
}
10161017
)
10171018
})
@@ -1177,7 +1178,7 @@ describe('AgenticChatController', () => {
11771178
extractDocumentContextStub.restore()
11781179
})
11791180

1180-
it('leaves editor state as undefined if cursorState is not passed', async () => {
1181+
it('leaves cursorState as undefined if cursorState is not passed', async () => {
11811182
const documentContextObject = {
11821183
programmingLanguage: 'typescript',
11831184
cursorState: undefined,
@@ -1198,12 +1199,12 @@ describe('AgenticChatController', () => {
11981199

11991200
assert.strictEqual(
12001201
calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
1201-
?.editorState,
1202+
?.editorState?.cursorState,
12021203
undefined
12031204
)
12041205
})
12051206

1206-
it('leaves editor state as undefined if relative file path is undefined', async () => {
1207+
it('leaves document as undefined if relative file path is undefined', async () => {
12071208
const documentContextObject = {
12081209
programmingLanguage: 'typescript',
12091210
cursorState: [],
@@ -1223,7 +1224,7 @@ describe('AgenticChatController', () => {
12231224

12241225
assert.strictEqual(
12251226
calledRequestInput.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext
1226-
?.editorState,
1227+
?.editorState?.document,
12271228
undefined
12281229
)
12291230
})
@@ -1257,6 +1258,7 @@ describe('AgenticChatController', () => {
12571258
relativeFilePath: 'file:///test.ts',
12581259
text: undefined,
12591260
},
1261+
workspaceFolders: [],
12601262
}
12611263
)
12621264
})

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ export class AgenticChatController implements ChatHandlers {
119119
) {
120120
this.#features = features
121121
this.#chatSessionManagementService = chatSessionManagementService
122-
this.#triggerContext = new AgenticChatTriggerContext(features.workspace, features.logging)
122+
this.#triggerContext = new AgenticChatTriggerContext(features)
123123
this.#telemetryController = new ChatTelemetryController(features, telemetryService)
124124
this.#telemetryService = telemetryService
125125
this.#amazonQServiceManager = amazonQServiceManager

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

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import {
2222
} from '@aws/language-server-runtimes/server-interface'
2323
import { Features } from '../../types'
2424
import { DocumentContext, DocumentContextExtractor } from '../../chat/contexts/documentContext'
25+
import { workspaceUtils } from '@aws/lsp-core'
2526

2627
export interface TriggerContext extends Partial<DocumentContext> {
2728
userIntent?: UserIntent
@@ -37,11 +38,13 @@ export class AgenticChatTriggerContext {
3738
private static readonly DEFAULT_CURSOR_STATE: CursorState = { position: { line: 0, character: 0 } }
3839

3940
#workspace: Features['workspace']
41+
#lsp: Features['lsp']
4042
#documentContextExtractor: DocumentContextExtractor
4143

42-
constructor(workspace: Features['workspace'], logger: Features['logging']) {
44+
constructor({ workspace, lsp, logging }: Pick<Features, 'workspace' | 'lsp' | 'logging'> & Partial<Features>) {
4345
this.#workspace = workspace
44-
this.#documentContextExtractor = new DocumentContextExtractor({ logger, workspace })
46+
this.#lsp = lsp
47+
this.#documentContextExtractor = new DocumentContextExtractor({ logger: logging, workspace })
4548
}
4649

4750
async getNewTriggerContext(params: ChatParams | InlineChatParams): Promise<TriggerContext> {
@@ -63,7 +66,7 @@ export class AgenticChatTriggerContext {
6366
additionalContent?: AdditionalContentEntryAddition[]
6467
): GenerateAssistantResponseCommandInput {
6568
const { prompt } = params
66-
69+
const defaultEditorState = { workspaceFolders: workspaceUtils.getWorkspaceFolders(this.#lsp) }
6770
const data: GenerateAssistantResponseCommandInput = {
6871
conversationState: {
6972
chatTriggerType: chatTriggerType,
@@ -80,13 +83,17 @@ export class AgenticChatTriggerContext {
8083
programmingLanguage: triggerContext.programmingLanguage,
8184
relativeFilePath: triggerContext.relativeFilePath,
8285
},
86+
...defaultEditorState,
8387
},
8488
tools,
8589
additionalContext: additionalContent,
8690
}
8791
: {
8892
tools,
8993
additionalContext: additionalContent,
94+
editorState: {
95+
...defaultEditorState,
96+
},
9097
},
9198
userIntent: triggerContext.userIntent,
9299
origin: 'IDE',

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

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,15 @@ import { TextDocument } from 'vscode-languageserver-textdocument'
99
import sinon = require('sinon')
1010
import { AgenticChatTriggerContext } from './agenticChatTriggerContext'
1111
import { DocumentContext, DocumentContextExtractor } from '../../chat/contexts/documentContext'
12+
import { ChatTriggerType, CursorState } from '@amzn/codewhisperer-streaming'
13+
import { URI } from 'vscode-uri'
14+
import { InitializeParams } from '@aws/language-server-runtimes/protocol'
1215

1316
describe('AgenticChatTriggerContext', () => {
1417
let testFeatures: TestFeatures
1518

1619
const filePath = 'file://test.ts'
20+
const mockWorkspaceFolders = [{ uri: URI.file('/path/to/my/workspace/').toString(), name: 'myWorkspace' }]
1721
const mockTSDocument = TextDocument.create(filePath, 'typescript', 1, '')
1822
const mockDocumentContext: DocumentContext = {
1923
text: '',
@@ -25,6 +29,9 @@ describe('AgenticChatTriggerContext', () => {
2529

2630
beforeEach(() => {
2731
testFeatures = new TestFeatures()
32+
testFeatures.lsp.getClientInitializeParams.returns({
33+
workspaceFolders: mockWorkspaceFolders,
34+
} as InitializeParams)
2835
sinon.stub(DocumentContextExtractor.prototype, 'extractDocumentContext').resolves(mockDocumentContext)
2936
})
3037

@@ -33,7 +40,7 @@ describe('AgenticChatTriggerContext', () => {
3340
})
3441

3542
it('returns null if text document is not defined in params', async () => {
36-
const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
43+
const triggerContext = new AgenticChatTriggerContext(testFeatures)
3744

3845
const documentContext = await triggerContext.extractDocumentContext({
3946
cursorState: [
@@ -51,7 +58,7 @@ describe('AgenticChatTriggerContext', () => {
5158
})
5259

5360
it('returns null if text document is not found', async () => {
54-
const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
61+
const triggerContext = new AgenticChatTriggerContext(testFeatures)
5562

5663
const documentContext = await triggerContext.extractDocumentContext({
5764
cursorState: [
@@ -71,7 +78,7 @@ describe('AgenticChatTriggerContext', () => {
7178
})
7279

7380
it('passes default cursor state if no cursor is found', async () => {
74-
const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
81+
const triggerContext = new AgenticChatTriggerContext(testFeatures)
7582

7683
const documentContext = await triggerContext.extractDocumentContext({
7784
cursorState: [],
@@ -84,7 +91,7 @@ describe('AgenticChatTriggerContext', () => {
8491
})
8592

8693
it('includes cursor state from the parameters and text document if found', async () => {
87-
const triggerContext = new AgenticChatTriggerContext(testFeatures.workspace, testFeatures.logging)
94+
const triggerContext = new AgenticChatTriggerContext(testFeatures)
8895

8996
testFeatures.openDocument(mockTSDocument)
9097
const documentContext = await triggerContext.extractDocumentContext({
@@ -96,4 +103,29 @@ describe('AgenticChatTriggerContext', () => {
96103

97104
assert.deepStrictEqual(documentContext, mockDocumentContext)
98105
})
106+
107+
it('includes workspace folders as part of editor state in chat params', async () => {
108+
const triggerContext = new AgenticChatTriggerContext(testFeatures)
109+
const chatParams = triggerContext.getChatParamsFromTrigger(
110+
{ tabId: 'tab', prompt: {} },
111+
{},
112+
ChatTriggerType.MANUAL
113+
)
114+
const chatParamsWithMore = triggerContext.getChatParamsFromTrigger(
115+
{ tabId: 'tab', prompt: {} },
116+
{ cursorState: {} as CursorState, relativeFilePath: '' },
117+
ChatTriggerType.MANUAL
118+
)
119+
120+
assert.deepStrictEqual(
121+
chatParams.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
122+
?.workspaceFolders,
123+
mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
124+
)
125+
assert.deepStrictEqual(
126+
chatParamsWithMore.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
127+
?.workspaceFolders,
128+
mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
129+
)
130+
})
99131
})

0 commit comments

Comments
 (0)