diff --git a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt index 25630d63bab..fb583742dad 100644 --- a/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt +++ b/plugins/amazonq/chat/jetbrains-community/src/software/aws/toolkits/jetbrains/services/cwc/controller/ChatController.kt @@ -84,6 +84,7 @@ import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings 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 @@ class ChatController private constructor( var prompt = message.chatMessage var queryResult: List = 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 { 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 index 49141741de0..91c9fd8a8b9 100644 --- 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,12 +35,22 @@ 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 @@ -47,6 +58,15 @@ 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.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 +345,196 @@ private class AmazonQServerInstance(private val project: Project, private val cs private val LOG = getLogger() } } + + +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 + + val languageServer: EncoderServerLspInterface + get() = launcher.remoteProxy + + @Suppress("ForbiddenVoid") + private val launcherFuture: Future + 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 = + 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() + .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? { + 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() + } +} + +interface EncoderServerLspInterface : LanguageServer { + @JsonRequest("lsp/queryInlineProjectContext") + fun queryInline(request: QueryInlineCompletionRequest): CompletableFuture> + + @JsonRequest("lsp/getUsage") + fun getUsageMetrics(): CompletableFuture + + @JsonRequest("lsp/query") + fun queryChat(request: QueryChatRequest): CompletableFuture> + + @JsonNotification("lsp/updateIndexV2") + fun updateIndex(request: UpdateIndexRequest): CompletableFuture + + @JsonNotification("lsp/buildIndex") + fun buildIndex(request: IndexRequest): CompletableFuture> +} 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 index 2da14c2dd50..928fb6de439 100644 --- 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 @@ -3,28 +3,21 @@ package software.aws.toolkits.jetbrains.services.amazonq.project -import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper import com.intellij.execution.configurations.GeneralCommandLine -import com.intellij.execution.process.KillableProcessHandler import com.intellij.openapi.Disposable import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.project.Project -import com.intellij.openapi.util.process.ProcessCloseUtil +import com.intellij.openapi.util.Disposer import com.intellij.util.io.HttpRequests import com.intellij.util.io.createDirectories import com.intellij.util.net.NetUtils -import com.nimbusds.jose.JOSEObjectType -import com.nimbusds.jose.JWSAlgorithm -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 @@ -33,25 +26,16 @@ import java.nio.file.Files import java.nio.file.Path import java.nio.file.Paths import java.nio.file.attribute.PosixFilePermission -import java.security.Key -import java.security.SecureRandom -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) { @@ -61,92 +45,36 @@ class EncoderServer(val project: Project) : Disposable { start() } - fun isNodeProcessRunning() = processHandler != null && processHandler?.process?.isAlive == true - - private fun generateHmacKey(): Key { - val keyBytes = ByteArray(32) - SecureRandom().nextBytes(keyBytes) - return SecretKeySpec(keyBytes, "HmacSHA256") - } - - fun encrypt(data: String): String { - val header = JWSHeader.Builder(JWSAlgorithm.HS256) - .type(JOSEObjectType.JWT) - .build() - - val claimsSet = JWTClaimsSet.Builder() - .subject(Base64.getUrlEncoder().withoutPadding().encodeToString(data.toByteArray())) - .build() - - val signedJWT = SignedJWT(header, claimsSet) - signedJWT.sign(MACSigner(key.encoded)) - - return signedJWT.serialize() - } - - data class EncryptionRequest( - val version: String = "1.0", - val mode: String = "JWT", - val key: String, - ) - - 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 - } - } + fun isNodeProcessRunning() = encoderServer2.initializer.isCompleted private fun getCommand(): GeneralCommandLine { val threadCount = CodeWhispererSettings.getInstance().getProjectContextIndexThreadCount() val isGpuEnabled = CodeWhispererSettings.getInstance().isProjectContextGpu() - val map = mutableMapOf( - "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() { - processHandler?.process?.let { ProcessCloseUtil.close(it) } + Disposer.dispose(encoderServer2) } private fun downloadArtifactsIfNeeded() { 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 index dabe282f2fd..707b227f946 100644 --- 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 index 9190e9d1313..37204709fff 100644 --- 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 @@ -6,7 +6,6 @@ package software.aws.toolkits.jetbrains.services.amazonq.project import com.fasterxml.jackson.annotation.JsonIgnoreProperties import com.fasterxml.jackson.annotation.JsonProperty import com.fasterxml.jackson.module.kotlin.jacksonObjectMapper -import com.fasterxml.jackson.module.kotlin.readValue import com.intellij.openapi.Disposable import com.intellij.openapi.application.ApplicationManager import com.intellij.openapi.project.BaseProjectDirectories.Companion.getBaseDirectories @@ -16,8 +15,8 @@ import com.intellij.openapi.vfs.VirtualFile import com.intellij.openapi.vfs.VirtualFileVisitor import com.intellij.openapi.vfs.isFile import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.async import kotlinx.coroutines.delay +import kotlinx.coroutines.future.await import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import kotlinx.coroutines.withTimeout @@ -31,12 +30,9 @@ import software.aws.toolkits.jetbrains.services.amazonq.SUPPLEMENTAL_CONTEXT_TIM import software.aws.toolkits.jetbrains.services.cwc.controller.chat.telemetry.getStartUrl import software.aws.toolkits.jetbrains.settings.CodeWhispererSettings import software.aws.toolkits.telemetry.AmazonqTelemetry -import java.io.OutputStreamWriter -import java.net.HttpURLConnection -import java.net.URL import java.util.concurrent.atomic.AtomicBoolean import java.util.concurrent.atomic.AtomicInteger -import kotlin.time.Duration.Companion.minutes +import kotlin.time.measureTimedValue class ProjectContextProvider(val project: Project, private val encoderServer: EncoderServer, private val cs: CoroutineScope) : Disposable { private val retryCount = AtomicInteger(0) @@ -106,8 +102,8 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En cs.launch { while (retryCount.get() < 5) { try { - logger.info { "project context: about to init key" } - val isInitSuccess = initEncryption() + logger.info { "project context: waiting for server to start" } + val isInitSuccess = encoderServer.encoderServer2.initializer.isCompleted if (isInitSuccess) { logger.info { "project context index starting" } delay(300) @@ -127,12 +123,6 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En } } - private fun initEncryption(): Boolean { - val request = encoderServer.getEncryptionRequest() - val response = sendMsgToLsp(LspMessage.Initialize, request) - return response?.responseCode == 200 - } - fun index(): Boolean { val projectRoot = project.basePath ?: return false @@ -146,48 +136,48 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En logger.debug { "time elapsed to collect project context files: ${duration}ms, collected ${filesResult.files.size} files" } val indexOption = if (CodeWhispererSettings.getInstance().isProjectContextEnabled()) IndexOption.ALL else IndexOption.DEFAULT - val encrypted = encryptRequest(IndexRequest(filesResult.files, projectRoot, indexOption.command, "")) - val response = sendMsgToLsp(LspMessage.Index, encrypted) - - duration = (System.currentTimeMillis() - indexStartTime).toDouble() - logger.debug { "project context index time: ${duration}ms" } - val startUrl = getStartUrl(project) - if (response?.responseCode == 200) { - val usage = getUsage() - recordIndexWorkspace(duration, filesResult.files.size, filesResult.fileSize, true, usage?.memoryUsage, usage?.cpuUsage, startUrl) - logger.debug { "project context index finished for ${project.name}" } - return true - } else { + try { + val (_, indexTime) = measureTimedValue { + encoderServer.encoderServer2.languageServer.buildIndex(IndexRequest(filesResult.files, projectRoot, indexOption.command, "")) + } + + logger.debug { "project context index time: ${indexTime}ms" } + } catch (e: Exception) { logger.debug { "project context index failed" } recordIndexWorkspace(duration, filesResult.files.size, filesResult.fileSize, false, null, null, startUrl) return false } + + val usage = getUsage() + recordIndexWorkspace(duration, filesResult.files.size, filesResult.fileSize, true, usage?.memoryUsage, usage?.cpuUsage, startUrl) + logger.debug { "project context index finished for ${project.name}" } + return true } // TODO: rename queryChat suspend fun query(prompt: String, timeout: Long?): List = withTimeout(timeout ?: CHAT_EXPLICIT_PROJECT_CONTEXT_TIMEOUT) { - cs.async { - val encrypted = encryptRequest(QueryChatRequest(prompt)) - val response = sendMsgToLsp(LspMessage.QueryChat, encrypted) ?: return@async emptyList() - val parsedResponse = mapper.readValue>(response.responseBody) - queryResultToRelevantDocuments(parsedResponse) - }.await() + val response = try { + encoderServer.encoderServer2.languageServer.queryChat(QueryChatRequest(prompt)).await() + } catch (e: Exception) { + logger.warn { "error querying chat ${e.message}" } + return@withTimeout emptyList() + } + queryResultToRelevantDocuments(response) } suspend fun queryInline(query: String, filePath: String, target: InlineContextTarget): List = withTimeout(SUPPLEMENTAL_CONTEXT_TIMEOUT) { - cs.async { - val encrypted = encryptRequest(QueryInlineCompletionRequest(query, filePath, target.toString())) - val r = sendMsgToLsp(LspMessage.QueryInlineCompletion, encrypted) ?: return@async emptyList() - return@async mapper.readValue>(r.responseBody) - }.await() + try { + encoderServer.encoderServer2.languageServer.queryInline(QueryInlineCompletionRequest(query, filePath, target.toString())).await() + } catch (e: Exception) { + logger.warn { "error querying chat ${e.message}" } + return@withTimeout emptyList() + } } - fun getUsage(): Usage? { - val response = sendMsgToLsp(LspMessage.GetUsageMetrics, request = null) ?: return null - return try { - val parsedResponse = mapper.readValue(response.responseBody) - parsedResponse + fun getUsage(): Usage? = runBlocking { + try { + encoderServer.encoderServer2.languageServer.getUsageMetrics().await() } catch (e: Exception) { logger.warn { "error parsing query response ${e.message}" } null @@ -195,8 +185,7 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En } fun updateIndex(filePaths: List, mode: IndexUpdateMode) { - val encrypted = encryptRequest(UpdateIndexRequest(filePaths, mode.command)) - sendMsgToLsp(LspMessage.UpdateIndex, encrypted) + encoderServer.encoderServer2.languageServer.updateIndex(UpdateIndexRequest(filePaths, mode.command)) } private fun recordIndexWorkspace( @@ -220,26 +209,6 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En ) } - private fun setConnectionTimeout(connection: HttpURLConnection, timeout: Int) { - connection.connectTimeout = timeout - connection.readTimeout = timeout - } - - private fun setConnectionProperties(connection: HttpURLConnection) { - connection.requestMethod = "POST" - connection.setRequestProperty("Content-Type", "text/plain") - connection.setRequestProperty("Accept", "text/plain") - } - - private fun setConnectionRequest(connection: HttpURLConnection, payload: String) { - connection.doOutput = true - connection.outputStream.use { outputStream -> - OutputStreamWriter(outputStream).use { writer -> - writer.write(payload) - } - } - } - private fun willExceedPayloadLimit(currentTotalFileSize: Long, currentFileSize: Long): Boolean { val maxSize = CodeWhispererSettings.getInstance().getProjectContextIndexMaxSize() return currentTotalFileSize.let { totalSize -> totalSize > (maxSize * 1024 * 1024 - currentFileSize) } @@ -305,39 +274,6 @@ class ProjectContextProvider(val project: Project, private val encoderServer: En return documents } - private fun encryptRequest(r: LspRequest): String { - val payloadJson = mapper.writeValueAsString(r) - return encoderServer.encrypt(payloadJson) - } - - private fun sendMsgToLsp(msgType: LspMessage, request: String?): LspResponse? { - logger.info { "sending message: ${msgType.endpoint} to lsp on port ${encoderServer.port}" } - val url = URL("http://localhost:${encoderServer.port}/${msgType.endpoint}") - if (!encoderServer.isNodeProcessRunning()) { - logger.warn { "language server for ${project.name} is not running" } - return null - } - // use 1h as timeout for index, 5 seconds for other APIs - val timeoutMs = if (msgType is LspMessage.Index) 60.minutes.inWholeMilliseconds.toInt() else 5000 - return with(url.openConnection() as HttpURLConnection) { - setConnectionProperties(this) - setConnectionTimeout(this, timeoutMs) - request?.let { r -> - setConnectionRequest(this, r) - } - val responseCode = this.responseCode - logger.info { "receiving response for $msgType with responseCode $responseCode" } - - val responseBody = if (responseCode == 200) { - this.inputStream.bufferedReader().use { reader -> reader.readText() } - } else { - "" - } - - LspResponse(responseCode, responseBody) - } - } - override fun dispose() { retryCount.set(0) }