Skip to content

Commit e4d8f61

Browse files
guvishlEge Ozcan
and
Ege Ozcan
authored
feat: integrate server side project context into agentic chat (#1405)
* 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 <eoozca@amazon.com>
1 parent c8d6edf commit e4d8f61

File tree

7 files changed

+118
-38
lines changed

7 files changed

+118
-38
lines changed
Binary file not shown.

package-lock.json

Lines changed: 56 additions & 31 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

server/aws-lsp-codewhisperer/package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@
2828
},
2929
"dependencies": {
3030
"@amzn/amazon-q-developer-streaming-client": "file:../../core/q-developer-streaming-client/amzn-amazon-q-developer-streaming-client-1.0.0.tgz",
31-
"@amzn/codewhisperer-streaming": "file:../../core/codewhisperer-streaming/amzn-codewhisperer-streaming-1.0.6.tgz",
31+
"@amzn/codewhisperer-streaming": "file:../../core/codewhisperer-streaming/amzn-codewhisperer-streaming-1.0.7.tgz",
3232
"@aws-sdk/util-arn-parser": "^3.723.0",
3333
"@aws-sdk/util-retry": "^3.374.0",
3434
"@aws/chat-client-ui-types": "^0.1.33",

server/aws-lsp-codewhisperer/src/client/token/bearer-token-service.json

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1506,6 +1506,10 @@
15061506
"shape": "ConversationId",
15071507
"documentation": "<p>Unique identifier for the chat conversation stream</p>"
15081508
},
1509+
"workspaceId": {
1510+
"shape": "UUID",
1511+
"documentation": "<p>Unique identifier for remote workspace</p>"
1512+
},
15091513
"history": {
15101514
"shape": "ChatHistory",
15111515
"documentation": "<p>Holds the history of chat messages.</p>"

server/aws-lsp-codewhisperer/src/client/token/codewhispererbearertokenclient.d.ts

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -82,11 +82,11 @@ declare class CodeWhispererBearerTokenClient extends Service {
8282
*/
8383
getCodeAnalysis(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetCodeAnalysisResponse) => void): Request<CodeWhispererBearerTokenClient.Types.GetCodeAnalysisResponse, AWSError>;
8484
/**
85-
*
85+
*
8686
*/
8787
getCodeFixJob(params: CodeWhispererBearerTokenClient.Types.GetCodeFixJobRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse, AWSError>;
8888
/**
89-
*
89+
*
9090
*/
9191
getCodeFixJob(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.GetCodeFixJobResponse, AWSError>;
9292
/**
@@ -130,11 +130,11 @@ declare class CodeWhispererBearerTokenClient extends Service {
130130
*/
131131
listAvailableCustomizations(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.ListAvailableCustomizationsResponse) => void): Request<CodeWhispererBearerTokenClient.Types.ListAvailableCustomizationsResponse, AWSError>;
132132
/**
133-
*
133+
*
134134
*/
135135
listAvailableProfiles(params: CodeWhispererBearerTokenClient.Types.ListAvailableProfilesRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse) => void): Request<CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse, AWSError>;
136136
/**
137-
*
137+
*
138138
*/
139139
listAvailableProfiles(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse) => void): Request<CodeWhispererBearerTokenClient.Types.ListAvailableProfilesResponse, AWSError>;
140140
/**
@@ -186,11 +186,11 @@ declare class CodeWhispererBearerTokenClient extends Service {
186186
*/
187187
startCodeAnalysis(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.StartCodeAnalysisResponse) => void): Request<CodeWhispererBearerTokenClient.Types.StartCodeAnalysisResponse, AWSError>;
188188
/**
189-
*
189+
*
190190
*/
191191
startCodeFixJob(params: CodeWhispererBearerTokenClient.Types.StartCodeFixJobRequest, callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse, AWSError>;
192192
/**
193-
*
193+
*
194194
*/
195195
startCodeFixJob(callback?: (err: AWSError, data: CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse) => void): Request<CodeWhispererBearerTokenClient.Types.StartCodeFixJobResponse, AWSError>;
196196
/**
@@ -463,6 +463,10 @@ declare namespace CodeWhispererBearerTokenClient {
463463
* Unique identifier for the chat conversation stream
464464
*/
465465
conversationId?: ConversationId;
466+
/**
467+
* Unique identifier for remote workspace
468+
*/
469+
workspaceId?: UUID;
466470
/**
467471
* Holds the history of chat messages.
468472
*/

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

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import { RelevantTextDocument } from '@amzn/codewhisperer-streaming'
3333
import { languageByExtension } from '../../../shared/languageDetection'
3434
import { AgenticChatResultStream } from '../agenticChatResultStream'
3535
import { ContextInfo, mergeFileLists, mergeRelevantTextDocuments } from './contextUtils'
36+
import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
3637

3738
export interface TriggerContext extends Partial<DocumentContext> {
3839
userIntent?: UserIntent
@@ -127,6 +128,16 @@ export class AgenticChatTriggerContext {
127128
promptContent = promptContent?.replace(/\*\*@workspace\*\*/, '')
128129
}
129130

131+
// Append remote workspaceId if it exists
132+
// Only append workspaceId to GenerateCompletions when WebSocket client is connected
133+
const remoteWsFolderManager = WorkspaceFolderManager.getInstance()
134+
const workspaceId =
135+
(remoteWsFolderManager &&
136+
remoteWsFolderManager.getWorkspaceState().webSocketClient?.isConnected &&
137+
remoteWsFolderManager.getWorkspaceState().workspaceId) ||
138+
undefined
139+
this.#logging.info(`remote workspaceId: ${workspaceId}`)
140+
130141
// Get workspace documents if @workspace is used
131142
let relevantDocuments = hasWorkspace
132143
? await this.#getRelevantDocuments(promptContent ?? '', chatResultStream)
@@ -176,6 +187,7 @@ export class AgenticChatTriggerContext {
176187

177188
const data: GenerateAssistantResponseCommandInput = {
178189
conversationState: {
190+
workspaceId: workspaceId,
179191
chatTriggerType: chatTriggerType,
180192
currentMessage: {
181193
userInputMessage: {

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

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@ import { ChatTriggerType, CursorState } from '@amzn/codewhisperer-streaming'
1414
import { URI } from 'vscode-uri'
1515
import { InitializeParams } from '@aws/language-server-runtimes/protocol'
1616
import { TestFolder } from '@aws/lsp-core/out/test/testFolder'
17+
import { WorkspaceFolderManager } from '../../workspaceContext/workspaceFolderManager'
1718

1819
describe('AgenticChatTriggerContext', () => {
1920
let testFeatures: TestFeatures
@@ -28,6 +29,7 @@ describe('AgenticChatTriggerContext', () => {
2829
hasCodeSnippet: false,
2930
totalEditorCharacters: 0,
3031
}
32+
let mockWorkspaceFolderManager: any
3133

3234
beforeEach(() => {
3335
testFeatures = new TestFeatures()
@@ -127,6 +129,39 @@ describe('AgenticChatTriggerContext', () => {
127129
mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
128130
)
129131
})
132+
it('includes remote workspaceId if it exists and is connected', async () => {
133+
mockWorkspaceFolderManager = {
134+
getWorkspaceState: sinon.stub(),
135+
}
136+
sinon.stub(WorkspaceFolderManager, 'getInstance').returns(mockWorkspaceFolderManager)
137+
mockWorkspaceFolderManager.getWorkspaceState.returns({
138+
webSocketClient: { isConnected: true },
139+
workspaceId: 'test-workspace-123',
140+
})
141+
const triggerContext = new AgenticChatTriggerContext(testFeatures)
142+
const chatParams = await triggerContext.getChatParamsFromTrigger(
143+
{ tabId: 'tab', prompt: {} },
144+
{},
145+
ChatTriggerType.MANUAL
146+
)
147+
const chatParamsWithMore = await triggerContext.getChatParamsFromTrigger(
148+
{ tabId: 'tab', prompt: {} },
149+
{ cursorState: {} as CursorState, relativeFilePath: '' },
150+
ChatTriggerType.MANUAL
151+
)
152+
153+
assert.deepStrictEqual(
154+
chatParams.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
155+
?.workspaceFolders,
156+
mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
157+
)
158+
assert.deepStrictEqual(
159+
chatParamsWithMore.conversationState?.currentMessage?.userInputMessage?.userInputMessageContext?.editorState
160+
?.workspaceFolders,
161+
mockWorkspaceFolders.map(f => URI.parse(f.uri).fsPath)
162+
)
163+
assert.deepStrictEqual(chatParamsWithMore.conversationState?.workspaceId, 'test-workspace-123')
164+
})
130165
describe('getTextDocument', function () {
131166
let tempFolder: TestFolder
132167

0 commit comments

Comments
 (0)