Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Implement user-selected context files #367

Merged
merged 8 commits into from
Feb 6, 2024
Merged

Implement user-selected context files #367

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
68ae69e
basic context file menu
beyang Feb 1, 2024
ae7d1ec
respond to PR comments
beyang Feb 2, 2024
585905e
fix displayPath; re-register webview message listener when agent sess…
beyang Feb 3, 2024
d6c3dd1
fix list appearance
beyang Feb 3, 2024
ed300e4
refactor
beyang Feb 3, 2024
43f608d
Merge remote-tracking branch 'origin/main' into bl/wip
beyang Feb 3, 2024
2192c16
Merge remote-tracking branch 'origin/main' into bl/wip
beyang Feb 5, 2024
27dfaa6
make webview message handling consistent with other message handling
beyang Feb 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
301 changes: 222 additions & 79 deletions src/main/java/com/sourcegraph/cody/PromptPanel.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,67 +1,72 @@
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
import com.intellij.openapi.application.ApplicationManager
import com.intellij.openapi.project.DumbAwareAction
import com.intellij.ui.DocumentAdapter
import com.intellij.ui.components.JBList
import com.intellij.ui.components.JBScrollPane
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
import com.sourcegraph.cody.ui.AutoGrowingTextArea
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 javax.swing.JLayeredPane
import javax.swing.KeyStroke
import java.awt.event.*
import java.io.File
import javax.swing.*
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 contextFilesListViewModel = DefaultListModel<DisplayedContextFile>()
private val contextFilesListView = JBList(contextFilesListViewModel)
private val contextFilesContainer = JBScrollPane(contextFilesListView)

/** Externally updated state */
private val selectedContextFiles: ArrayList<ContextFile> = 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())
}
}
}
contextFilesListView.disableEmptyText()
add(contextFilesContainer, PALETTE_LAYER, 0)

add(sendButton, PALETTE_LAYER, 0)

preferredSize = Dimension(scrollPane.width, scrollPane.height)

/** Add listeners */
addAncestorListener(
object : AncestorListener {
override fun ancestorAdded(event: AncestorEvent?) {
Expand All @@ -73,87 +78,225 @@ 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()
contextFilesListView.addMouseListener(
object : MouseAdapter() {
override fun mouseClicked(e: MouseEvent) {
contextFilesListView.selectedIndex = contextFilesListView.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 (contextFilesListView.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 (contextFilesListView.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 = contextFilesListView.model.getElementAt(contextFilesListView.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 =
(contextFilesListView.selectedIndex + increment) % contextFilesListView.model.size
if (newSelectedIndex < 0) {
newSelectedIndex += contextFilesListView.model.size
}
contextFilesListView.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 contextFilesContainerHeight =
if (contextFilesListViewModel.isEmpty) 0 else contextFilesListView.preferredSize.height + 2
if (contextFilesContainerHeight == 0) {
contextFilesContainer.isVisible = false
} else {
contextFilesContainer.size = Dimension(scrollPane.width, contextFilesContainerHeight)
contextFilesContainer.isVisible = true
}

scrollPane.setBounds(0, contextFilesContainerHeight, width, scrollPane.preferredSize.height)
preferredSize = Dimension(scrollPane.width, scrollPane.height + contextFilesContainerHeight)

sendButton.setLocation(
scrollPane.width - sendButton.preferredSize.width,
scrollPane.height + contextFilesContainerHeight - 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<ContextFile>) {
val changed = contextFilesListViewModel.elements().toList() != newUserContextFiles
if (changed) {
val newModel = DefaultListModel<DisplayedContextFile>()
newModel.addAll(newUserContextFiles.map { f -> DisplayedContextFile(f) })
contextFilesListView.model = newModel
contextFilesListViewModel = newModel

private fun getTextAndReset(): String {
val text = textArea.text
promptMessageHistory.messageSent(text)
textArea.text = ""
return text
if (newUserContextFiles.isNotEmpty()) {
contextFilesListView.selectedIndex = 0
} else {
contextFilesListView.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<AtExpression> {
val matches = atExpressionPattern.findAll(text)
val expressions = ArrayList<AtExpression>()
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<ContextFile>, text: String): List<ContextFile> {
val atExpressions = findAtExpressions(text)
return contextFiles.filter { f -> atExpressions.any { it.value == displayPath(f) } }
}

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)
// TODO(beyang): temporary displayPath implementation. This should be replaced by acquiring the
// display path from the agent
// Current behavior: if the path contains more than three components, display the last three.
fun displayPath(contextFile: ContextFile): String {
val path = contextFile.uri.path
val pathComponents = path.split("/") // uri path is posix-style
if (pathComponents.size > 3) {
return "...${File.separator}${pathComponents.subList(pathComponents.size - 3, pathComponents.size).joinToString(File.separator)}"
}
return path.replace("/", File.separator)
}
Loading
Loading