From 68ae69e72087350d96d8af22914250e1326c8026 Mon Sep 17 00:00:00 2001 From: Beyang Liu Date: Thu, 1 Feb 2024 13:54:19 -0800 Subject: [PATCH] basic context file menu pass context files on chat submit proper URI serialization --- .../java/com/sourcegraph/cody/PromptPanel.kt | 294 +++++++++++++----- .../cody/agent/CodyAgentClient.java | 17 +- .../cody/agent/WebviewPostMessageParams.kt | 9 +- .../sourcegraph/cody/CodyToolWindowContent.kt | 2 + .../com/sourcegraph/cody/agent/CodyAgent.kt | 6 +- .../cody/agent/protocol/ClientInfo.kt | 3 +- .../cody/agent/protocol/ContextFile.kt | 74 ++++- .../sourcegraph/cody/chat/AgentChatSession.kt | 37 ++- .../com/sourcegraph/cody/chat/ChatSession.kt | 7 +- .../com/sourcegraph/cody/chat/ui/ChatPanel.kt | 2 +- .../cody/ui/AutoGrowingTextArea.kt | 13 + .../com/sourcegraph/cody/PromptPanelTest.kt | 42 +++ 12 files changed, 401 insertions(+), 105 deletions(-) create mode 100644 src/test/kotlin/com/sourcegraph/cody/PromptPanelTest.kt diff --git a/src/main/java/com/sourcegraph/cody/PromptPanel.kt b/src/main/java/com/sourcegraph/cody/PromptPanel.kt index 09ab5d8180..4384fa01b6 100644 --- a/src/main/java/com/sourcegraph/cody/PromptPanel.kt +++ b/src/main/java/com/sourcegraph/cody/PromptPanel.kt @@ -1,6 +1,5 @@ package com.sourcegraph.cody -import com.intellij.openapi.actionSystem.AnAction import com.intellij.openapi.actionSystem.AnActionEvent import com.intellij.openapi.actionSystem.CustomShortcutSet import com.intellij.openapi.actionSystem.KeyboardShortcut @@ -10,6 +9,8 @@ import com.intellij.ui.DocumentAdapter import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.util.ui.JBUI import com.intellij.util.ui.UIUtil +import com.sourcegraph.cody.agent.WebviewMessage +import com.sourcegraph.cody.agent.protocol.ContextFile import com.sourcegraph.cody.chat.ChatSession import com.sourcegraph.cody.chat.CodyChatMessageHistory import com.sourcegraph.cody.chat.ui.SendButton @@ -18,50 +19,60 @@ import com.sourcegraph.cody.vscode.CancellationToken import java.awt.Dimension import java.awt.event.ComponentAdapter import java.awt.event.ComponentEvent -import java.awt.event.KeyAdapter import java.awt.event.KeyEvent +import java.awt.event.MouseAdapter +import java.awt.event.MouseEvent +import java.io.File +import javax.swing.DefaultListModel import javax.swing.JLayeredPane +import javax.swing.JList +import javax.swing.JScrollPane import javax.swing.KeyStroke import javax.swing.border.EmptyBorder import javax.swing.event.AncestorEvent import javax.swing.event.AncestorListener import javax.swing.event.DocumentEvent -class PromptPanel(private val chatSession: ChatSession) : JLayeredPane() { +class PromptPanel( + private val chatSession: ChatSession, +) : JLayeredPane() { + /** View components */ private val autoGrowingTextArea = AutoGrowingTextArea(5, 9, this) private val scrollPane = autoGrowingTextArea.scrollPane private val textArea = autoGrowingTextArea.textArea + private val sendButton = SendButton() + private var contextFilesSelectorModel = DefaultListModel() + private val contextFilesSelector = JList(contextFilesSelectorModel) + private val contextFilesScroller = JScrollPane(contextFilesSelector) + + /** Externally updated state */ + private val selectedContextFiles: ArrayList = ArrayList() + + /** Related components */ private val promptMessageHistory = CodyChatMessageHistory(CHAT_MESSAGE_HISTORY_CAPACITY, chatSession) - private val sendButton = SendButton() init { + /** Initialize view */ textArea.emptyText.text = "Ask a question about this code..." + scrollPane.border = EmptyBorder(JBUI.emptyInsets()) + scrollPane.background = UIUtil.getPanelBackground() - sendButton.addActionListener { _ -> chatSession.sendMessage(getTextAndReset()) } + // Set initial bounds for the scrollPane (100x100) to ensure proper initialization; + // later adjusted dynamically based on component resizing in the component listener. + scrollPane.setBounds(0, 0, 100, 100) + add(scrollPane, DEFAULT_LAYER) + scrollPane.setBounds(0, 0, width, scrollPane.preferredSize.height) - val upperMessageAction: AnAction = - object : DumbAwareAction() { - override fun actionPerformed(e: AnActionEvent) { - promptMessageHistory.popUpperMessage(textArea) - } - } - val lowerMessageAction: AnAction = - object : DumbAwareAction() { - override fun actionPerformed(e: AnActionEvent) { - promptMessageHistory.popLowerMessage(textArea) - } - } - val sendMessageAction: AnAction = - object : DumbAwareAction() { - override fun actionPerformed(e: AnActionEvent) { - if (sendButton.isEnabled) { - chatSession.sendMessage(getTextAndReset()) - } - } - } + contextFilesSelector.border = EmptyBorder(JBUI.emptyInsets()) + add(contextFilesScroller, PALETTE_LAYER, 0) + + add(sendButton, PALETTE_LAYER, 0) + preferredSize = Dimension(scrollPane.width, scrollPane.height) + + /** Add listeners */ addAncestorListener( object : AncestorListener { override fun ancestorAdded(event: AncestorEvent?) { @@ -73,87 +84,222 @@ class PromptPanel(private val chatSession: ChatSession) : JLayeredPane() { override fun ancestorMoved(event: AncestorEvent?) {} }) - - sendMessageAction.registerCustomShortcutSet(DEFAULT_SUBMIT_ACTION_SHORTCUT, textArea) - upperMessageAction.registerCustomShortcutSet(POP_UPPER_MESSAGE_ACTION_SHORTCUT, textArea) - lowerMessageAction.registerCustomShortcutSet(POP_LOWER_MESSAGE_ACTION_SHORTCUT, textArea) - - textArea.addKeyListener( - object : KeyAdapter() { - override fun keyReleased(e: KeyEvent) { - val keyCode = e.keyCode - if (keyCode != KeyEvent.VK_UP && keyCode != KeyEvent.VK_DOWN) {} + addComponentListener( + object : ComponentAdapter() { + override fun componentResized(e: ComponentEvent?) { + // HACK + val jButtonPreferredSize = sendButton.preferredSize + sendButton.setBounds( + scrollPane.width - jButtonPreferredSize.width, + scrollPane.height - jButtonPreferredSize.height, + jButtonPreferredSize.width, + jButtonPreferredSize.height) + refreshViewLayout() } }) + + // Add user action listeners + sendButton.addActionListener { _ -> didSubmitChatMessage() } textArea.document.addDocumentListener( object : DocumentAdapter() { override fun textChanged(e: DocumentEvent) { refreshSendButton() + didUserInputChange(textArea.text) } }) - scrollPane.border = EmptyBorder(JBUI.emptyInsets()) - scrollPane.background = UIUtil.getPanelBackground() + contextFilesSelector.addMouseListener( + object : MouseAdapter() { + override fun mouseClicked(e: MouseEvent) { + contextFilesSelector.selectedIndex = contextFilesSelector.locationToIndex(e.getPoint()) + didSelectContextFile() + textArea.requestFocusInWindow() + } + }) + for (shortcut in listOf(ENTER, UP, DOWN, TAB)) { // key listeners + object : DumbAwareAction() { + override fun actionPerformed(e: AnActionEvent) { + didUseShortcut(shortcut) + } + } + .registerCustomShortcutSet(shortcut, textArea) + } + } - // Set initial bounds for the scrollPane (100x100) to ensure proper initialization; - // later adjusted dynamically based on component resizing in the component listener. - scrollPane.setBounds(0, 0, 100, 100) + private fun didUseShortcut(shortcut: CustomShortcutSet) { + if (contextFilesSelector.model.size > 0) { + when (shortcut) { + UP -> setSelectedContextFileIndex(-1) + DOWN -> setSelectedContextFileIndex(1) + ENTER, + TAB -> didSelectContextFile() + } + return + } + when (shortcut) { + ENTER -> if (sendButton.isEnabled) didSubmitChatMessage() + UP -> promptMessageHistory.popUpperMessage(textArea) + DOWN -> promptMessageHistory.popLowerMessage(textArea) + } + } - add(scrollPane, DEFAULT_LAYER) + /** View handlers */ + private fun didSubmitChatMessage() { + val cf = findContextFiles(selectedContextFiles, textArea.text) + val text = textArea.text - add(sendButton, PALETTE_LAYER, 0) + // Reset text + promptMessageHistory.messageSent(text) + textArea.text = "" + selectedContextFiles.clear() - scrollPane.setBounds(0, 0, width, scrollPane.preferredSize.height) + chatSession.sendMessage(text, cf) + } - preferredSize = Dimension(scrollPane.width, scrollPane.height) + private fun didSelectContextFile() { + if (contextFilesSelector.selectedIndex == -1) return - addComponentListener( - object : ComponentAdapter() { - override fun componentResized(e: ComponentEvent?) { - revalidate() - val jButtonPreferredSize = sendButton.preferredSize - sendButton.setBounds( - scrollPane.width - jButtonPreferredSize.width, - scrollPane.height - jButtonPreferredSize.height, - jButtonPreferredSize.width, - jButtonPreferredSize.height) - } - }) + val selected = contextFilesSelector.model.getElementAt(contextFilesSelector.selectedIndex) + this.selectedContextFiles.add(selected.contextFile) + val cfDisplayPath = selected.toString() + val expr = findAtExpressions(textArea.text).lastOrNull() ?: return + + textArea.replaceRange("@${cfDisplayPath} ", expr.startIndex, expr.endIndex) + + setContextFilesSelector(listOf()) + refreshViewLayout() } + private fun didUserInputChange(text: String) { + val exp = findAtExpressions(text).lastOrNull() + if (exp == null || + exp.endIndex < + text.length) { // TODO(beyang): instead of text.length, should be current cursor index + setContextFilesSelector(listOf()) + refreshViewLayout() + return + } + this.chatSession.sendWebviewMessage( + WebviewMessage(command = "getUserContext", submitType = "user", query = exp.value)) + } + + /** State updaters */ + private fun setSelectedContextFileIndex(increment: Int) { + var newSelectedIndex = + (contextFilesSelector.selectedIndex + increment) % contextFilesSelector.model.size + if (newSelectedIndex < 0) { + newSelectedIndex += contextFilesSelector.model.size + } + contextFilesSelector.selectedIndex = newSelectedIndex + refreshViewLayout() + } + + /** View updaters */ @RequiresEdt - fun refreshSendButton() { + private fun refreshViewLayout() { + // get the height of the context files list based on font height and number of context files + val contextFilesHeight = contextFilesSelector.preferredSize.height + contextFilesScroller.size = Dimension(scrollPane.width, contextFilesHeight) + + val margin = 10 + scrollPane.setBounds(0, contextFilesHeight, width, scrollPane.preferredSize.height + margin) + preferredSize = Dimension(scrollPane.width, scrollPane.height + contextFilesHeight) + + sendButton.setLocation( + scrollPane.width - sendButton.preferredSize.width, + scrollPane.height + contextFilesSelector.height - sendButton.preferredSize.height) + + revalidate() + } + + @RequiresEdt + private fun refreshSendButton() { sendButton.isEnabled = textArea.getText().isNotEmpty() && chatSession.getCancellationToken().isDone } + /** External prop setters */ fun registerCancellationToken(cancellationToken: CancellationToken) { cancellationToken.onFinished { ApplicationManager.getApplication().invokeLater { refreshSendButton() } } } - override fun revalidate() { - super.revalidate() - - scrollPane.setBounds(0, 0, width, scrollPane.preferredSize.height) - preferredSize = Dimension(scrollPane.width, scrollPane.height) - } + @RequiresEdt + fun setContextFilesSelector(newUserContextFiles: List) { + val changed = contextFilesSelectorModel.elements().toList() != newUserContextFiles + if (changed) { + val newModel = DefaultListModel() + newModel.addAll(newUserContextFiles.map { f -> DisplayedContextFile(f) }) + contextFilesSelector.model = newModel + contextFilesSelectorModel = newModel - private fun getTextAndReset(): String { - val text = textArea.text - promptMessageHistory.messageSent(text) - textArea.text = "" - return text + if (newUserContextFiles.isNotEmpty()) { + contextFilesSelector.selectedIndex = 0 + } else { + contextFilesSelector.selectedIndex = -1 + } + refreshViewLayout() + } } companion object { private const val CHAT_MESSAGE_HISTORY_CAPACITY = 100 - private val JUST_ENTER = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_ENTER, 0), null) + private val KEY_ENTER = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_ENTER, 0), null) + private val KEY_UP = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_UP, 0), null) + private val KEY_DOWN = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_DOWN, 0), null) + private val KEY_TAB = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_TAB, 0), null) + + val ENTER = CustomShortcutSet(KEY_ENTER) + val UP = CustomShortcutSet(KEY_UP) + val DOWN = CustomShortcutSet(KEY_DOWN) + val TAB = CustomShortcutSet(KEY_TAB) + } +} + +data class DisplayedContextFile(val contextFile: ContextFile) { + override fun toString(): String { + return displayPath(contextFile) + } +} + +data class AtExpression( + val startIndex: Int, + val endIndex: Int, + val rawValue: String, + val value: String +) + +val atExpressionPattern = """(@(?:\\\s|[^\s])+)(?:\s|$)""".toRegex() + +fun findAtExpressions(text: String): List { + val matches = atExpressionPattern.findAll(text) + val expressions = ArrayList() + for (match in matches) { + val subMatch = match.groups.get(1) + if (subMatch != null) { + val value = subMatch.value.substring(1).replace("\\ ", " ") + expressions.add( + AtExpression(subMatch.range.first, subMatch.range.last + 1, subMatch.value, value)) + } + } + return expressions +} + +fun findContextFiles(contextFiles: List, text: String): List { + val atExpressions = findAtExpressions(text) + return contextFiles.filter { f -> atExpressions.any { it.value == displayPath(f) } } +} + +// TODO(beyang): temporary displayPath implementation, should be updated to mirror what the VS Code +// plugin does +fun displayPath(contextFile: ContextFile): String { + // if the path contains more than three components, display the last three + val path = contextFile.uri.path - val UP = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_UP, 0), null) - val DOWN = KeyboardShortcut(KeyStroke.getKeyStroke(KeyEvent.VK_DOWN, 0), null) - val DEFAULT_SUBMIT_ACTION_SHORTCUT = CustomShortcutSet(JUST_ENTER) - val POP_UPPER_MESSAGE_ACTION_SHORTCUT = CustomShortcutSet(UP) - val POP_LOWER_MESSAGE_ACTION_SHORTCUT = CustomShortcutSet(DOWN) + // split path on separator (OS agnostic) + val pathComponents = path.split(File.separator) + if (pathComponents.size > 3) { + return "...${File.separator}${pathComponents.subList(pathComponents.size - 3, pathComponents.size).joinToString(File.separator)}" } + return path } diff --git a/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java b/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java index cdf065008c..b0a9fe38fe 100644 --- a/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java +++ b/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java @@ -4,6 +4,8 @@ import com.intellij.openapi.diagnostic.Logger; import com.intellij.openapi.editor.Editor; import com.sourcegraph.cody.agent.protocol.DebugMessage; +import java.util.HashMap; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.function.Consumer; import java.util.function.Supplier; @@ -63,8 +65,19 @@ public void webviewPostMessage(WebviewPostMessageParams params) { && extensionMessage.getType().equals(ExtensionMessage.Type.TRANSCRIPT)) { ApplicationManager.getApplication().invokeLater(() -> onNewMessage.accept(params)); } else { - logger.debug("onNewMessage is null or message type is not transcript"); - logger.debug(String.format("webview/postMessage %s: %s", params.getId(), extensionMessage)); + var listener = this.webviewMessageListeners.get(params.getId()); + if (listener != null) { + listener.accept(params.getMessage()); + } else { + logger.warn( + String.format("webview/postMessage %s: %s", params.getId(), params.getMessage())); + } } } + + private final Map> webviewMessageListeners = new HashMap<>(); + + public void onWebviewMessage(String panelID, Consumer callback) { + this.webviewMessageListeners.put(panelID, callback); + } } diff --git a/src/main/java/com/sourcegraph/cody/agent/WebviewPostMessageParams.kt b/src/main/java/com/sourcegraph/cody/agent/WebviewPostMessageParams.kt index 30b19fa67e..18d962c4a2 100644 --- a/src/main/java/com/sourcegraph/cody/agent/WebviewPostMessageParams.kt +++ b/src/main/java/com/sourcegraph/cody/agent/WebviewPostMessageParams.kt @@ -3,6 +3,8 @@ package com.sourcegraph.cody.agent import com.sourcegraph.cody.agent.protocol.ChatError import com.sourcegraph.cody.agent.protocol.ChatMessage import com.sourcegraph.cody.agent.protocol.ContextFile +import kotlinx.serialization.* +import kotlinx.serialization.modules.* /** * A message sent from the webview to the extension host. See vscode/src/chat/protocol.ts for the @@ -15,6 +17,7 @@ data class WebviewMessage( val addEnhancedContext: Boolean? = null, val contextFiles: List? = null, val error: ChatError? = null, + val query: String? = null, ) data class WebviewReceiveMessageParams(val id: String, val message: WebviewMessage) @@ -30,13 +33,15 @@ data class ExtensionMessage( val chatID: String? = null, val isTranscriptError: Boolean? = null, val customPrompts: List>? = null, - val context: Any? = null, - val errors: String? + val context: List? = null, + val errors: String?, + val query: String? = null, ) { object Type { const val TRANSCRIPT = "transcript" const val ERRORS = "errors" + const val USER_CONTEXT_FILES = "userContextFiles" } } diff --git a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt b/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt index be5ef5c9f3..aa0f84fb59 100644 --- a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt +++ b/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt @@ -7,7 +7,9 @@ import com.intellij.openapi.project.Project import com.intellij.ui.components.JBTabbedPane import com.intellij.util.concurrency.annotations.RequiresEdt import com.jetbrains.rd.util.AtomicReference +import com.sourcegraph.cody.agent.* import com.sourcegraph.cody.agent.CodyAgentService +import com.sourcegraph.cody.chat.* import com.sourcegraph.cody.chat.AgentChatSession import com.sourcegraph.cody.chat.AgentChatSessionService import com.sourcegraph.cody.chat.SignInWithSourcegraphPanel diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt index 6336b58a7c..0d35584466 100644 --- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt +++ b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgent.kt @@ -10,6 +10,7 @@ import com.sourcegraph.cody.agent.protocol.* import com.sourcegraph.config.ConfigUtil import java.io.* import java.net.Socket +import java.net.URI import java.nio.file.* import java.util.* import java.util.concurrent.* @@ -103,7 +104,8 @@ private constructor( .initialize( ClientInfo( version = ConfigUtil.getPluginVersion(), - workspaceRootUri = ConfigUtil.getWorkspaceRootPath(project).toUri(), + workspaceRootUri = + ConfigUtil.getWorkspaceRootPath(project).toUri().toString(), extensionConfiguration = ConfigUtil.getAgentConfiguration(project))) .thenApply { info -> logger.info("Connected to Cody agent " + info.name) @@ -180,6 +182,8 @@ private constructor( .registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer) .registerTypeAdapter(ContextFile::class.java, contextFileDeserializer) .registerTypeAdapter(Speaker::class.java, SpeakerSerializer) + .registerTypeAdapter(URI::class.java, uriDeserializer) + .registerTypeAdapter(URI::class.java, uriSerializer) } .setRemoteInterface(CodyAgentServer::class.java) .traceMessages(traceWriter()) diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ClientInfo.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ClientInfo.kt index bc8da6d447..926b8f117e 100644 --- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ClientInfo.kt +++ b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ClientInfo.kt @@ -1,11 +1,10 @@ package com.sourcegraph.cody.agent.protocol import com.sourcegraph.cody.agent.ExtensionConfiguration -import java.net.URI data class ClientInfo( var version: String, - var workspaceRootUri: URI, + var workspaceRootUri: String, var extensionConfiguration: ExtensionConfiguration? = null ) { val name = "JetBrains" diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ContextFile.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ContextFile.kt index a562dd715a..62556d8952 100644 --- a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ContextFile.kt +++ b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/ContextFile.kt @@ -3,28 +3,72 @@ package com.sourcegraph.cody.agent.protocol import com.google.gson.JsonDeserializationContext import com.google.gson.JsonDeserializer import com.google.gson.JsonElement +import com.google.gson.JsonObject +import com.google.gson.JsonSerializer import java.lang.reflect.Type import java.net.URI +import kotlinx.serialization.json.jsonObject -data class ContextFile( - val uri: URI, - val repoName: String?, - val revision: String?, -) +sealed class ContextFile() { + abstract val type: String + abstract val uri: URI + abstract val repoName: String? + abstract val revision: String? +} + +data class ContextFileFile( + override val uri: URI, + override val repoName: String?, + override val revision: String?, +) : ContextFile() { + override val type: String = "file" +} val contextFileDeserializer = - JsonDeserializer { jsonElement: JsonElement, _: Type, _: JsonDeserializationContext -> + JsonDeserializer { jsonElement: JsonElement, typ: Type, context: JsonDeserializationContext -> val jsonObject = jsonElement.asJsonObject - - val uriObj = jsonObject["uri"].asJsonObject - val uri = - URI( - uriObj["scheme"]?.asString, - uriObj["host"]?.asString, - uriObj["path"]?.asString, - /* fragment= */ null) + val uri = context.deserialize(jsonObject["uri"], URI::class.java) val repoName = jsonObject["repoName"]?.asString val revision = jsonObject["revision"]?.asString - ContextFile(uri, repoName, revision) + when (jsonObject["type"]?.asString) { + "file" -> ContextFileFile(uri, repoName, revision) + + // TODO(beyang): should throw an exception here, but we don't because the context field is + // overloaded in the protocol + else -> null + } + } + +val uriDeserializer = + JsonDeserializer { jsonElement: JsonElement, typ: Type, context: JsonDeserializationContext -> + val j = jsonElement.asJsonObject + URI( + j["scheme"]?.asString, + j["authority"]?.asString, + j["path"]?.asString, + j["query"]?.asString, + j["fragment"]?.asString, + ) + } + +val uriSerializer = JsonSerializer { uri: URI, type, context -> + val obj = JsonObject() + obj.addProperty("scheme", uri.scheme) + obj.addProperty("authority", uri.authority) + obj.addProperty("path", uri.path) + obj.addProperty("query", uri.query) + obj.addProperty("fragment", uri.fragment) + obj +} + +fun contextFilesFromList(list: List): List { + val contextFiles = ArrayList() + for (item in list) { + if (item is Map<*, *> && item.get("type") == "file") { + val path = (item.get("uri") as Map<*, *>).get("path") + contextFiles.add(path as String) } + } + return contextFiles +} diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt b/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt index bb79a660fe..f78d32d5dd 100644 --- a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt +++ b/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt @@ -6,10 +6,7 @@ import com.intellij.util.concurrency.annotations.RequiresEdt import com.intellij.xml.util.XmlStringUtil import com.jetbrains.rd.util.AtomicReference import com.sourcegraph.cody.agent.* -import com.sourcegraph.cody.agent.protocol.ChatMessage -import com.sourcegraph.cody.agent.protocol.ChatRestoreParams -import com.sourcegraph.cody.agent.protocol.ChatSubmitMessageParams -import com.sourcegraph.cody.agent.protocol.Speaker +import com.sourcegraph.cody.agent.protocol.* import com.sourcegraph.cody.chat.ui.ChatPanel import com.sourcegraph.cody.commands.CommandId import com.sourcegraph.cody.config.RateLimitStateManager @@ -24,6 +21,7 @@ import com.sourcegraph.telemetry.GraphQlLogger import java.util.* import java.util.concurrent.CompletableFuture import java.util.concurrent.ExecutionException +import java.util.function.Consumer import org.slf4j.LoggerFactory class AgentChatSession @@ -46,6 +44,25 @@ private constructor( init { cancellationToken.get().dispose() + + CodyAgentService.applyAgentOnBackgroundThread(project) { agent -> + agent.client.onWebviewMessage( + this.sessionId.get().get(), + Consumer { message -> this.receiveWebviewExtensionMessage(message) }) + } + } + + private fun receiveWebviewExtensionMessage(message: ExtensionMessage) { + when (message.type) { + ExtensionMessage.Type.USER_CONTEXT_FILES -> { + if (message.context is List<*>) { + this.chatPanel.promptPanel.setContextFilesSelector(message.context) + } + } + else -> { + logger.warn(String.format("CodyToolWindowContent: unknown message type: %s", message.type)) + } + } } fun getPanel(): ChatPanel = chatPanel @@ -68,8 +85,15 @@ private constructor( sessionId.getAndSet(newSessionId) } + override fun sendWebviewMessage(message: WebviewMessage) { + CodyAgentService.applyAgentOnBackgroundThread(project) { agent -> + agent.server.webviewReceiveMessage( + WebviewReceiveMessageParams(this.sessionId.get().get(), message)) + } + } + @RequiresEdt - override fun sendMessage(text: String) { + override fun sendMessage(text: String, contextFiles: List) { val displayText = XmlStringUtil.escapeString(text) val humanMessage = ChatMessage(Speaker.HUMAN, text, displayText) addMessage(humanMessage) @@ -81,8 +105,7 @@ private constructor( text = humanMessage.actualMessage(), submitType = "user", addEnhancedContext = chatPanel.isEnhancedContextEnabled(), - // TODO(#242): allow to manually add files to the context via `@` - contextFiles = listOf()) + contextFiles = contextFiles) val request = agent.server.chatSubmitMessage(ChatSubmitMessageParams(sessionId.get().get(), message)) diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt index f1d5597269..009211037f 100644 --- a/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt +++ b/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt @@ -2,12 +2,17 @@ package com.sourcegraph.cody.chat import com.intellij.util.concurrency.annotations.RequiresEdt import com.sourcegraph.cody.agent.ExtensionMessage +import com.sourcegraph.cody.agent.WebviewMessage +import com.sourcegraph.cody.agent.protocol.* import com.sourcegraph.cody.vscode.CancellationToken typealias SessionId = String interface ChatSession { - @RequiresEdt fun sendMessage(text: String) + + fun sendWebviewMessage(message: WebviewMessage) + + @RequiresEdt fun sendMessage(text: String, contextFiles: List) fun receiveMessage(extensionMessage: ExtensionMessage) diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ui/ChatPanel.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ui/ChatPanel.kt index 15c0a27669..5ec6ecd571 100644 --- a/src/main/kotlin/com/sourcegraph/cody/chat/ui/ChatPanel.kt +++ b/src/main/kotlin/com/sourcegraph/cody/chat/ui/ChatPanel.kt @@ -22,7 +22,7 @@ class ChatPanel(project: Project, val chatSession: ChatSession) : JPanel(VerticalFlowLayout(VerticalFlowLayout.CENTER, 0, 0, true, false)) { private val messagesPanel = MessagesPanel(project) private val chatPanel = ChatScrollPane(messagesPanel) - private val promptPanel: PromptPanel = PromptPanel(chatSession) + val promptPanel: PromptPanel = PromptPanel(chatSession) private val contextView: EnhancedContextPanel = EnhancedContextPanel(project) private val stopGeneratingButton = diff --git a/src/main/kotlin/com/sourcegraph/cody/ui/AutoGrowingTextArea.kt b/src/main/kotlin/com/sourcegraph/cody/ui/AutoGrowingTextArea.kt index 766d8e57b1..4e9fd285fa 100644 --- a/src/main/kotlin/com/sourcegraph/cody/ui/AutoGrowingTextArea.kt +++ b/src/main/kotlin/com/sourcegraph/cody/ui/AutoGrowingTextArea.kt @@ -20,6 +20,7 @@ import javax.swing.plaf.basic.BasicTextAreaUI import javax.swing.text.AttributeSet import javax.swing.text.Document import javax.swing.text.PlainDocument +import javax.swing.undo.UndoManager import kotlin.math.max import kotlin.math.min @@ -28,6 +29,7 @@ class AutoGrowingTextArea(private val minRows: Int, maxRows: Int, outerPanel: JC val scrollPane: JBScrollPane private val initialPreferredSize: Dimension private val autoGrowUpToRow: Int + val undoManager = UndoManager() init { autoGrowUpToRow = maxRows + 1 @@ -48,10 +50,21 @@ class AutoGrowingTextArea(private val minRows: Int, maxRows: Int, outerPanel: JC outerPanel.revalidate() } } + textArea.document = document + document.addUndoableEditListener { event -> undoManager.addEdit(event.edit) } + updateTextAreaSize() } + fun getText(): String { + return this.textArea.text + } + + fun setText(newText: String) { + textArea.text = newText + } + private fun createTextArea(): JBTextArea { val promptInput: JBTextArea = RoundedJBTextArea(minRows, 10) val textUI = DarculaTextAreaUI.createUI(promptInput) as BasicTextAreaUI diff --git a/src/test/kotlin/com/sourcegraph/cody/PromptPanelTest.kt b/src/test/kotlin/com/sourcegraph/cody/PromptPanelTest.kt new file mode 100644 index 0000000000..a7e83eab34 --- /dev/null +++ b/src/test/kotlin/com/sourcegraph/cody/PromptPanelTest.kt @@ -0,0 +1,42 @@ +package com.sourcegraph.cody + +import junit.framework.TestCase + +class PromptPanelTest : TestCase() { + fun `test findAtExpressions`() { + data class Case(val text: String, val expected: List) + val cases = + listOf( + Case( + "@some-file what does this file do?", + listOf(AtExpression(0, "@some-file".length, "@some-file", "some-file"))), + Case( + "foo @file1 @file2 bar @file3", + listOf( + AtExpression("foo ".length, "foo ".length + "@file1".length, "@file1", "file1"), + AtExpression( + "foo @file1 ".length, + "foo @file1 ".length + "@file2".length, + "@file2", + "file2"), + AtExpression( + "foo @file1 @file2 bar ".length, + "foo @file1 @file2 bar ".length + "@file3".length, + "@file3", + "file3"), + )), + Case( + """foo @file\ with\ spaces bar""", + listOf( + AtExpression( + "foo ".length, + "foo ".length + "@file\\ with\\ spaces".length, + "@file\\ with\\ spaces", + "file with spaces"), + ))) + + for (case in cases) { + assertEquals(case.expected, findAtExpressions(case.text)) + } + } +}