wip: use json-rpc for local workspace context

rli · rli · commit 6e63a5c6415c · 2025-02-28T10:44:31.000-08:00
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/lsp/AmazonQLspService.kt
@@ -3,6 +3,7 @@
 
 package software.aws.toolkits.jetbrains.services.amazonq.lsp
 
+import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper
 import com.google.gson.ToNumberPolicy
 import com.intellij.execution.configurations.GeneralCommandLine
 import com.intellij.execution.impl.ExecutionManagerImpl
@@ -34,19 +35,40 @@ import org.eclipse.lsp4j.ClientInfo
 import org.eclipse.lsp4j.FileOperationsWorkspaceCapabilities
 import org.eclipse.lsp4j.InitializeParams
 import org.eclipse.lsp4j.InitializedParams
+import org.eclipse.lsp4j.MessageActionItem
+import org.eclipse.lsp4j.MessageParams
+import org.eclipse.lsp4j.PublishDiagnosticsParams
+import org.eclipse.lsp4j.ShowMessageRequestParams
 import org.eclipse.lsp4j.SynchronizationCapabilities
 import org.eclipse.lsp4j.TextDocumentClientCapabilities
 import org.eclipse.lsp4j.WorkspaceClientCapabilities
 import org.eclipse.lsp4j.WorkspaceFolder
 import org.eclipse.lsp4j.jsonrpc.Launcher
+import org.eclipse.lsp4j.jsonrpc.MessageConsumer
+import org.eclipse.lsp4j.jsonrpc.messages.RequestMessage
+import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
+import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.launch.LSPLauncher
+import org.eclipse.lsp4j.services.LanguageClient
+import org.eclipse.lsp4j.services.LanguageServer
 import org.slf4j.event.Level
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.warn
 import software.aws.toolkits.jetbrains.isDeveloperMode
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.encryption.JwtEncryptionManager
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.model.createExtendedClientMetadata
+import software.aws.toolkits.jetbrains.services.amazonq.project.EncoderServer
+import software.aws.toolkits.jetbrains.services.amazonq.project.IndexRequest
+import software.aws.toolkits.jetbrains.services.amazonq.project.IndexUpdateMode
+import software.aws.toolkits.jetbrains.services.amazonq.project.InlineBm25Chunk
+import software.aws.toolkits.jetbrains.services.amazonq.project.LspMessage
+import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider
+import software.aws.toolkits.jetbrains.services.amazonq.project.ProjectContextProvider.Usage
+import software.aws.toolkits.jetbrains.services.amazonq.project.QueryChatRequest
+import software.aws.toolkits.jetbrains.services.amazonq.project.QueryInlineCompletionRequest
+import software.aws.toolkits.jetbrains.services.amazonq.project.RelevantDocument
+import software.aws.toolkits.jetbrains.services.amazonq.project.UpdateIndexRequest
 import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
 import java.io.IOException
 import java.io.OutputStreamWriter
@@ -56,6 +78,7 @@ import java.io.PrintWriter
 import java.io.StringWriter
 import java.net.URI
 import java.nio.charset.StandardCharsets
+import java.util.concurrent.CompletableFuture
 import java.util.concurrent.Future
 import kotlin.time.Duration.Companion.seconds
 
@@ -324,3 +347,193 @@ private class AmazonQServerInstance(private val project: Project, private val cs
         private val LOG = getLogger<AmazonQServerInstance>()
     }
 }
+
+
+class EncoderServer2(private val encoderServer: EncoderServer, private val commandLine: GeneralCommandLine, private val project: Project, private val cs: CoroutineScope) : Disposable {
+    private val launcher: Launcher<EncoderServerLspInterface>
+
+    val languageServer: EncoderServerLspInterface
+        get() = launcher.remoteProxy
+
+    @Suppress("ForbiddenVoid")
+    private val launcherFuture: Future<Void>
+    private val launcherHandler: KillableProcessHandler
+    val initializer: Job
+
+    private fun createClientCapabilities(): ClientCapabilities =
+        ClientCapabilities().apply {
+            textDocument = TextDocumentClientCapabilities().apply {
+                // For didSaveTextDocument, other textDocument/ messages always mandatory
+                synchronization = SynchronizationCapabilities().apply {
+                    didSave = true
+                }
+            }
+
+            workspace = WorkspaceClientCapabilities().apply {
+                applyEdit = false
+
+                // For workspace folder changes
+                workspaceFolders = true
+
+                // For file operations (create, delete)
+                fileOperations = FileOperationsWorkspaceCapabilities().apply {
+                    didCreate = true
+                    didDelete = true
+                }
+            }
+        }
+
+    // needs case handling when project's base path is null: default projects/unit tests
+    private fun createWorkspaceFolders(): List<WorkspaceFolder> =
+        project.basePath?.let { basePath ->
+            listOf(
+                WorkspaceFolder(
+                    URI("file://$basePath").toString(),
+                    project.name
+                )
+            )
+        }.orEmpty() // no folders to report or workspace not folder based
+
+    private fun createClientInfo(): ClientInfo {
+        val metadata = ClientMetadata.getDefault()
+        return ClientInfo().apply {
+            name = metadata.awsProduct.toString()
+            version = metadata.awsVersion
+        }
+    }
+
+    private fun createInitializeParams(): InitializeParams =
+        InitializeParams().apply {
+            processId = ProcessHandle.current().pid().toInt()
+            capabilities = createClientCapabilities()
+            clientInfo = createClientInfo()
+            workspaceFolders = createWorkspaceFolders()
+            initializationOptions = mapOf(
+                "extensionPath" to encoderServer.cachePath.toAbsolutePath().toString()
+            )
+        }
+
+    init {
+        launcherHandler = KillableColoredProcessHandler.Silent(commandLine)
+        val inputWrapper = LSPProcessListener()
+        launcherHandler.addProcessListener(inputWrapper)
+        launcherHandler.startNotify()
+
+        launcher = LSPLauncher.Builder<EncoderServerLspInterface>()
+            .setLocalService(object : LanguageClient {
+                override fun telemetryEvent(p0: Any?) {
+                    println(p0)
+                }
+
+                override fun publishDiagnostics(p0: PublishDiagnosticsParams?) {
+                    println(p0)
+                }
+
+                override fun showMessage(p0: MessageParams?) {
+                    println(p0)
+                }
+
+                override fun showMessageRequest(p0: ShowMessageRequestParams?): CompletableFuture<MessageActionItem?>? {
+                    println(p0)
+
+                    return CompletableFuture.completedFuture(null)
+                }
+
+                override fun logMessage(p0: MessageParams?) {
+                    println(p0)
+                }
+
+            })
+            .setRemoteInterface(EncoderServerLspInterface::class.java)
+            .configureGson {
+                // TODO: maybe need adapter for initialize:
+                //   https://github.com/aws/amazon-q-eclipse/blob/b9d5bdcd5c38e1dd8ad371d37ab93a16113d7d4b/plugin/src/software/aws/toolkits/eclipse/amazonq/lsp/QLspTypeAdapterFactory.java
+
+                // otherwise Gson treats all numbers as double which causes deser issues
+                it.setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE)
+            }.traceMessages(
+                PrintWriter(
+                    object : StringWriter() {
+                        private val traceLogger = LOG.atLevel(if (isDeveloperMode()) Level.INFO else Level.DEBUG)
+
+                        override fun flush() {
+                            traceLogger.log { buffer.toString() }
+                            buffer.setLength(0)
+                        }
+                    }
+                )
+            )
+            .wrapMessages { consumer ->
+                MessageConsumer { message ->
+                    if (message is RequestMessage && message.params is LspMessage) {
+                        message.params = encoderServer.encrypt(jacksonObjectMapper().writeValueAsString(message))
+                    }
+                    consumer.consume(message)
+                }
+            }
+            .setInput(inputWrapper.inputStream)
+            .setOutput(launcherHandler.process.outputStream)
+            .create()
+
+        launcherFuture = launcher.startListening()
+
+        initializer = cs.launch {
+            // encryption info must be sent within 5s or Flare process will exit
+            launcherHandler.process.outputStream.write(encoderServer.getEncryptionRequest().toByteArray())
+
+            val initializeResult = try {
+                withTimeout(5.seconds) {
+                    languageServer.initialize(createInitializeParams()).await()
+                }
+            } catch (_: TimeoutCancellationException) {
+                LOG.warn { "LSP initialization timed out" }
+                null
+            } catch (e: Exception) {
+                LOG.warn(e) { "LSP initialization failed" }
+                null
+            }
+
+            // then if this succeeds then we can allow the client to send requests
+            if (initializeResult == null) {
+                launcherHandler.destroyProcess()
+            }
+            languageServer.initialized(InitializedParams())
+        }
+    }
+
+    override fun dispose() {
+        if (!launcherFuture.isDone) {
+            try {
+                languageServer.apply {
+                    shutdown().thenRun { exit() }
+                }
+            } catch (e: Exception) {
+                LOG.warn(e) { "LSP shutdown failed" }
+                launcherHandler.destroyProcess()
+            }
+        } else if (!launcherHandler.isProcessTerminated) {
+            launcherHandler.destroyProcess()
+        }
+    }
+
+    companion object {
+        private val LOG = getLogger<AmazonQServerInstance>()
+    }
+}
+
+interface EncoderServerLspInterface : LanguageServer {
+    @JsonRequest("lsp/queryInlineProjectContext")
+    fun queryInline(request: QueryInlineCompletionRequest): CompletableFuture<List<InlineBm25Chunk>>
+
+    @JsonRequest("lsp/getUsage")
+    fun getUsageMetrics(): CompletableFuture<Usage>
+
+    @JsonRequest("lsp/query")
+    fun queryChat(request: QueryChatRequest): CompletableFuture<List<ProjectContextProvider.Chunk>>
+
+    @JsonNotification("lsp/updateIndexV2")
+    fun updateIndex(request: UpdateIndexRequest): CompletableFuture<Void>
+
+    @JsonNotification("lsp/buildIndex")
+    fun buildIndex(request: IndexRequest): CompletableFuture<List<InlineBm25Chunk>>
+}
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/EncoderServer.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/EncoderServer.kt
@@ -19,12 +19,14 @@ import com.nimbusds.jose.JWSHeader
 import com.nimbusds.jose.crypto.MACSigner
 import com.nimbusds.jwt.JWTClaimsSet
 import com.nimbusds.jwt.SignedJWT
+import kotlinx.coroutines.CoroutineScope
 import org.apache.commons.codec.digest.DigestUtils
 import software.amazon.awssdk.utils.UserHomeDirectoryUtils
 import software.aws.toolkits.core.utils.getLogger
 import software.aws.toolkits.core.utils.info
 import software.aws.toolkits.core.utils.tryDirOp
 import software.aws.toolkits.core.utils.warn
+import software.aws.toolkits.jetbrains.services.amazonq.lsp.EncoderServer2
 import software.aws.toolkits.jetbrains.services.amazonq.lsp.artifacts.extractZipFile
 import software.aws.toolkits.jetbrains.services.amazonq.project.manifest.ManifestManager
 import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings
@@ -39,19 +41,18 @@ import java.util.Base64
 import java.util.concurrent.atomic.AtomicInteger
 import javax.crypto.spec.SecretKeySpec
 
-class EncoderServer(val project: Project) : Disposable {
-    private val cachePath = Paths.get(
+class EncoderServer(val project: Project, private val cs: CoroutineScope) : Disposable {
+    val cachePath = Paths.get(
         UserHomeDirectoryUtils.userHomeDirectory()
     ).resolve(".aws").resolve("amazonq").resolve("cache")
     val manifestManager = ManifestManager()
     private val serverDirectoryName = "qserver-${manifestManager.currentVersion}.zip"
-    private val numberOfRetry = AtomicInteger(0)
     val port by lazy { NetUtils.findAvailableSocketPort() }
     private val nodeRunnableName = if (manifestManager.getOs() == "windows") "node.exe" else "node"
-    private val maxRetry: Int = 3
     private val key = generateHmacKey()
     private var processHandler: KillableProcessHandler? = null
     private val mapper = jacksonObjectMapper()
+    lateinit var encoderServer2: EncoderServer2
 
     fun downloadArtifactsAndStartServer() {
         if (ApplicationManager.getApplication().isUnitTestMode) {
@@ -92,57 +93,33 @@ class EncoderServer(val project: Project) : Disposable {
 
     fun getEncryptionRequest(): String {
         val request = EncryptionRequest(key = Base64.getUrlEncoder().withoutPadding().encodeToString(key.encoded))
-        return mapper.writeValueAsString(request)
-    }
-
-    private fun runCommand(command: GeneralCommandLine): Boolean {
-        try {
-            logger.info { "starting encoder server for project context on $port for ${project.name}" }
-            processHandler = KillableProcessHandler(command)
-            val exitCode = processHandler?.waitFor()
-            if (exitCode == true) {
-                throw Exception("Encoder server exited")
-            } else {
-                return true
-            }
-        } catch (e: Exception) {
-            logger.warn(e) { "error running encoder server:" }
-            processHandler?.destroyProcess()
-            numberOfRetry.incrementAndGet()
-            return false
-        }
+        return mapper.writeValueAsString(request) + "\n"
     }
 
     private fun getCommand(): GeneralCommandLine {
         val threadCount = CodeWhispererSettings.getInstance().getProjectContextIndexThreadCount()
         val isGpuEnabled = CodeWhispererSettings.getInstance().isProjectContextGpu()
-        val map = mutableMapOf<String, String>(
-            "PORT" to port.toString(),
-            "START_AMAZONQ_LSP" to "true",
-            "CACHE_DIR" to cachePath.toString(),
-            "MODEL_DIR" to cachePath.resolve("qserver").toString()
-        )
-        if (threadCount > 0) {
-            map["Q_WORKER_THREADS"] = threadCount.toString()
-        }
-        if (isGpuEnabled) {
-            map["Q_ENABLE_GPU"] = "true"
+        val environment = buildMap {
+            if (threadCount > 0) {
+                put("Q_WORKER_THREADS", threadCount.toString())
+            }
+
+            if (isGpuEnabled) {
+                put("Q_ENABLE_GPU", "true")
+            }
         }
-        val jsPath = cachePath.resolve("qserver").resolve("dist").resolve("extension.js").toString()
+
+        val jsPath = cachePath.resolve("qserver").resolve("lspServer.js").toString()
         val nodePath = cachePath.resolve(nodeRunnableName).toString()
         val command = GeneralCommandLine(nodePath, jsPath)
+            .withParameters("--stdio")
             .withParentEnvironmentType(GeneralCommandLine.ParentEnvironmentType.CONSOLE)
-            .withEnvironment(map)
+            .withEnvironment(environment)
         return command
     }
 
     fun start() {
-        while (numberOfRetry.get() < maxRetry) {
-            val isSuccess = runCommand(getCommand())
-            if (isSuccess) {
-                return
-            }
-        }
+        encoderServer2 = EncoderServer2(this, getCommand(), project, cs)
     }
 
     private fun close() {
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextController.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextController.kt
@@ -26,7 +26,7 @@ import java.util.concurrent.TimeoutException
 @Service(Service.Level.PROJECT)
 class ProjectContextController(private val project: Project, private val cs: CoroutineScope) : Disposable {
     // TODO: Ideally we should inject dependencies via constructor for easier testing, refer to how [TelemetryService] inject publisher and batcher
-    private val encoderServer: EncoderServer = EncoderServer(project)
+    private val encoderServer: EncoderServer = EncoderServer(project, cs)
     private val projectContextProvider: ProjectContextProvider = ProjectContextProvider(project, encoderServer, cs)
     val initJob: Job = cs.launch {
         encoderServer.downloadArtifactsAndStartServer()
diff --git a/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt b/plugins/amazonq/shared/jetbrains-community/src/software/aws/toolkits/jetbrains/services/amazonq/project/ProjectContextProvider.kt
