wip: use json-rpc for local workspace context

plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt

@@ -84,6 +84,7 @@
 import software.aws.toolkits.telemetry.CwsprChatCommandType
 import java.time.Instant
 import java.util.UUID
+import kotlin.time.Duration.Companion.milliseconds
 
 data class TestCommandMessage(
     val sender: String = "codetest",
@@ -136,15 +137,15 @@
         var prompt = message.chatMessage
         var queryResult: List<RelevantDocument> = emptyList()
         val triggerId = UUID.randomUUID().toString()
-        var shouldAddIndexInProgressMessage: Boolean = false
-        var shouldUseWorkspaceContext: Boolean = false
+        var shouldAddIndexInProgressMessage = false
+        var shouldUseWorkspaceContext = false
 
         if (prompt.contains("@workspace")) {
             if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) {
                 shouldUseWorkspaceContext = true
                 prompt = prompt.replace("@workspace", "")
                 val projectContextController = ProjectContextController.getInstance(context.project)
-                queryResult = projectContextController.queryChat(prompt, timeout = null)
+                queryResult = projectContextController.queryChat(prompt, timeout = 500)
                 if (!projectContextController.getProjectContextIndexComplete()) shouldAddIndexInProgressMessage = true
                 logger.info { "project context relevant document count: ${queryResult.size}" }
             } else {

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,38 @@
 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.InlineBm25Chunk
+import software.aws.toolkits.jetbrains.services.amazonq.project.LspRequest
+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.UpdateIndexRequest
 import software.aws.toolkits.jetbrains.services.telemetry.ClientMetadata
 import java.io.IOException
 import java.io.OutputStreamWriter
@@ -56,6 +76,7 @@
 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 +345,196 @@
         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 encryptionManager = JwtEncryptionManager()
+
+    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 LspRequest) {
+                        println(message.params)
+                        message.params = encryptionManager.encrypt(jacksonObjectMapper().writeValueAsString(message.params))
+                    }
+                    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
+            encryptionManager.writeInitializationPayload(launcherHandler.process.outputStream)
+
+            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>>
+}