Merge branch 'feature/q-lsp' into samgst/q-lsp-fix-241-test

samgst-amazon · web-flow · commit cee0d53daa85 · 2025-04-01T16:14:37.000-07:00
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`	`)`