From 4729e4173566550324ffa8acd313819d93383c86 Mon Sep 17 00:00:00 2001 From: Piotr Kukielka Date: Mon, 19 Aug 2024 16:27:57 +0200 Subject: [PATCH] Cleanup callbacks implementation in CodyAgentClient --- .../com/sourcegraph/cody/agent/CodyAgent.kt | 4 +- .../sourcegraph/cody/agent/CodyAgentClient.kt | 245 +++++++----------- .../cody/agent/CodyAgentService.kt | 73 ------ .../cody/agent/protocol/DebugMessage.kt | 3 - .../agent/protocol/UntitledTextDocument.kt | 7 - .../cody/context/RemoteRepoSearcher.kt | 17 +- .../com/sourcegraph/cody/error/CodyConsole.kt | 2 +- .../cody/agent/CodyAgentClientTest.kt | 7 +- 8 files changed, 109 insertions(+), 249 deletions(-) delete mode 100644 src/main/kotlin/com/sourcegraph/cody/agent/protocol/DebugMessage.kt delete mode 100644 src/main/kotlin/com/sourcegraph/cody/agent/protocol/UntitledTextDocument.kt diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt index 161014355b..9a049733cd 100644 --- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt +++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt @@ -2,7 +2,6 @@ package com.sourcegraph.cody.agent import com.intellij.ide.plugins.PluginManagerCore import com.intellij.openapi.application.ApplicationInfo -import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.Logger import com.intellij.openapi.extensions.PluginId import com.intellij.openapi.project.Project @@ -103,8 +102,7 @@ private constructor( fun create(project: Project): CompletableFuture { try { val conn = startAgentProcess() - val client = CodyAgentClient(WebUIServiceWebviewProvider(project)) - client.onSetConfigFeatures = project.service() + val client = CodyAgentClient(project, WebUIServiceWebviewProvider(project)) val launcher = startAgentLauncher(conn, client) val server = launcher.remoteProxy val listeningToJsonRpc = launcher.startListening() diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt index ae55481fff..e1ee7b485a 100644 --- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt +++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt @@ -1,19 +1,27 @@ package com.sourcegraph.cody.agent import com.intellij.openapi.application.ApplicationManager +import com.intellij.openapi.application.runInEdt +import com.intellij.openapi.components.service import com.intellij.openapi.diagnostic.Logger -import com.sourcegraph.cody.agent.protocol.DebugMessage -import com.sourcegraph.cody.agent.protocol.OpenExternalParams +import com.intellij.openapi.project.Project import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument -import com.sourcegraph.cody.agent.protocol.RemoteRepoFetchState -import com.sourcegraph.cody.agent.protocol.TextDocumentShowParams -import com.sourcegraph.cody.agent.protocol.UntitledTextDocument import com.sourcegraph.cody.agent.protocol.WebviewCreateWebviewPanelParams +import com.sourcegraph.cody.agent.protocol_generated.DebugMessage import com.sourcegraph.cody.agent.protocol_generated.DisplayCodeLensParams -import com.sourcegraph.cody.agent.protocol_generated.EditTask +import com.sourcegraph.cody.agent.protocol_generated.Env_OpenExternalParams +import com.sourcegraph.cody.agent.protocol_generated.Null import com.sourcegraph.cody.agent.protocol_generated.TextDocumentEditParams +import com.sourcegraph.cody.agent.protocol_generated.TextDocument_ShowParams +import com.sourcegraph.cody.agent.protocol_generated.UntitledTextDocument import com.sourcegraph.cody.agent.protocol_generated.WorkspaceEditParams +import com.sourcegraph.cody.edit.EditService +import com.sourcegraph.cody.edit.LensesService +import com.sourcegraph.cody.error.CodyConsole +import com.sourcegraph.cody.ignore.IgnoreOracle import com.sourcegraph.cody.ui.NativeWebviewProvider +import com.sourcegraph.common.BrowserOpener +import com.sourcegraph.utils.CodyEditorUtil import java.util.concurrent.CompletableFuture import org.eclipse.lsp4j.jsonrpc.services.JsonNotification import org.eclipse.lsp4j.jsonrpc.services.JsonRequest @@ -22,124 +30,27 @@ import org.eclipse.lsp4j.jsonrpc.services.JsonRequest * Implementation of the client part of the Cody agent protocol. This class dispatches the requests * and notifications sent by the agent. */ -@Suppress("unused") -class CodyAgentClient(private val webview: NativeWebviewProvider) { +@Suppress("unused", "FunctionName") +class CodyAgentClient(private val project: Project, private val webview: NativeWebviewProvider) { companion object { private val logger = Logger.getInstance(CodyAgentClient::class.java) } - // TODO: Remove this once we stop sniffing postMessage. - // Callback that is invoked when the agent sends a "chat/updateMessageInProgress" notification. - var onNewMessage: ((WebviewPostMessageParams) -> Unit)? = null - // Callback that is invoked when the agent sends a "setConfigFeatures" message. - var onSetConfigFeatures: ConfigFeaturesObserver? = null - - // Callback that is invoked on webview messages which aren't handled by onNewMessage or - // onSetConfigFeatures - var onReceivedWebviewMessageTODODeleteThis: ((WebviewPostMessageParams) -> Unit)? = null - - // Callback for the "editTask/didUpdate" notification from the agent. - var onEditTaskDidUpdate: ((EditTask) -> Unit)? = null - - // Callback for the "editTask/didDelete" notification from the agent. - var onEditTaskDidDelete: ((EditTask) -> Unit)? = null - - // Callback for the "editTask/codeLensesDisplay" notification from the agent. - var onCodeLensesDisplay: ((DisplayCodeLensParams) -> Unit)? = null - - // Callback for the "textDocument/edit" request from the agent. - var onTextDocumentEdit: ((TextDocumentEditParams) -> Boolean)? = null - - // Callback for the "textDocument/show" request from the agent. - var onTextDocumentShow: ((TextDocumentShowParams) -> Boolean)? = null - - // Callback for the "textDocument/openUntitledDocument" request from the agent. - var onOpenUntitledDocument: ((UntitledTextDocument) -> ProtocolTextDocument)? = null - - // Callback for the "workspace/edit" request from the agent. - var onWorkspaceEdit: ((WorkspaceEditParams) -> Boolean)? = null - - var onDebugMessage: ((DebugMessage) -> Unit)? = null - - @JsonNotification("editTask/didUpdate") - fun editTaskDidUpdate(params: EditTask): CompletableFuture = - acceptOnEventThread("editTask/didUpdate", onEditTaskDidUpdate, params) - - @JsonNotification("editTask/didDelete") - fun editTaskDidDelete(params: EditTask): CompletableFuture = - acceptOnEventThread("editTask/didDelete", onEditTaskDidDelete, params) - - @JsonNotification("codeLenses/display") - fun codeLensesDisplay(params: DisplayCodeLensParams): CompletableFuture = - acceptOnEventThread("codeLenses/display", onCodeLensesDisplay, params) - - var onOpenExternal: ((OpenExternalParams) -> Boolean)? = null - - @JsonRequest("env/openExternal") - fun ignoreTest(params: OpenExternalParams): CompletableFuture = - acceptOnEventThreadAndGet("env/openExternal", onOpenExternal, params) - - var onRemoteRepoDidChange: (() -> Unit)? = null - - @JsonNotification("remoteRepo/didChange") - fun remoteRepoDidChange() { - onRemoteRepoDidChange?.invoke() - } - - var onRemoteRepoDidChangeState: ((RemoteRepoFetchState) -> Unit)? = null - - @JsonNotification("remoteRepo/didChangeState") - fun remoteRepoDidChangeState(state: RemoteRepoFetchState) { - onRemoteRepoDidChangeState?.invoke(state) - } - - var onIgnoreDidChange: (() -> Unit)? = null - - @JsonNotification("ignore/didChange") - fun ignoreDidChange() { - onIgnoreDidChange?.invoke() - } - - @JsonRequest("textDocument/edit") - fun textDocumentEdit(params: TextDocumentEditParams): CompletableFuture = - acceptOnEventThreadAndGet("textDocument/edit", onTextDocumentEdit, params) - - @JsonRequest("textDocument/show") - fun textDocumentShow(params: TextDocumentShowParams): CompletableFuture = - acceptOnEventThreadAndGet("textDocument/show", onTextDocumentShow, params) - - @JsonRequest("textDocument/openUntitledDocument") - fun openUntitledDocument(params: UntitledTextDocument): CompletableFuture = - if (onOpenUntitledDocument == null) { - CompletableFuture.failedFuture( - Exception("No callback registered for textDocument/openUntitledDocument")) - } else { - CompletableFuture.completedFuture(onOpenUntitledDocument!!.invoke(params)) - } - - @JsonRequest("workspace/edit") - fun workspaceEdit(params: WorkspaceEditParams): CompletableFuture = - acceptOnEventThreadAndGet("workspace/edit", onWorkspaceEdit, params) + var onSetConfigFeatures: ConfigFeaturesObserver = project.service() /** * Helper to run client request/notification handlers on the IntelliJ event thread. Use this * helper for handlers that require access to the IntelliJ editor, for example to read the text * contents of the open editor. */ - private fun acceptOnEventThreadAndGet( - name: String, - callback: ((T) -> R)?, - params: T + private fun acceptOnEventThreadAndGet( + callback: (() -> R), ): CompletableFuture { val result = CompletableFuture() ApplicationManager.getApplication().invokeLater { try { - if (callback != null) { - result.complete(callback.invoke(params)) - } else { - result.completeExceptionally(Exception("No callback registered for $name")) - } + result.complete(callback.invoke()) } catch (e: Exception) { result.completeExceptionally(e) } @@ -147,31 +58,86 @@ class CodyAgentClient(private val webview: NativeWebviewProvider) { return result } - private fun acceptOnEventThread( - name: String, - callback: ((T) -> R)?, - params: T - ): CompletableFuture { - val fun1: ((T) -> R)? = callback?.let { cb -> { t: T -> cb.invoke(t) } } - return acceptOnEventThreadAndGet(name, fun1, params) + // ============= + // Requests + // ============= + + @JsonRequest("env/openExternal") + fun env_openExternal(params: Env_OpenExternalParams): CompletableFuture { + return acceptOnEventThreadAndGet { + BrowserOpener.openInBrowser(project, params.uri) + true + } } - // TODO: Delete this - // Webviews - @JsonRequest("webview/create") - fun webviewCreate(params: WebviewCreateParams): CompletableFuture { - logger.error("webview/create This request should not happen if you are using chat/new.") - return CompletableFuture.completedFuture(null) + @JsonRequest("workspace/edit") + fun workspace_edit(params: WorkspaceEditParams): CompletableFuture { + return acceptOnEventThreadAndGet { + try { + EditService.getInstance(project).performWorkspaceEdit(params) + } catch (e: RuntimeException) { + logger.error(e) + false + } + } + } + + @JsonRequest("textDocument/edit") + fun textDocument_edit(params: TextDocumentEditParams): CompletableFuture { + return acceptOnEventThreadAndGet { + try { + EditService.getInstance(project).performTextEdits(params.uri, params.edits) + } catch (e: RuntimeException) { + logger.error(e) + false + } + } + } + + @JsonRequest("textDocument/show") + fun textDocument_show(params: TextDocument_ShowParams): CompletableFuture { + return acceptOnEventThreadAndGet { + val selection = params.options?.selection + val preserveFocus = params.options?.preserveFocus + val vf = CodyEditorUtil.findFileOrScratch(project, params.uri) + if (vf != null) { + CodyEditorUtil.showDocument(project, vf, selection, preserveFocus) + true + } else { + false + } + } + } + + @JsonRequest("textDocument/openUntitledDocument") + fun textDocument_openUntitledDocument( + params: UntitledTextDocument + ): CompletableFuture { + return acceptOnEventThreadAndGet { + val vf = CodyEditorUtil.createFileOrScratchFromUntitled(project, params.uri, params.content) + vf?.let { ProtocolTextDocument.fromVirtualFile(it) } + } } // ============= // Notifications // ============= + @JsonNotification("codeLenses/display") + fun codeLenses_display(params: DisplayCodeLensParams) { + runInEdt { LensesService.getInstance(project).updateLenses(params.uri, params.codeLenses) } + } + + @JsonNotification("ignore/didChange") + fun ignore_didChange(params: Null?) { + IgnoreOracle.getInstance(project).onIgnoreDidChange() + } + @JsonNotification("debug/message") - fun debugMessage(msg: DebugMessage) { - logger.warn("${msg.channel}: ${msg.message}") - onDebugMessage?.invoke(msg) + fun debug_message(params: DebugMessage) { + if (!project.isDisposed) { + CodyConsole.getInstance(project).addMessage(params) + } } // ================================================ @@ -229,28 +195,17 @@ class CodyAgentClient(private val webview: NativeWebviewProvider) { // TODO: Remove this @JsonNotification("webview/postMessage") fun webviewPostMessage(params: WebviewPostMessageParams) { - val extensionMessage = params.message - - if (onNewMessage != null && extensionMessage.type == ExtensionMessage.Type.TRANSCRIPT) { - ApplicationManager.getApplication().invokeLater { onNewMessage?.invoke(params) } - return - } - - if (onSetConfigFeatures != null && - extensionMessage.type == ExtensionMessage.Type.SET_CONFIG_FEATURES) { - ApplicationManager.getApplication().invokeLater { - onSetConfigFeatures?.update(extensionMessage.configFeatures) - } - return - } - - if (onReceivedWebviewMessageTODODeleteThis != null) { - ApplicationManager.getApplication().invokeLater { - onReceivedWebviewMessageTODODeleteThis?.invoke(params) - } - return + if (params.message.type == ExtensionMessage.Type.SET_CONFIG_FEATURES) { + runInEdt { onSetConfigFeatures.update(params.message.configFeatures) } } logger.debug("webview/postMessage ${params.id}: ${params.message}") } + + // TODO: Remove this + @JsonRequest("webview/create") + fun webviewCreate(params: WebviewCreateParams): CompletableFuture { + logger.error("webview/create This request should not happen if you are using chat/new.") + return CompletableFuture.completedFuture(null) + } } diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentService.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentService.kt index efa3bb58b9..02557e1837 100644 --- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentService.kt +++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentService.kt @@ -10,18 +10,10 @@ import com.intellij.openapi.diagnostic.Logger import com.intellij.openapi.progress.ProcessCanceledException import com.intellij.openapi.project.Project import com.intellij.util.net.HttpConfigurable -import com.sourcegraph.cody.agent.protocol.ProtocolTextDocument import com.sourcegraph.cody.config.CodyApplicationSettings -import com.sourcegraph.cody.context.RemoteRepoSearcher -import com.sourcegraph.cody.edit.EditService -import com.sourcegraph.cody.edit.LensesService -import com.sourcegraph.cody.error.CodyConsole -import com.sourcegraph.cody.ignore.IgnoreOracle import com.sourcegraph.cody.listeners.CodyFileEditorListener import com.sourcegraph.cody.statusbar.CodyStatusService -import com.sourcegraph.common.BrowserOpener import com.sourcegraph.common.CodyBundle -import com.sourcegraph.utils.CodyEditorUtil import java.util.Timer import java.util.TimerTask import java.util.concurrent.CompletableFuture @@ -56,71 +48,6 @@ class CodyAgentService(private val project: Project) : Disposable { 0, 5000) // Check every 5 seconds onStartup { agent -> - agent.client.onOpenExternal = { params -> - BrowserOpener.openInBrowser(project, params.uri) - true - } - - agent.client.onWorkspaceEdit = { params -> - try { - EditService.getInstance(project).performWorkspaceEdit(params) - } catch (e: RuntimeException) { - logger.error(e) - false - } - } - - agent.client.onCodeLensesDisplay = { params -> - LensesService.getInstance(project).updateLenses(params.uri, params.codeLenses) - } - - agent.client.onTextDocumentEdit = { params -> - try { - EditService.getInstance(project).performTextEdits(params.uri, params.edits) - } catch (e: RuntimeException) { - logger.error(e) - false - } - } - - agent.client.onTextDocumentShow = { params -> - val selection = params.options?.selection - val preserveFocus = params.options?.preserveFocus - val vf = CodyEditorUtil.findFileOrScratch(project, params.uri) - if (vf != null) { - CodyEditorUtil.showDocument(project, vf, selection, preserveFocus) - true - } else { - false - } - } - - agent.client.onOpenUntitledDocument = { params -> - val result = CompletableFuture() - ApplicationManager.getApplication().invokeAndWait { - val vf = - CodyEditorUtil.createFileOrScratchFromUntitled(project, params.uri, params.content) - result.complete(if (vf == null) null else ProtocolTextDocument.fromVirtualFile(vf)) - } - result.get() - } - - agent.client.onRemoteRepoDidChange = { - RemoteRepoSearcher.getInstance(project).remoteRepoDidChange() - } - - agent.client.onRemoteRepoDidChangeState = { state -> - RemoteRepoSearcher.getInstance(project).remoteRepoDidChangeState(state) - } - - agent.client.onIgnoreDidChange = { IgnoreOracle.getInstance(project).onIgnoreDidChange() } - - agent.client.onDebugMessage = { message -> - if (!project.isDisposed) { - CodyConsole.getInstance(project).addMessage(message) - } - } - if (!project.isDisposed) { CodyFileEditorListener.registerAllOpenedFiles(project, agent) } diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/DebugMessage.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/DebugMessage.kt deleted file mode 100644 index 1ba1b029ed..0000000000 --- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/DebugMessage.kt +++ /dev/null @@ -1,3 +0,0 @@ -package com.sourcegraph.cody.agent.protocol - -data class DebugMessage(var channel: String, var message: String) diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/UntitledTextDocument.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/UntitledTextDocument.kt deleted file mode 100644 index 8c52bd0188..0000000000 --- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/UntitledTextDocument.kt +++ /dev/null @@ -1,7 +0,0 @@ -package com.sourcegraph.cody.agent.protocol - -data class UntitledTextDocument( - val uri: String, - val content: String?, - val language: String?, -) diff --git a/src/main/kotlin/com/sourcegraph/cody/context/RemoteRepoSearcher.kt b/src/main/kotlin/com/sourcegraph/cody/context/RemoteRepoSearcher.kt index e924fbdf6e..a15186d41d 100644 --- a/src/main/kotlin/com/sourcegraph/cody/context/RemoteRepoSearcher.kt +++ b/src/main/kotlin/com/sourcegraph/cody/context/RemoteRepoSearcher.kt @@ -7,7 +7,9 @@ import com.intellij.openapi.progress.ProgressManager import com.intellij.openapi.project.Project import com.sourcegraph.cody.agent.CodyAgentException import com.sourcegraph.cody.agent.CodyAgentService -import com.sourcegraph.cody.agent.protocol.* +import com.sourcegraph.cody.agent.protocol.RemoteRepoHasParams +import com.sourcegraph.cody.agent.protocol.RemoteRepoListParams +import com.sourcegraph.cody.agent.protocol.Repo import java.util.concurrent.CompletableFuture import java.util.concurrent.TimeUnit import java.util.concurrent.TimeoutException @@ -94,17 +96,4 @@ class RemoteRepoSearcher(private val project: Project) { } return result } - - private fun fetchDone(state: RemoteRepoFetchState): Boolean { - return state.state == "complete" || state.state == "errored" - } - - // Callbacks for CodyAgentService - fun remoteRepoDidChange() { - // Ignore this. `search` uses the earliest available result. - } - - fun remoteRepoDidChangeState(state: RemoteRepoFetchState) { - // No-op. - } } diff --git a/src/main/kotlin/com/sourcegraph/cody/error/CodyConsole.kt b/src/main/kotlin/com/sourcegraph/cody/error/CodyConsole.kt index 8493e0bc75..b4cd6932c5 100644 --- a/src/main/kotlin/com/sourcegraph/cody/error/CodyConsole.kt +++ b/src/main/kotlin/com/sourcegraph/cody/error/CodyConsole.kt @@ -8,7 +8,7 @@ import com.intellij.openapi.components.service import com.intellij.openapi.project.Project import com.intellij.openapi.wm.ToolWindowManager import com.intellij.ui.content.Content -import com.sourcegraph.cody.agent.protocol.DebugMessage +import com.sourcegraph.cody.agent.protocol_generated.DebugMessage @Service(Service.Level.PROJECT) class CodyConsole(project: Project) { diff --git a/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt b/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt index deea9b819e..d7c2e614f9 100644 --- a/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt +++ b/src/test/kotlin/com/sourcegraph/cody/agent/CodyAgentClientTest.kt @@ -1,5 +1,6 @@ package com.sourcegraph.cody.agent +import com.intellij.openapi.project.Project import com.intellij.testFramework.PlatformTestUtil import com.intellij.testFramework.fixtures.BasePlatformTestCase import java.util.concurrent.locks.ReentrantLock @@ -20,8 +21,8 @@ class CodyAgentClientTest : BasePlatformTestCase() { private val lock = ReentrantLock() private val condition = lock.newCondition() - private fun client(): CodyAgentClient { - val client = CodyAgentClient(StubWebviewProvider()) + private fun client(project: Project): CodyAgentClient { + val client = CodyAgentClient(project, StubWebviewProvider()) client.onSetConfigFeatures = ConfigFeaturesObserver { lock.lock() try { @@ -37,7 +38,7 @@ class CodyAgentClientTest : BasePlatformTestCase() { @Test fun `notifies observer`() { val expected = ConfigFeatures(attribution = true, serverSentModels = true) - client() + client(myFixture.project) .webviewPostMessage( WebviewPostMessageParams( id = WEBVIEW_ID,