Skip to content

Commit

Permalink
Checkout conflicting files
Browse files Browse the repository at this point in the history
  • Loading branch information
@cbart
cbart committed Jan 31, 2024
1 parent 3a7bb87 commit 19698e1
Show file tree
Hide file tree
Showing 2 changed files with 84 additions and 71 deletions.
47 changes: 7 additions & 40 deletions src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,12 @@
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.editor.Editor;
import com.sourcegraph.cody.agent.protocol.ChatMessage;
import com.sourcegraph.cody.agent.protocol.DebugMessage;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Supplier;
import org.eclipse.lsp4j.jsonrpc.services.JsonNotification;
import org.eclipse.lsp4j.jsonrpc.services.JsonRequest;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

/** Implementation of the client part of the Cody agent protocol. */
Expand All @@ -20,9 +17,7 @@ public class CodyAgentClient {

private static final Logger logger = Logger.getInstance(CodyAgentClient.class);
// Callback that is invoked when the agent sends a "chat/updateMessageInProgress" notification.
@Nullable public Consumer<ChatMessage> onChatUpdateMessageInProgress;
@Nullable public ConfigFeaturesObserver onConfigFeatures;
@NotNull public Runnable onFinishedProcessing = () -> {};
@Nullable public Consumer<WebviewPostMessageParams> onNewMessage;
@Nullable public Editor editor;

/**
Expand All @@ -48,17 +43,6 @@ private <T> CompletableFuture<T> onEventThread(Supplier<T> handler) {
// Notifications
// =============

@JsonNotification("chat/updateMessageInProgress")
public void chatUpdateMessageInProgress(ChatMessage params) {
if (onChatUpdateMessageInProgress != null && params != null) {
ApplicationManager.getApplication()
.invokeLater(() -> onChatUpdateMessageInProgress.accept(params));
}
if (params == null) {
onFinishedProcessing.run();
}
}

@JsonNotification("debug/message")
public void debugMessage(DebugMessage msg) {
logger.warn(String.format("%s: %s", msg.getChannel(), msg.getMessage()));
Expand All @@ -73,31 +57,14 @@ public CompletableFuture<Void> webviewCreate(WebviewCreateParams params) {

@JsonNotification("webview/postMessage")
public void webviewPostMessage(WebviewPostMessageParams params) {
if (params.getMessage().getType().equals("setConfigFeatures")
&& params.getMessage().getConfigFeatures() != null) {
if (onConfigFeatures != null) {
onConfigFeatures.update(params.getMessage().getConfigFeatures());
}
}
if (onChatUpdateMessageInProgress != null
&& params.getMessage().getType().equals(ExtensionMessage.Type.TRANSCRIPT)) {
if (Boolean.FALSE.equals(params.getMessage().isMessageInProgress())) {
onFinishedProcessing.run();
} else if (params.getMessage().getMessages() != null
&& !params.getMessage().getMessages().isEmpty()) {
ApplicationManager.getApplication()
.invokeLater(
() ->
onChatUpdateMessageInProgress.accept(
Objects.requireNonNull(params.getMessage().getMessages())
.get(params.getMessage().getMessages().size() - 1)));
ExtensionMessage extensionMessage = params.getMessage();

} else {
logger.warn("webview/postMessage: no messages in transcript");
}
if (onNewMessage != null
&& extensionMessage.getType().equals(ExtensionMessage.Type.TRANSCRIPT)) {
ApplicationManager.getApplication().invokeLater(() -> onNewMessage.accept(params));
} else {
logger.warn("onChatUpdateMessageInProgress is null or message type is not transcript");
logger.warn(String.format("webview/postMessage %s: %s", params.getId(), params.getMessage()));
logger.debug("onNewMessage is null or message type is not transcript");
logger.debug(String.format("webview/postMessage %s: %s", params.getId(), extensionMessage));
}
}
}
108 changes: 77 additions & 31 deletions src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt
View file
Original file line number Diff line number Diff line change
@@ -1,19 +1,15 @@
package com.sourcegraph.cody.agent

import com.intellij.ide.plugins.PluginManagerCore
import com.intellij.openapi.components.service
import com.intellij.openapi.diagnostic.Logger
import com.intellij.openapi.extensions.PluginId
import com.intellij.openapi.project.Project
import com.intellij.openapi.util.SystemInfoRt
import com.intellij.util.system.CpuArch
import com.sourcegraph.cody.agent.CodyAgentClient
import com.sourcegraph.cody.agent.CodyAgentException
import com.sourcegraph.cody.agent.CodyAgentServer
import com.sourcegraph.cody.agent.CurrentConfigFeatures
import com.sourcegraph.cody.agent.protocol.*
import com.sourcegraph.config.ConfigUtil
import java.io.File
import java.io.IOException
import java.io.PrintWriter
import java.io.*
import java.net.Socket
import java.nio.file.*
import java.util.*
import java.util.concurrent.*
Expand All @@ -29,7 +25,7 @@ private constructor(
val client: CodyAgentClient,
val server: CodyAgentServer,
val launcher: Launcher<CodyAgentServer>,
private val agentProcess: Process,
private val connection: AgentConnection,
private val listeningToJsonRpc: Future<Void?>
) {

Expand All @@ -38,28 +34,67 @@ private constructor(
server.exit()
logger.warn("Cody Agent shut down")
listeningToJsonRpc.cancel(true)
agentProcess.destroyForcibly()
connection.close()
}
}

fun isConnected(): Boolean {
// NOTE(olafurpg): there are probably too many conditions below. We test multiple conditions
// because we don't know 100% yet what exactly constitutes a "connected" state. Out of
// abundance of caution, we check everything we can think of.
return agentProcess.isAlive && !listeningToJsonRpc.isDone && !listeningToJsonRpc.isCancelled
return connection.isConnected() && !listeningToJsonRpc.isDone && !listeningToJsonRpc.isCancelled
}

/** Abstracts over the Process and Socket types to the extent we need it. */
sealed class AgentConnection {
abstract fun isConnected(): Boolean

abstract fun close()

abstract fun getInputStream(): InputStream

abstract fun getOutputStream(): OutputStream

class ProcessConnection(val process: Process) : AgentConnection() {
override fun isConnected(): Boolean = process.isAlive

override fun close() {
process.destroy()
}

override fun getInputStream(): InputStream = process.inputStream

override fun getOutputStream(): OutputStream = process.outputStream
}

class SocketConnection(val socket: Socket) : AgentConnection() {
override fun isConnected(): Boolean = socket.isConnected && !socket.isClosed

override fun close() {
socket.close()
}

override fun getInputStream(): InputStream = socket.getInputStream()

override fun getOutputStream(): OutputStream = socket.getOutputStream()
}
}

companion object {
private val logger = Logger.getInstance(CodyAgent::class.java)
private val PLUGIN_ID = PluginId.getId("com.sourcegraph.jetbrains")
private const val DEFAULT_AGENT_DEBUG_PORT = 3113 // Also defined in agent/src/cli/jsonrpc.ts
@JvmField val executorService: ExecutorService = Executors.newCachedThreadPool()

private fun shouldConnectToDebugAgent() = System.getenv("CODY_AGENT_DEBUG_REMOTE") == "true"

private fun shouldSpawnDebuggableAgent() = System.getenv("CODY_AGENT_DEBUG_INSPECT") == "true"

fun create(project: Project): CompletableFuture<CodyAgent> {
try {
val agentProcess = startAgentProcess()
val conn = startAgentProcess()
val client = CodyAgentClient()
client.onConfigFeatures = project.service<CurrentConfigFeatures>()
val launcher = startAgentLauncher(agentProcess, client)
val launcher = startAgentLauncher(conn, client)
val server = launcher.remoteProxy
val listeningToJsonRpc = launcher.startListening()

Expand All @@ -73,7 +108,7 @@ private constructor(
.thenApply { info ->
logger.info("Connected to Cody agent " + info.name)
server.initialized()
CodyAgent(client, server, launcher, agentProcess, listeningToJsonRpc)
CodyAgent(client, server, launcher, conn, listeningToJsonRpc)
}
} catch (e: Exception) {
logger.warn("Failed to send 'initialize' JSON-RPC request Cody agent", e)
Expand All @@ -85,15 +120,22 @@ private constructor(
}
}

private fun startAgentProcess(): Process {
val binary = agentBinary()
logger.info("starting Cody agent " + binary.absolutePath)
private fun startAgentProcess(): AgentConnection {
if (shouldConnectToDebugAgent()) {
return connectToDebugAgent()
}
val command: List<String> =
if (System.getenv("CODY_DIR") != null) {
val script = File(System.getenv("CODY_DIR"), "agent/dist/index.js")
logger.info("using Cody agent script " + script.absolutePath)
listOf("node", "--enable-source-maps", script.absolutePath)
if (shouldSpawnDebuggableAgent()) {
listOf("node", "--inspect", "--enable-source-maps", script.absolutePath)
} else {
listOf("node", "--enable-source-maps", script.absolutePath)
}
} else {
val binary = agentBinary()
logger.info("starting Cody agent " + binary.absolutePath)
listOf(binary.absolutePath)
}

Expand All @@ -114,38 +156,36 @@ private constructor(
.start()

// Redirect agent stderr into idea.log by buffering line by line into `logger.warn()`
// statements. Without this logic, the stderr output of the agent process is lost if the
// process
// fails to start for some reason. We use `logger.warn()` because the agent shouldn't print
// much
// normally (excluding a few noisy messages during initialization), it's mostly used to report
// unexpected errors.
// statements. Without this logic, the stderr output of the agent process is lost if
// the process fails to start for some reason. We use `logger.warn()` because the
// agent shouldn't print much normally (excluding a few noisy messages during
// initialization), it's mostly used to report unexpected errors.
Thread { process.errorStream.bufferedReader().forEachLine { line -> logger.warn(line) } }
.start()

return process
return AgentConnection.ProcessConnection(process)
}

@Throws(IOException::class, CodyAgentException::class)
private fun startAgentLauncher(
agentProcess: Process,
process: AgentConnection,
client: CodyAgentClient
): Launcher<CodyAgentServer> {
return Launcher.Builder<CodyAgentServer>()
.configureGson { gsonBuilder ->
gsonBuilder
// emit `null` instead of leaving fields undefined because Cody
// in VSC has
// many `=== null` checks that return false for undefined fields.
// VSC has many `=== null` checks that return false for undefined fields.
.serializeNulls()
.registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer)
.registerTypeAdapter(ContextFile::class.java, contextFileDeserializer)
.registerTypeAdapter(Speaker::class.java, SpeakerSerializer)
}
.setRemoteInterface(CodyAgentServer::class.java)
.traceMessages(traceWriter())
.setExecutorService(executorService)
.setInput(agentProcess.inputStream)
.setOutput(agentProcess.outputStream)
.setInput(process.getInputStream())
.setOutput(process.getOutputStream())
.setLocalService(client)
.create()
}
Expand All @@ -163,6 +203,7 @@ private constructor(
}

private fun agentDirectory(): Path? {
// N.B. this is the default/production setting. CODY_DIR overrides it locally.
val fromProperty = System.getProperty("cody-agent.directory", "")
if (fromProperty.isNotEmpty()) {
return Paths.get(fromProperty)
Expand Down Expand Up @@ -212,5 +253,10 @@ private constructor(
}
return null
}

private fun connectToDebugAgent(): AgentConnection {
val port = System.getenv("CODY_AGENT_DEBUG_PORT")?.toInt() ?: DEFAULT_AGENT_DEBUG_PORT
return AgentConnection.SocketConnection(Socket("localhost", port))
}
}
}

0 comments on commit 19698e1

Please sign in to comment.