feat(lsp): Addition of LSP for Amazon Q

feat(lsp): Addition of LSP for Amazon Q

Author: manodnyab

gradle/libs.versions.toml

@@ -106,6 +106,7 @@ kotlin-reflect = { module = "org.jetbrains.kotlin:kotlin-reflect", version.ref =
 kotlin-stdLibJdk8 = { module = "org.jetbrains.kotlin:kotlin-stdlib-jdk8", version.ref = "kotlin" }
 kotlin-test = { module = "org.jetbrains.kotlin:kotlin-test-junit", version.ref = "kotlin" }
 mockito-core = { module = "org.mockito:mockito-core", version.ref = "mockito" }
+mockito-junit-jupiter = { module = "org.mockito:mockito-junit-jupiter", version.ref = "mockito" }
 mockito-kotlin = { module = "org.mockito.kotlin:mockito-kotlin", version.ref = "mockitoKotlin" }
 mockk = { module = "io.mockk:mockk", version.ref="mockk" }
 nimbus-jose-jwt = {module = "com.nimbusds:nimbus-jose-jwt", version.ref = "nimbus-jose-jwt"}
@@ -121,7 +122,7 @@ zjsonpatch = { module = "com.flipkart.zjsonpatch:zjsonpatch", version.ref = "zjs
 [bundles]
 jackson = ["jackson-datetime", "jackson-kotlin", "jackson-yaml", "jackson-xml"]
 kotlin = ["kotlin-stdLibJdk8", "kotlin-reflect"]
-mockito = ["mockito-core", "mockito-kotlin"]
+mockito = ["mockito-core", "mockito-junit-jupiter", "mockito-kotlin"]
 sshd = ["sshd-core", "sshd-scp", "sshd-sftp"]
 
 [plugins]

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/startup/AmazonQStartupActivity.kt

@@ -19,6 +19,7 @@ import software.aws.toolkits.jetbrains.core.credentials.ToolkitConnectionManager
 import software.aws.toolkits.jetbrains.core.credentials.pinning.QConnection
 import software.aws.toolkits.jetbrains.core.gettingstarted.emitUserState
 import software.aws.toolkits.jetbrains.services.amazonq.CodeWhispererFeatureConfigService
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextController
 import software.aws.toolkits.jetbrains.services.amazonq.toolwindow.AmazonQToolWindow
@@ -56,6 +57,7 @@ class AmazonQStartupActivity : ProjectActivity {
 
         QRegionProfileManager.getInstance().validateProfile(project)
 
+        AmazonQLspService.getInstance(project)
         startLsp(project)
         if (runOnce.get()) return
         emitUserState(project)

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.amazonq.profile.QRegionProfileManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererClientAdaptor
 import software.aws.toolkits.jetbrains.services.codewhisperer.customization.CodeWhispererModelConfigurator
@@ -92,6 +95,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
@@ -233,7 +239,8 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                         requestContext.fileContextInfo,
                         requestContext.awaitSupplementalContext(),
                         requestContext.customizationArn,
-                        requestContext.profileArn
+                        requestContext.profileArn,
+                        requestContext.workspaceId,
                     )
                 )
 
@@ -670,10 +677,42 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
 
         val profileArn = QRegionProfileManager.getInstance().activeProfile(project)?.arn
 
+        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, profileArn
+            project,
+            editor,
+            triggerTypeInfo,
+            caretPosition,
+            fileContext,
+            supplementalContext,
+            connection,
+            latencyContext,
+            customizationArn,
+            profileArn,
+            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 {
@@ -808,6 +847,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
             supplementalContext: SupplementalContextInfo?,
             customizationArn: String?,
             profileArn: String?,
+            workspaceId: String?,
         ): GenerateCompletionsRequest {
             val programmingLanguage = ProgrammingLanguage.builder()
                 .languageName(fileContextInfo.programmingLanguage.toCodeWhispererRuntimeLanguage().languageId)
@@ -837,6 +877,7 @@ class CodeWhispererService(private val cs: CoroutineScope) : Disposable {
                 .customizationArn(customizationArn)
                 .optOutPreference(getTelemetryOptOutPreference())
                 .profileArn(profileArn)
+                .workspaceId(workspaceId)
                 .build()
         }
     }
@@ -853,6 +894,7 @@ data class RequestContext(
     val latencyContext: LatencyContext,
     val customizationArn: String?,
     val profileArn: String?,
+    val workspaceId: String?,
 ) {
     // TODO: should make the entire getRequestContext() suspend function instead of making supplemental context only
     var supplementalContext: SupplementalContextInfo? = null

plugins/amazonq/codewhisperer/jetbrains-community/src/software/aws/toolkits/jetbrains/services/codewhisperer/settings/CodeWhispererConfigurable.kt

@@ -5,13 +5,16 @@ package software.aws.toolkits.jetbrains.services.codewhisperer.settings
 
 import com.intellij.icons.AllIcons
 import com.intellij.ide.DataManager
+import com.intellij.openapi.fileChooser.FileChooserDescriptorFactory
 import com.intellij.openapi.options.BoundConfigurable
 import com.intellij.openapi.options.Configurable
 import com.intellij.openapi.options.SearchableConfigurable
 import com.intellij.openapi.options.ex.Settings
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.emptyText
 import com.intellij.ui.components.ActionLink
 import com.intellij.ui.components.fields.ExpandableTextField
+import com.intellij.ui.dsl.builder.Align
 import com.intellij.ui.dsl.builder.bindIntText
 import com.intellij.ui.dsl.builder.bindSelected
 import com.intellij.ui.dsl.builder.bindText
@@ -24,6 +27,7 @@ import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWh
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
+import software.aws.toolkits.jetbrains.settings.LspSettings
 import software.aws.toolkits.resources.message
 import java.awt.Font
 import java.util.concurrent.TimeUnit
@@ -61,6 +65,24 @@ class CodeWhispererConfigurable(private val project: Project) :
             }
         }
 
+        group(message("amazonqFeatureDev.placeholder.lsp")) {
+            row(message("amazonqFeatureDev.placeholder.select_lsp_artifact")) {
+                val fileChooserDescriptor = FileChooserDescriptorFactory.createSingleFileDescriptor()
+                fileChooserDescriptor.isForcedToUseIdeaFileChooser = true
+
+                textFieldWithBrowseButton(fileChooserDescriptor = fileChooserDescriptor)
+                    .bindText(
+                        { LspSettings.getInstance().getArtifactPath().orEmpty() },
+                        { LspSettings.getInstance().setArtifactPath(it) }
+                    )
+                    .applyToComponent {
+                        emptyText.text = message("executableCommon.auto_managed")
+                    }
+                    .resizableColumn()
+                    .align(Align.FILL)
+            }
+        }
+
         group(message("aws.settings.codewhisperer.group.general")) {
             row {
                 checkBox(message("aws.settings.codewhisperer.include_code_with_reference")).apply {
@@ -116,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"))
+            }.visible(false)
             row {
                 checkBox(message("aws.settings.codewhisperer.project_context")).apply {
                     connect.subscribe(

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererCodeCoverageTrackerTest.kt

@@ -177,7 +177,8 @@ internal class CodeWhispererCodeCoverageTrackerTestPython : CodeWhispererCodeCov
             null,
             mock(),
             aString(),
-            aString()
+            aString(),
+            aString(),
         )
         val responseContext = ResponseContext("sessionId")
         val recommendationContext = RecommendationContext(

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)

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererServiceTest.kt

@@ -213,7 +213,8 @@ class CodeWhispererServiceTest {
                 connection = ToolkitConnectionManager.getInstance(projectRule.project).activeConnection(),
                 latencyContext = LatencyContext(),
                 customizationArn = "fake-arn",
-                profileArn = "fake-arn"
+                profileArn = "fake-arn",
+                workspaceId = null,
             )
         )
 

plugins/amazonq/codewhisperer/jetbrains-community/tst/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererSettingsTest.kt

@@ -13,17 +13,22 @@ import com.intellij.openapi.wm.impl.status.widget.StatusBarWidgetsManager
 import com.intellij.testFramework.replaceService
 import com.intellij.testFramework.runInEdtAndWait
 import com.intellij.util.xmlb.XmlSerializer
+import io.mockk.every
+import io.mockk.junit4.MockKRule
+import io.mockk.mockkObject
 import org.assertj.core.api.Assertions.assertThat
 import org.jdom.output.XMLOutputter
 import org.junit.Before
 import org.junit.Ignore
+import org.junit.Rule
 import org.junit.Test
 import org.mockito.kotlin.any
 import org.mockito.kotlin.never
 import org.mockito.kotlin.spy
 import org.mockito.kotlin.verify
 import org.mockito.kotlin.whenever
 import software.aws.toolkits.jetbrains.core.ToolWindowHeadlessManagerImpl
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.AmazonQLspService
 import software.aws.toolkits.jetbrains.services.codewhisperer.credentials.CodeWhispererLoginType
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExploreActionState
 import software.aws.toolkits.jetbrains.services.codewhisperer.explorer.isCodeWhispererEnabled
@@ -40,6 +45,9 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
     private lateinit var codewhispererServiceSpy: CodeWhispererService
     private lateinit var toolWindowHeadlessManager: ToolWindowHeadlessManagerImpl
 
+    @get:Rule
+    val mockkRule = MockKRule(this)
+
     @Before
     override fun setUp() {
         super.setUp()
@@ -246,6 +254,18 @@ class CodeWhispererSettingsTest : CodeWhispererTestBase() {
             assertThat(sut.getProjectContextIndexMaxSize()).isEqualTo(expected)
         }
     }
+
+    @Test
+    fun `toggleMetricOptIn should trigger LSP didChangeConfiguration`() {
+        mockkObject(AmazonQLspService)
+        every { AmazonQLspService.didChangeConfiguration(any()) } returns Unit
+        settingsManager.toggleMetricOptIn(true)
+        settingsManager.toggleMetricOptIn(false)
+
+        io.mockk.verify(atLeast = 2) {
+            AmazonQLspService.didChangeConfiguration(any())
+        }
+    }
 }
 
 class CodeWhispererSettingUnitTest {

plugins/amazonq/codewhisperer/jetbrains-community/tstFixtures/software/aws/toolkits/jetbrains/services/codewhisperer/CodeWhispererTestUtil.kt

@@ -258,7 +258,8 @@ fun aRequestContext(
             aString()
         ),
         customizationArn = null,
-        profileArn = null
+        profileArn = null,
+        workspaceId = null,
     )
 }
 

plugins/amazonq/shared/jetbrains-community/resources/META-INF/module-amazonq.xml

@@ -9,4 +9,5 @@
     <extensions defaultExtensionNs="com.intellij">
         <applicationService serviceImplementation="migration.software.aws.toolkits.jetbrains.services.codewhisperer.explorer.CodeWhispererExplorerActionManager"/>
     </extensions>
+
 </idea-plugin>