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 index 111ccdc509f..05d814b94a9 100644 --- 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 { + 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 index 3440d5fa8c8..7d31d9b7bfd 100644 --- 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 index a82f8700eb4..1e81a695308 100644 --- 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 index 5224284b42a..67aff7a14e3 100644 --- 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() - 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() - 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 index 1d5646d3869..add4a9909cd 100644 --- 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 index 330ce9a092e..3bf772e86d5 100644 --- 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 index 3ea36e4c82b..50b1be3626d 100644 --- 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 index 6470584da8e..18cf95d3973 100644 --- 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 + + @JsonRequest("aws/getConfigurationFromServer") + fun getConfigurationFromServer(params: GetConfigurationFromServerParams): CompletableFuture } 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 index 1cb6e2d391e..ca8fffcbb51 100644 --- 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 index 9feb23a2fa4..d54acf55fbe 100644 --- 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 new file mode 100644 index 00000000000..551ddfa97b0 --- /dev/null +++ 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 new file mode 100644 index 00000000000..e762b938d28 --- /dev/null +++ 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) 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 index 275411172f2..7f9e443c14f 100644 --- 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