feat: adding a check before invoking workspace context (#1227)

suprajaven · Supraja Venkatesh · web-flow · commit 3202b6e0654a · 2025-05-07T10:11:38.000+02:00
Co-authored-by: Supraja Venkatesh &lt;suprajav@amazon.com&gt;
diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/chatController.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/chatController.ts
@@ -86,7 +86,7 @@ export class ChatController implements ChatHandlers {
     ) {
         this.#features = features
         this.#chatSessionManagementService = chatSessionManagementService
-        this.#triggerContext = new QChatTriggerContext(features.workspace, features.logging)
+        this.#triggerContext = new QChatTriggerContext(features.workspace, features.logging, serviceManager)
         this.#telemetryController = new ChatTelemetryController(features, telemetryService)
         this.#telemetryService = telemetryService
         this.#serviceManager = serviceManager
diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/contexts/triggerContext.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/contexts/triggerContext.ts
@@ -6,6 +6,7 @@ import { DocumentContext, DocumentContextExtractor } from './documentContext'
 import { SendMessageCommandInput } from '../../../shared/streamingClientService'
 import { LocalProjectContextController } from '../../../shared/localProjectContextController'
 import { convertChunksToRelevantTextDocuments } from '../tools/relevantTextDocuments'
+import { AmazonQBaseServiceManager as AmazonQServiceManager } from '../../../shared/amazonQServiceManager/BaseAmazonQServiceManager'
 
 export interface TriggerContext extends Partial<DocumentContext> {
     userIntent?: UserIntent
@@ -21,7 +22,11 @@ export class QChatTriggerContext {
     #documentContextExtractor: DocumentContextExtractor
     #logger: Features['logging']
 
-    constructor(workspace: Features['workspace'], logger: Features['logging']) {
+    constructor(
+        workspace: Features['workspace'],
+        logger: Features['logging'],
+        private amazonQServiceManager?: AmazonQServiceManager
+    ) {
         this.#workspace = workspace
         this.#documentContextExtractor = new DocumentContextExtractor({ logger, workspace })
         this.#logger = logger
@@ -120,6 +125,19 @@ export class QChatTriggerContext {
     async extractProjectContext(query?: string): Promise<RelevantTextDocument[]> {
         if (query) {
             try {
+                let enableWorkspaceContext = true
+
+                if (this.amazonQServiceManager) {
+                    const config = this.amazonQServiceManager.getConfiguration()
+                    if (config.projectContext?.enableLocalIndexing === false) {
+                        enableWorkspaceContext = false
+                    }
+                }
+
+                if (!enableWorkspaceContext) {
+                    this.#logger.debug('Workspace context is disabled, skipping project context extraction')
+                    return []
+                }
                 const contextController = await LocalProjectContextController.getInstance()
                 const resp = await contextController.queryVectorIndex({ query })
                 return convertChunksToRelevantTextDocuments(resp)
diff --git a/server/aws-lsp-codewhisperer/src/language-server/chat/contexts/triggerContexts.test.ts b/server/aws-lsp-codewhisperer/src/language-server/chat/contexts/triggerContexts.test.ts
@@ -4,9 +4,11 @@ import assert = require('assert')
 import { TextDocument } from 'vscode-languageserver-textdocument'
 import { DocumentContext, DocumentContextExtractor } from './documentContext'
 import sinon = require('sinon')
+import { LocalProjectContextController } from '../../../shared/localProjectContextController'
 
 describe('QChatTriggerContext', () => {
     let testFeatures: TestFeatures
+    let amazonQServiceManager: any
 
     const filePath = 'file://test.ts'
     const mockTSDocument = TextDocument.create(filePath, 'typescript', 1, '')
@@ -20,6 +22,13 @@ describe('QChatTriggerContext', () => {
 
     beforeEach(() => {
         testFeatures = new TestFeatures()
+        amazonQServiceManager = {
+            getConfiguration: sinon.stub().returns({
+                projectContext: {
+                    enableLocalIndexing: true,
+                },
+            }),
+        }
         sinon.stub(DocumentContextExtractor.prototype, 'extractDocumentContext').resolves(mockDocumentContext)
     })
 
@@ -28,7 +37,11 @@ describe('QChatTriggerContext', () => {
     })
 
     it('returns null if text document is not defined in params', async () => {
-        const triggerContext = new QChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new QChatTriggerContext(
+            testFeatures.workspace,
+            testFeatures.logging,
+            amazonQServiceManager
+        )
 
         const documentContext = await triggerContext.extractDocumentContext({
             cursorState: [
@@ -46,7 +59,11 @@ describe('QChatTriggerContext', () => {
     })
 
     it('returns null if text document is not found', async () => {
-        const triggerContext = new QChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new QChatTriggerContext(
+            testFeatures.workspace,
+            testFeatures.logging,
+            amazonQServiceManager
+        )
 
         const documentContext = await triggerContext.extractDocumentContext({
             cursorState: [
@@ -66,7 +83,11 @@ describe('QChatTriggerContext', () => {
     })
 
     it('passes default cursor state if no cursor is found', async () => {
-        const triggerContext = new QChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new QChatTriggerContext(
+            testFeatures.workspace,
+            testFeatures.logging,
+            amazonQServiceManager
+        )
 
         const documentContext = await triggerContext.extractDocumentContext({
             cursorState: [],
@@ -79,7 +100,11 @@ describe('QChatTriggerContext', () => {
     })
 
     it('includes cursor state from the parameters and text document if found', async () => {
-        const triggerContext = new QChatTriggerContext(testFeatures.workspace, testFeatures.logging)
+        const triggerContext = new QChatTriggerContext(
+            testFeatures.workspace,
+            testFeatures.logging,
+            amazonQServiceManager
+        )
 
         testFeatures.openDocument(mockTSDocument)
         const documentContext = await triggerContext.extractDocumentContext({
@@ -91,4 +116,26 @@ describe('QChatTriggerContext', () => {
 
         assert.deepStrictEqual(documentContext, mockDocumentContext)
     })
+
+    it('should not extract project context when workspace context is disabled', async () => {
+        amazonQServiceManager.getConfiguration.returns({
+            projectContext: {
+                enableLocalIndexing: false,
+            },
+        })
+
+        const triggerContext = new QChatTriggerContext(
+            testFeatures.workspace,
+            testFeatures.logging,
+            amazonQServiceManager
+        )
+
+        const getInstanceStub = sinon.stub(LocalProjectContextController, 'getInstance')
+
+        const result = await triggerContext.extractProjectContext('test query')
+
+        sinon.assert.notCalled(getInstanceStub)
+
+        assert.deepStrictEqual(result, [])
+    })
 })
diff --git a/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts b/server/aws-lsp-codewhisperer/src/language-server/inline-completion/codeWhispererServer.ts
@@ -336,7 +336,7 @@ export const CodewhispererServerFactory =
                     // supplementalContext available only via token authentication
                     const supplementalContextPromise =
                         codeWhispererService instanceof CodeWhispererServiceToken
-                            ? fetchSupplementalContext(textDocument, params.position, workspace, logging, token)
+                            ? fetchSupplementalContext(textDocument, params.position, workspace, logging, token, amazonQServiceManager)
                             : Promise.resolve(undefined)
 
                     let requestContext: GenerateSuggestionsRequest = {
diff --git a/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.test.ts b/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.test.ts
@@ -249,8 +249,16 @@ describe('crossFileContextUtil', function () {
     })
     describe('codemapContext', () => {
         let sandbox: sinon.SinonSandbox
+        let amazonQServiceManager: any
         beforeEach(() => {
             sandbox = sinon.createSandbox()
+            amazonQServiceManager = {
+                getConfiguration: sandbox.stub().returns({
+                    projectContext: {
+                        enableLocalIndexing: true,
+                    },
+                }),
+            }
         })
         afterEach(() => {
             sandbox.restore()
@@ -271,13 +279,47 @@ describe('crossFileContextUtil', function () {
                 document,
                 { line: 0, character: 0 },
                 features.workspace,
-                fakeCancellationToken
+                fakeCancellationToken,
+                amazonQServiceManager
+            )
+            assert.deepStrictEqual(result, {
+                supplementalContextItems: [],
+                strategy: 'Empty',
+            })
+        })
+
+        it('should return Empty strategy when workspace context is disabled', async () => {
+            const document = TextDocument.create(
+                'file:///testfile.java',
+                'java',
+                1,
+                'line_1\nline_2\nline_3\nline_4\nline_5\nline_6\nline_7'
             )
+
+            amazonQServiceManager.getConfiguration.returns({
+                projectContext: {
+                    enableLocalIndexing: false,
+                },
+            })
+
+            const instanceStub = sandbox.stub(LocalProjectContextController, 'getInstance')
+
+            const result = await crossFile.codemapContext(
+                document,
+                { line: 0, character: 0 },
+                features.workspace,
+                fakeCancellationToken,
+                amazonQServiceManager
+            )
+
+            sinon.assert.notCalled(instanceStub)
+
             assert.deepStrictEqual(result, {
                 supplementalContextItems: [],
                 strategy: 'Empty',
             })
         })
+
         it('should return codemap strategy when project context exists', async () => {
             const document = TextDocument.create(
                 'file:///testfile.java',
@@ -296,7 +338,8 @@ describe('crossFileContextUtil', function () {
                 document,
                 { line: 0, character: 0 },
                 features.workspace,
-                fakeCancellationToken
+                fakeCancellationToken,
+                amazonQServiceManager
             )
             assert.deepStrictEqual(result, {
                 supplementalContextItems: [{ content: 'someOtherContet', filePath: '/path/', score: 29.879 }],
@@ -321,7 +364,8 @@ describe('crossFileContextUtil', function () {
                 document,
                 { line: 0, character: 0 },
                 features.workspace,
-                fakeCancellationToken
+                fakeCancellationToken,
+                amazonQServiceManager
             )
             assert.deepStrictEqual(result, {
                 supplementalContextItems: [
diff --git a/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.ts b/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/crossFileContextUtil.ts
@@ -29,6 +29,7 @@ import { LocalProjectContextController } from '../../shared/localProjectContextC
 import { QueryInlineProjectContextRequestV2 } from 'local-indexing'
 import { URI } from 'vscode-uri'
 import { waitUntil } from '@aws/lsp-core/out/util/timeoutUtils'
+import { AmazonQBaseServiceManager } from '../amazonQServiceManager/BaseAmazonQServiceManager'
 
 type CrossFileSupportedLanguage =
     | 'java'
@@ -65,7 +66,8 @@ export async function fetchSupplementalContextForSrc(
     document: TextDocument,
     position: Position,
     workspace: Workspace,
-    cancellationToken: CancellationToken
+    cancellationToken: CancellationToken,
+    amazonQServiceManager?: AmazonQBaseServiceManager
 ): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
     const supplementalContextConfig = getSupplementalContextConfig(document.languageId)
 
@@ -74,15 +76,17 @@ export async function fetchSupplementalContextForSrc(
     }
     //TODO: add logic for other strategies once available
     if (supplementalContextConfig === 'codemap') {
-        return await codemapContext(document, position, workspace, cancellationToken)
+        return await codemapContext(document, position, workspace, cancellationToken, amazonQServiceManager)
     }
+    return { supplementalContextItems: [], strategy: 'Empty' }
 }
 
 export async function codemapContext(
     document: TextDocument,
     position: Position,
     workspace: Workspace,
-    cancellationToken: CancellationToken
+    cancellationToken: CancellationToken,
+    amazonQServiceManager?: AmazonQBaseServiceManager
 ): Promise<Pick<CodeWhispererSupplementalContext, 'supplementalContextItems' | 'strategy'> | undefined> {
     let strategy: SupplementalContextStrategy = 'Empty'
 
@@ -95,7 +99,7 @@ export async function codemapContext(
 
     const projectContextPromise = waitUntil(
         async function () {
-            return await fetchProjectContext(document, position, 'codemap')
+            return await fetchProjectContext(document, position, 'codemap', amazonQServiceManager)
         },
         { timeout: supplementalContextTimeoutInMs, interval: 5, truthy: false }
     )
@@ -133,7 +137,8 @@ export async function codemapContext(
 export async function fetchProjectContext(
     document: TextDocument,
     position: Position,
-    target: 'default' | 'codemap' | 'bm25'
+    target: 'default' | 'codemap' | 'bm25',
+    amazonQServiceManager?: AmazonQBaseServiceManager
 ): Promise<CodeWhispererSupplementalContextItem[]> {
     const inputChunk: Chunk = getInputChunk(document, position, crossFileContextConfig.numberOfLinesEachChunk)
     const fsPath = URI.parse(document.uri).fsPath
@@ -143,6 +148,17 @@ export async function fetchProjectContext(
         filePath: fsPath,
         target,
     }
+    let enableWorkspaceContext = true
+
+    if (amazonQServiceManager) {
+        const config = amazonQServiceManager.getConfiguration()
+        if (config.projectContext?.enableLocalIndexing === false) {
+            enableWorkspaceContext = false
+        }
+    }
+    if (!enableWorkspaceContext) {
+        return []
+    }
 
     try {
         controller = await LocalProjectContextController.getInstance()
diff --git a/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/supplementalContextUtil.test.ts b/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/supplementalContextUtil.test.ts
diff --git a/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/supplementalContextUtil.ts b/server/aws-lsp-codewhisperer/src/shared/supplementalContextUtil/supplementalContextUtil.ts
