Skip to content

Commit

Permalink
Cleanup callbacks implementation in CodyAgentClient
Browse files Browse the repository at this point in the history
  • Loading branch information
@pkukielka
pkukielka committed Aug 19, 2024
1 parent f9021d6 commit 4729e41
Show file tree
Hide file tree
Showing 8 changed files with 109 additions and 249 deletions.
4 changes: 1 addition & 3 deletions src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -103,8 +102,7 @@ private constructor(
fun create(project: Project): CompletableFuture<CodyAgent> {
try {
val conn = startAgentProcess()
val client = CodyAgentClient(WebUIServiceWebviewProvider(project))
client.onSetConfigFeatures = project.service<CurrentConfigFeatures>()
val client = CodyAgentClient(project, WebUIServiceWebviewProvider(project))
val launcher = startAgentLauncher(conn, client)
val server = launcher.remoteProxy
val listeningToJsonRpc = launcher.startListening()
Expand Down
245 changes: 100 additions & 145 deletions src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentClient.kt
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -22,156 +30,114 @@ 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<Unit> =
acceptOnEventThread("editTask/didUpdate", onEditTaskDidUpdate, params)

@JsonNotification("editTask/didDelete")
fun editTaskDidDelete(params: EditTask): CompletableFuture<Unit> =
acceptOnEventThread("editTask/didDelete", onEditTaskDidDelete, params)

@JsonNotification("codeLenses/display")
fun codeLensesDisplay(params: DisplayCodeLensParams): CompletableFuture<Unit> =
acceptOnEventThread("codeLenses/display", onCodeLensesDisplay, params)

var onOpenExternal: ((OpenExternalParams) -> Boolean)? = null

@JsonRequest("env/openExternal")
fun ignoreTest(params: OpenExternalParams): CompletableFuture<Boolean> =
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<Boolean> =
acceptOnEventThreadAndGet("textDocument/edit", onTextDocumentEdit, params)

@JsonRequest("textDocument/show")
fun textDocumentShow(params: TextDocumentShowParams): CompletableFuture<Boolean> =
acceptOnEventThreadAndGet("textDocument/show", onTextDocumentShow, params)

@JsonRequest("textDocument/openUntitledDocument")
fun openUntitledDocument(params: UntitledTextDocument): CompletableFuture<ProtocolTextDocument> =
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<Boolean> =
acceptOnEventThreadAndGet("workspace/edit", onWorkspaceEdit, params)
var onSetConfigFeatures: ConfigFeaturesObserver = project.service<CurrentConfigFeatures>()

/**
* 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 <T, R> acceptOnEventThreadAndGet(
name: String,
callback: ((T) -> R)?,
params: T
private fun <R> acceptOnEventThreadAndGet(
callback: (() -> R),
): CompletableFuture<R> {
val result = CompletableFuture<R>()
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)
}
}
return result
}

private fun <T, R> acceptOnEventThread(
name: String,
callback: ((T) -> R)?,
params: T
): CompletableFuture<R> {
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<Boolean> {
return acceptOnEventThreadAndGet {
BrowserOpener.openInBrowser(project, params.uri)
true
}
}

// TODO: Delete this
// Webviews
@JsonRequest("webview/create")
fun webviewCreate(params: WebviewCreateParams): CompletableFuture<Void> {
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<Boolean> {
return acceptOnEventThreadAndGet {
try {
EditService.getInstance(project).performWorkspaceEdit(params)
} catch (e: RuntimeException) {
logger.error(e)
false
}
}
}

@JsonRequest("textDocument/edit")
fun textDocument_edit(params: TextDocumentEditParams): CompletableFuture<Boolean> {
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<Boolean> {
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<ProtocolTextDocument?> {
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)
}
}

// ================================================
Expand Down Expand Up @@ -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<Void> {
logger.error("webview/create This request should not happen if you are using chat/new.")
return CompletableFuture.completedFuture(null)
}
}
Loading

0 comments on commit 4729e41

Please sign in to comment.