feat(amazonq): workspace context settings & workspaceId (#5526)

leigaol · web-flow · commit 5456fa0f5fc2 · 2025-04-01T15:37:41.000-07:00
1. Add the settings of workspace context. (String of the settings to be adjusted based on product spec) 2. Let inline completion use workspaceId. Before launch date, future workspace context PRs will target branch feature/q-lsp directly. VSC ref: aws/aws-toolkit-vscode#6894
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/service/CodeWhispererService.kt
@@ -53,6 +53,9 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnection
 import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.CodeWhispererConnection
 import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIMEOUT
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
 import software.aws.toolkits.jetbrains.services.codewhisperer.editor.CodeWhispererEditorManager
@@ -91,6 +94,9 @@ import software.aws.toolkits.resources.message
 import software.aws.toolkits.telemetry.CodewhispererCompletionType
 import software.aws.toolkits.telemetry.CodewhispererSuggestionState
 import software.aws.toolkits.telemetry.CodewhispererTriggerType
+import java.net.URI
+import java.nio.file.Paths
+import java.util.concurrent.CompletableFuture
 import java.util.concurrent.TimeUnit
 
 @Service
@@ -231,7 +237,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                     buildCodeWhispererRequest(
                         requestContext.fileContextInfo,
                         requestContext.awaitSupplementalContext(),
-                        requestContext.customizationArn
+                        requestContext.customizationArn,
+                        requestContext.workspaceId
                     )
                 )
 
@@ -666,7 +673,33 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
         // 5. customization
         val customizationArn = CodeWhispererModelConfigurator.getInstance().activeCustomization(project)?.arn
 
-        return RequestContext(project, editor, triggerTypeInfo, caretPosition, fileContext, supplementalContext, connection, latencyContext, customizationArn)
+        var workspaceId: String? = null
+        try {
+            val workspacesInfos = getWorkspaceIds(project).get().workspaces
+            for (workspaceInfo in workspacesInfos) {
+                val workspaceRootPath = Paths.get(URI(workspaceInfo.workspaceRoot)).toString()
+                if (psiFile.virtualFile.path.startsWith(workspaceRootPath)) {
+                    workspaceId = workspaceInfo.workspaceId
+                    LOG.info { "Found workspaceId from LSP '$workspaceId'" }
+                    break
+                }
+            }
+        } catch (e: Exception) {
+            LOG.warn { "Cannot get workspaceId from LSP'$e'" }
+        }
+        return RequestContext(
+            project, editor, triggerTypeInfo, caretPosition,
+            fileContext, supplementalContext, connection, latencyContext, customizationArn, workspaceId
+        )
+    }
+
+    private fun getWorkspaceIds(project: Project): CompletableFuture<LspServerConfigurations> {
+        val payload = GetConfigurationFromServerParams(
+            section = "aws.q.workspaceContext"
+        )
+        return AmazonQLspService.executeIfRunning(project) { server ->
+            server.getConfigurationFromServer(payload)
+        } ?: (CompletableFuture.failedFuture(IllegalStateException("LSP Server not running")))
     }
 
     fun validateResponse(response: GenerateCompletionsResponse): GenerateCompletionsResponse {
@@ -800,6 +833,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
             fileContextInfo: FileContextInfo,
             supplementalContext: SupplementalContextInfo?,
             customizationArn: String?,
+            workspaceId: String?,
         ): GenerateCompletionsRequest {
             val programmingLanguage = ProgrammingLanguage.builder()
                 .languageName(fileContextInfo.programmingLanguage.toCodeWhispererRuntimeLanguage().languageId)
@@ -828,6 +862,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                 .referenceTrackerConfiguration { it.recommendationsWithReferences(includeCodeWithReference) }
                 .customizationArn(customizationArn)
                 .optOutPreference(getTelemetryOptOutPreference())
+                .workspaceId(workspaceId)
                 .build()
         }
     }
@@ -843,6 +878,7 @@ data class RequestContext(
     val connection: ToolkitConnection?,
     val latencyContext: LatencyContext,
     val customizationArn: String?,
+    val workspaceId: String?,
 ) {
     // TODO: should make the entire getRequestContext() suspend function instead of making supplemental context only
     var supplementalContext: SupplementalContextInfo? = null
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt b/plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt
@@ -138,6 +138,20 @@ class CodeWhispererConfigurable(private val project: Project) :
         }
 
         group(message("aws.settings.codewhisperer.group.q_chat")) {
+            row {
+                checkBox(message("aws.settings.codewhisperer.workspace_context")).apply {
+                    connect.subscribe(
+                        ToolkitConnectionManagerListener.TOPIC,
+                        object : ToolkitConnectionManagerListener {
+                            override fun activeConnectionChanged(newConnection: ToolkitConnection?) {
+                                enabled(isCodeWhispererEnabled(project))
+                            }
+                        }
+                    )
+                    enabled(invoke)
+                    bindSelected(codeWhispererSettings::isWorkspaceContextEnabled, codeWhispererSettings::toggleWorkspaceContextEnabled)
+                }.comment(message("aws.settings.codewhisperer.workspace_context.tooltip"))
+            }
             row {
                 checkBox(message("aws.settings.codewhisperer.project_context")).apply {
                     connect.subscribe(
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt
@@ -176,7 +176,8 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov
             },
             null,
             mock(),
-            aString()
+            aString(),
+            null
         )
         val responseContext = ResponseContext("sessionId")
         val recommendationContext = RecommendationContext(
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererConfigurableTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererConfigurableTest.kt
@@ -39,17 +39,18 @@ class CodeWhispererConfigurableTest : CodeWhispererTestBase() {
 
         val checkboxes = panel.components.filterIsInstance<JCheckBox>()
 
-        assertThat(checkboxes.size).isEqualTo(5)
+        assertThat(checkboxes.size).isEqualTo(6)
         assertThat(checkboxes.map { it.text }).containsExactlyInAnyOrder(
             message("aws.settings.codewhisperer.include_code_with_reference"),
             message("aws.settings.codewhisperer.configurable.opt_out.title"),
             message("aws.settings.codewhisperer.automatic_import_adder"),
+            "Workspace context",
             message("aws.settings.codewhisperer.project_context"),
             message("aws.settings.codewhisperer.project_context_gpu")
         )
 
         val comments = panel.components.filterIsInstance<DslLabel>()
-        assertThat(comments.size).isEqualTo(8)
+        assertThat(comments.size).isEqualTo(9)
 
         mockCodeWhispererEnabledStatus(false)
         ApplicationManager.getApplication().messageBus.syncPublisher(ToolkitConnectionManagerListener.TOPIC)
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt
@@ -212,7 +212,8 @@ class CodeWhispererServiceTest {
                 supplementalContextDeferred = async { mockSupContext },
                 connection = ToolkitConnectionManager.getInstance(projectRule.project).activeConnection(),
                 latencyContext = LatencyContext(),
-                customizationArn = "fake-arn"
+                customizationArn = "fake-arn",
+                workspaceId = null
             )
         )
 
diff --git a/plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt b/plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt
@@ -257,7 +257,8 @@ fun aRequestContext(
             Random.nextLong(),
             aString()
         ),
-        customizationArn = null
+        customizationArn = null,
+        workspaceId = null
     )
 }
 
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImpl.kt
@@ -84,6 +84,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC
                                 CodeWhispererLspConfiguration(
                                     shouldShareData = CodeWhispererSettings.getInstance().isMetricOptIn(),
                                     shouldShareCodeReferences = CodeWhispererSettings.getInstance().isIncludeCodeWithReference(),
+                                    shouldEnableWorkspaceContext = CodeWhispererSettings.getInstance().isWorkspaceContextEnabled()
                                 )
                             )
                         }
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageServer.kt
@@ -7,6 +7,8 @@ import org.eclipse.lsp4j.jsonrpc.messages.ResponseMessage
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
 import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.services.LanguageServer
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.GetConfigurationFromServerParams
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.LspServerConfigurations
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.credentials.UpdateCredentialsPayload
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws.dependencies.DidChangeDependencyPathsParams
 import java.util.concurrent.CompletableFuture
@@ -24,4 +26,7 @@ interface AmazonQLanguageServer : LanguageServer {
 
     @JsonNotification("aws/credentials/token/delete")
     fun deleteTokenCredentials(): CompletableFuture<Unit>
+
+    @JsonRequest("aws/getConfigurationFromServer")
+    fun getConfigurationFromServer(params: GetConfigurationFromServerParams): CompletableFuture<LspServerConfigurations>
 }
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspConstants.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspConstants.kt
@@ -8,4 +8,5 @@ object AmazonQLspConstants {
     const val LSP_CW_CONFIGURATION_KEY = "aws.codeWhisperer"
     const val LSP_CW_OPT_OUT_KEY = "shareCodeWhispererContentWithAWS"
     const val LSP_CODE_REFERENCES_OPT_OUT_KEY = "includeSuggestionsWithCodeReferences"
+    const val LSP_WORKSPACE_CONTEXT_ENABLED_KEY = "workspaceContext"
 }
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/CodeWhispererLspConfiguration.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/CodeWhispererLspConfiguration.kt
@@ -9,6 +9,9 @@ data class CodeWhispererLspConfiguration(
     @SerializedName(AmazonQLspConstants.LSP_CW_OPT_OUT_KEY)
     val shouldShareData: Boolean? = null,
 
+    @SerializedName(AmazonQLspConstants.LSP_WORKSPACE_CONTEXT_ENABLED_KEY)
+    val shouldEnableWorkspaceContext: Boolean? = null,
+
     @SerializedName(AmazonQLspConstants.LSP_CODE_REFERENCES_OPT_OUT_KEY)
     val shouldShareCodeReferences: Boolean? = null,
 )
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/GetConfigurationFromServerParams.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/GetConfigurationFromServerParams.kt
@@ -0,0 +1,8 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws
+
+data class GetConfigurationFromServerParams(
+    val section: String,
+)
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/LspServerConfigurations.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/model/aws/LspServerConfigurations.kt
@@ -0,0 +1,10 @@
+// Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package software.aws.toolkits.jetbrains.services.amazonq.lsp.model.aws
+
+// This represents each item in the array
+data class WorkspaceInfo(val workspaceRoot: String, val workspaceId: String)
+
+// This represents the entire array
+data class LspServerConfigurations(val workspaces: List<WorkspaceInfo>)
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/settings/CodeWhispererSettings.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/settings/CodeWhispererSettings.kt
@@ -91,6 +91,17 @@ class CodeWhispererSettings : PersistentStateComponent<CodeWhispererConfiguratio
         }
     }
 
+    fun toggleWorkspaceContextEnabled(value: Boolean) {
+        state.value[CodeWhispererConfigurationType.IsWorkspaceContextEnabled] = value
+        ProjectManager.getInstance().openProjects.forEach {
+            if (it.isDisposed) {
+                return@forEach
+            }
+            AmazonQLspService.didChangeConfiguration(it)
+        }
+    }
+
+    fun isWorkspaceContextEnabled() = state.value.getOrDefault(CodeWhispererConfigurationType.IsWorkspaceContextEnabled, true)
     fun isProjectContextEnabled() = state.value.getOrDefault(CodeWhispererConfigurationType.IsProjectContextEnabled, false)
 
     private fun hasEnabledProjectContextOnce() = state.value.getOrDefault(CodeWhispererConfigurationType.HasEnabledProjectContextOnce, false)
@@ -199,6 +210,7 @@ enum class CodeWhispererConfigurationType {
     HasEnabledProjectContextOnce,
     IsQPrioritizedForTabAccept,
     IsTabAcceptPriorityNotificationShownOnce,
+    IsWorkspaceContextEnabled,
 }
 
 enum class CodeWhispererStringConfigurationType {
diff --git a/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImplTest.kt b/plugins/amazonq/shared/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLanguageClientImplTest.kt
@@ -62,25 +62,43 @@ class AmazonQLanguageClientImplTest {
     @Test
     fun `configuration for codeWhisperer respects opt-out`() {
         CodeWhispererSettings.getInstance().toggleMetricOptIn(false)
+        CodeWhispererSettings.getInstance().toggleWorkspaceContextEnabled(true)
         assertThat(sut.configuration(configurationParams("aws.codeWhisperer")).get())
             .singleElement()
             .isEqualTo(
                 CodeWhispererLspConfiguration(
                     shouldShareData = false,
-                    shouldShareCodeReferences = false
+                    shouldShareCodeReferences = false,
+                    shouldEnableWorkspaceContext = true
                 )
             )
     }
 
     @Test
     fun `configuration for codeWhisperer respects opt-in`() {
         CodeWhispererSettings.getInstance().toggleMetricOptIn(true)
+        CodeWhispererSettings.getInstance().toggleWorkspaceContextEnabled(true)
         assertThat(sut.configuration(configurationParams("aws.codeWhisperer")).get())
             .singleElement()
             .isEqualTo(
                 CodeWhispererLspConfiguration(
                     shouldShareData = true,
-                    shouldShareCodeReferences = false
+                    shouldShareCodeReferences = false,
+                    shouldEnableWorkspaceContext = true
+                )
+            )
+    }
+
+    @Test
+    fun `configuration for workspace context respects opt-in`() {
+        CodeWhispererSettings.getInstance().toggleWorkspaceContextEnabled(false)
+        assertThat(sut.configuration(configurationParams("aws.codeWhisperer")).get())
+            .singleElement()
+            .isEqualTo(
+                CodeWhispererLspConfiguration(
+                    shouldShareData = true,
+                    shouldShareCodeReferences = false,
+                    shouldEnableWorkspaceContext = false
                 )
             )
     }
diff --git a/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties b/plugins/core/resources/resources/software/aws/toolkits/resources/MessagesBundle.properties
@@ -297,6 +297,8 @@ aws.settings.codewhisperer.project_context_index_max_size.tooltip=The maximum si
 aws.settings.codewhisperer.project_context_index_thread=Workspace index worker threads
 aws.settings.codewhisperer.project_context_index_thread.tooltip=Number of worker threads of Amazon Q local index process. Set to 0 to use system default worker threads for balanced performance. Please restart or reload IntelliJ after changing worker threads.
 aws.settings.codewhisperer.warning=To use Amazon Q, login with AWS Builder ID or AWS IAM Identity Center.
+aws.settings.codewhisperer.workspace_context=Workspace context
+aws.settings.codewhisperer.workspace_context.tooltip=When checked, Amazon Q will enable server side project context.
 aws.settings.dynamic_resources_configurable.clear_all=Clear All
 aws.settings.dynamic_resources_configurable.select_all=Select All
 aws.settings.dynamic_resources_configurable.suggest_types.dialog.message=Please suggest additional AWS resource types (e.g. AWS::S3::Bucket)\nyou would like to see supported in future releases.\n\n(max length: 2000 chars)

Original file line number	Diff line number	Diff line change
`@@ -212,7 +212,8 @@ class CodeWhispererServiceTest {`
`212`	`212`	`supplementalContextDeferred = async { mockSupContext },`
`213`	`213`	`connection = ToolkitConnectionManager.getInstance(projectRule.project).activeConnection(),`
`214`	`214`	`latencyContext = LatencyContext(),`
`215`		`- customizationArn = "fake-arn"`
	`215`	`+ customizationArn = "fake-arn",`
	`216`	`+ workspaceId = null`
`216`	`217`	`)`
`217`	`218`	`)`
`218`	`219`
Original file line number	Diff line number	Diff line change
`@@ -257,7 +257,8 @@ fun aRequestContext(`
`257`	`257`	`Random.nextLong(),`
`258`	`258`	`aString()`
`259`	`259`	`),`
`260`		`- customizationArn = null`
	`260`	`+ customizationArn = null,`
	`261`	`+ workspaceId = null`
`261`	`262`	`)`
`262`	`263`	`}`
`263`	`264`
Original file line number	Diff line number	Diff line change
`@@ -84,6 +84,7 @@ class AmazonQLanguageClientImpl(private val project: Project) : AmazonQLanguageC`
`84`	`84`	`CodeWhispererLspConfiguration(`
`85`	`85`	`shouldShareData = CodeWhispererSettings.getInstance().isMetricOptIn(),`
`86`	`86`	`shouldShareCodeReferences = CodeWhispererSettings.getInstance().isIncludeCodeWithReference(),`
	`87`	`+ shouldEnableWorkspaceContext = CodeWhispererSettings.getInstance().isWorkspaceContextEnabled()`
`87`	`88`	`)`
`88`	`89`	`)`
`89`	`90`	`}`
Original file line number	Diff line number	Diff line change
`@@ -8,4 +8,5 @@ object AmazonQLspConstants {`
`8`	`8`	`const val LSP_CW_CONFIGURATION_KEY = "aws.codeWhisperer"`
`9`	`9`	`const val LSP_CW_OPT_OUT_KEY = "shareCodeWhispererContentWithAWS"`
`10`	`10`	`const val LSP_CODE_REFERENCES_OPT_OUT_KEY = "includeSuggestionsWithCodeReferences"`
	`11`	`+ const val LSP_WORKSPACE_CONTEXT_ENABLED_KEY = "workspaceContext"`
`11`	`12`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,9 @@ data class CodeWhispererLspConfiguration(`
`9`	`9`	`@SerializedName(AmazonQLspConstants.LSP_CW_OPT_OUT_KEY)`
`10`	`10`	`val shouldShareData: Boolean? = null,`
`11`	`11`
	`12`	`+ @SerializedName(AmazonQLspConstants.LSP_WORKSPACE_CONTEXT_ENABLED_KEY)`
	`13`	`+ val shouldEnableWorkspaceContext: Boolean? = null,`
	`14`	`+`
`12`	`15`	`@SerializedName(AmazonQLspConstants.LSP_CODE_REFERENCES_OPT_OUT_KEY)`
`13`	`16`	`val shouldShareCodeReferences: Boolean? = null,`
`14`	`17`	`)`