Attribution Chat UI

src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java

@@ -1,13 +1,10 @@
 package com.sourcegraph.cody.agent;
 
-import com.intellij.openapi.Disposable;
 import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.diagnostic.Logger;
 import com.intellij.openapi.editor.Editor;
 import com.sourcegraph.cody.agent.protocol.DebugMessage;
-import java.lang.ref.WeakReference;
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification;
@@ -25,7 +22,8 @@ public class CodyAgentClient {
   // Callback that is invoked when the agent sends a "setConfigFeatures" message.
   @Nullable public ConfigFeaturesObserver onSetConfigFeatures;
 
-  // Callback that is invoked on webview messages which aren't handled by onNewMessage or onSetConfigFeatures
+  // Callback that is invoked on webview messages which aren't handled by onNewMessage or
+  // onSetConfigFeatures
   @Nullable public Consumer<WebviewPostMessageParams> onReceivedWebviewMessage;
 
   @Nullable public Editor editor;
@@ -83,7 +81,8 @@ public void webviewPostMessage(WebviewPostMessageParams params) {
     }
 
     if (onReceivedWebviewMessage != null) {
-      ApplicationManager.getApplication().invokeLater(() -> onReceivedWebviewMessage.accept(params));
+      ApplicationManager.getApplication()
+          .invokeLater(() -> onReceivedWebviewMessage.accept(params));
       return;
     }
 

src/main/java/com/sourcegraph/cody/attribution/AttributionListener.kt

@@ -0,0 +1,32 @@
+package com.sourcegraph.cody.attribution
+
+import com.intellij.openapi.application.ApplicationManager
+import com.sourcegraph.cody.agent.protocol.AttributionSearchResponse
+
+/**
+ * [AttributionListener] responds to attribution search state changes.
+ *
+ * The interface does not convey any contract about execution thread. The caller and callee should
+ * make sure of proper execution.
+ */
+interface AttributionListener {
+  /** Notifies the listener that attribution search has started. */
+  fun onAttributionSearchStart()
+
+  /** Notifies the listener of the result of attribution search. */
+  fun updateAttribution(attribution: AttributionSearchResponse)
+
+  /**
+   * Wraps given [AttributionListener] so that all notifications are delivered asynchronously on UI
+   * thread.
+   */
+  class UiThreadDecorator(private val delegate: AttributionListener) : AttributionListener {
+    override fun onAttributionSearchStart() {
+      ApplicationManager.getApplication().invokeLater { delegate.onAttributionSearchStart() }
+    }
+
+    override fun updateAttribution(attribution: AttributionSearchResponse) {
+      ApplicationManager.getApplication().invokeLater { delegate.updateAttribution(attribution) }
+    }
+  }
+}

src/main/java/com/sourcegraph/cody/attribution/AttributionSearchCommand.kt

@@ -0,0 +1,57 @@
+package com.sourcegraph.cody.attribution
+
+import com.intellij.openapi.application.ApplicationManager
+import com.intellij.openapi.project.Project
+import com.sourcegraph.cody.agent.CodyAgentService
+import com.sourcegraph.cody.agent.CurrentConfigFeatures
+import com.sourcegraph.cody.agent.protocol.AttributionSearchParams
+import com.sourcegraph.cody.agent.protocol.AttributionSearchResponse
+import com.sourcegraph.cody.chat.AgentChatSession
+import com.sourcegraph.cody.chat.AgentChatSessionService
+import com.sourcegraph.cody.chat.SessionId
+import com.sourcegraph.cody.chat.ui.CodeEditorPart
+import java.util.*
+import java.util.function.BiFunction
+
+/**
+ * [AttributionSearchCommand] performs attribution search on a code snippet, and then notifies of
+ * the result.
+ */
+class AttributionSearchCommand(private val project: Project) {
+
+  /**
+   * [onSnippetFinished] invoked when assistant finished writing a code snippet in a chat message,
+   * and triggers attribution search (if enabled). Once attribution returns, the
+   * [CodeEditorPart.attributionListener] is updated.
+   */
+  fun onSnippetFinished(snippet: String, sessionId: SessionId, listener: AttributionListener) {
+    if (attributionEnabled()) {
+      CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
+        ApplicationManager.getApplication().invokeLater { listener.onAttributionSearchStart() }
+        val params = AttributionSearchParams(id = sessionId, snippet = snippet)
+        agent.server.attributionSearch(params).handle(updateEditor(listener))
+      }
+    }
+  }
+
+  /**
+   * [updateEditor] returns a future handler for attribution search operation, which notifies the
+   * listener.
+   */
+  private fun updateEditor(listener: AttributionListener) =
+      BiFunction<AttributionSearchResponse?, Throwable?, Unit> { response, throwable ->
+        listener.updateAttribution(
+            response
+                ?: AttributionSearchResponse(
+                    error = throwable?.message ?: "Error searching for attribution.",
+                    repoNames = listOf(),
+                    limitHit = false,
+                ))
+      }
+
+  private fun attributionEnabled(): Boolean =
+      project.getService(CurrentConfigFeatures::class.java).get().attribution
+
+  private fun findChatSessionFor(messageId: UUID): AgentChatSession? =
+      AgentChatSessionService.getInstance(project).findByMessage(messageId)
+}

src/main/java/com/sourcegraph/cody/chat/CodeEditorFactory.java

@@ -19,6 +19,7 @@
 import com.intellij.util.ui.JBInsets;
 import com.sourcegraph.cody.chat.ui.CodeEditorButtons;
 import com.sourcegraph.cody.chat.ui.CodeEditorPart;
+import com.sourcegraph.cody.ui.AttributionButtonController;
 import com.sourcegraph.cody.ui.TransparentButton;
 import java.awt.Dimension;
 import java.awt.Insets;
@@ -64,6 +65,9 @@ public CodeEditorPart createCodeEditor(@NotNull String code, @Nullable String la
     insertAtCursorButton.setToolTipText("Insert text at current cursor position");
     insertAtCursorButton.addActionListener(insertAtCursorActionListener(editor));
 
+    AttributionButtonController attributionButtonController =
+        AttributionButtonController.Companion.setup(project);
+
     Dimension copyButtonPreferredSize = copyButton.getPreferredSize();
     int halfOfButtonHeight = copyButtonPreferredSize.height / 2;
     JLayeredPane layeredEditorPane = new JLayeredPane();
@@ -89,9 +93,16 @@ public CodeEditorPart createCodeEditor(@NotNull String code, @Nullable String la
         editorPreferredSize.height + halfOfButtonHeight);
     layeredEditorPane.add(editorComponent, JLayeredPane.DEFAULT_LAYER);
 
-    JButton[] buttons = new JButton[] {copyButton, insertAtCursorButton};
+    // Rendering order of buttons is right-to-left:
+    JButton[] buttons =
+        new JButton[] {attributionButtonController.getButton(), copyButton, insertAtCursorButton};
     CodeEditorButtons codeEditorButtons = new CodeEditorButtons(buttons);
     codeEditorButtons.addButtons(layeredEditorPane, editorComponent.getWidth());
+    attributionButtonController.onUpdate(
+        () -> {
+          // Resize buttons on text update.
+          codeEditorButtons.updateBounds(editorComponent.getWidth());
+        });
 
     // resize the editor and move the copy button when the parent panel is resized
     layeredEditorPane.addComponentListener(
@@ -144,7 +155,8 @@ public void mouseExited(@NotNull EditorMouseEvent event) {
 
     editor.addEditorMouseMotionListener(editorMouseMotionListener);
     editor.addEditorMouseListener(editorMouseListener);
-    CodeEditorPart codeEditorPart = new CodeEditorPart(layeredEditorPane, editor);
+    CodeEditorPart codeEditorPart =
+        new CodeEditorPart(layeredEditorPane, editor, attributionButtonController);
     codeEditorPart.updateLanguage(language);
     return codeEditorPart;
   }

src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt

@@ -100,4 +100,9 @@ interface CodyAgentServer {
   fun chatModels(params: ChatModelsParams): CompletableFuture<ChatModelsResponse>
 
   @JsonRequest("chat/restore") fun chatRestore(params: ChatRestoreParams): CompletableFuture<String>
+
+  @JsonRequest("attribution/search")
+  fun attributionSearch(
+      params: AttributionSearchParams
+  ): CompletableFuture<AttributionSearchResponse>
 }

src/main/kotlin/com/sourcegraph/cody/agent/protocol/AttributionSearchParams.kt

@@ -0,0 +1,6 @@
+package com.sourcegraph.cody.agent.protocol
+
+data class AttributionSearchParams(
+    val id: String,
+    val snippet: String,
+)

src/main/kotlin/com/sourcegraph/cody/agent/protocol/AttributionSearchResponse.kt

@@ -0,0 +1,7 @@
+package com.sourcegraph.cody.agent.protocol
+
+data class AttributionSearchResponse(
+    val error: String?,
+    val repoNames: List<String>,
+    val limitHit: Boolean,
+)

src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt

@@ -62,8 +62,11 @@ private constructor(
 
   fun getPanel(): ChatPanel = chatPanel
 
-  fun hasSessionId(thatSessionId: SessionId): Boolean =
-      sessionId.get().getNow(null) == thatSessionId
+  override fun getSessionId(): SessionId? = sessionId.get().getNow(null)
+
+  fun hasSessionId(thatSessionId: SessionId): Boolean = getSessionId() == thatSessionId
+
+  fun hasMessageId(messageId: UUID): Boolean = messages.any { it.id == messageId }
 
   override fun getInternalId(): String = internalId
 

src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSessionService.kt

@@ -6,6 +6,7 @@ import com.intellij.openapi.project.Project
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.history.state.ChatState
+import java.util.UUID
 import java.util.concurrent.ConcurrentLinkedQueue
 
 @Service(Service.Level.PROJECT)
@@ -34,6 +35,9 @@ class AgentChatSessionService(private val project: Project) {
   fun getSession(sessionId: SessionId): AgentChatSession? =
       chatSessions.find { it.hasSessionId(sessionId) }
 
+  fun findByMessage(messageId: UUID): AgentChatSession? =
+      chatSessions.find { it.hasMessageId(messageId) }
+
   fun restoreAllSessions(agent: CodyAgent) {
     chatSessions.forEach { it.restoreAgentSession(agent) }
   }

src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt

@@ -10,6 +10,8 @@ typealias SessionId = String
 
 interface ChatSession {
 
+  fun getSessionId(): SessionId?
+
   fun sendWebviewMessage(message: WebviewMessage)
 
   @RequiresEdt fun sendMessage(text: String, contextFiles: List<ContextFile>)

src/main/kotlin/com/sourcegraph/cody/chat/ui/ChatPanel.kt

@@ -20,7 +20,7 @@ import javax.swing.JPanel
 
 class ChatPanel(project: Project, val chatSession: ChatSession) :
     JPanel(VerticalFlowLayout(VerticalFlowLayout.CENTER, 0, 0, true, false)) {
-  private val messagesPanel = MessagesPanel(project)
+  private val messagesPanel = MessagesPanel(project, chatSession)
   private val chatPanel = ChatScrollPane(messagesPanel)
   val promptPanel: PromptPanel = PromptPanel(chatSession)
   private val contextView: EnhancedContextPanel = EnhancedContextPanel(project)

src/main/kotlin/com/sourcegraph/cody/chat/ui/CodeEditorButtons.kt

@@ -41,7 +41,9 @@ class CodeEditorButtons(val buttons: Array<JButton>) {
           0,
           jButtonPreferredSize.width,
           jButtonPreferredSize.height)
-      shift += jButtonPreferredSize.width + CodeEditorFactory.spaceBetweenButtons
+      if (jButtonPreferredSize.width > 0) { // Do not add space for collapsed button.
+        shift += jButtonPreferredSize.width + CodeEditorFactory.spaceBetweenButtons
+      }
     }
   }
 

src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagePart.kt

@@ -11,6 +11,8 @@ import com.intellij.openapi.fileTypes.PlainTextFileType
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Computable
 import com.intellij.util.ui.SwingHelper
+import com.sourcegraph.cody.ui.AttributionButtonController
+import java.util.concurrent.atomic.AtomicReference
 import javax.swing.JComponent
 import javax.swing.JEditorPane
 
@@ -22,7 +24,18 @@ class TextPart(val component: JEditorPane) : MessagePart {
   }
 }
 
-class CodeEditorPart(val component: JComponent, private val editor: EditorEx) : MessagePart {
+class CodeEditorPart(
+    val component: JComponent,
+    private val editor: EditorEx,
+    val attribution: AttributionButtonController
+) : MessagePart {
+
+  private val _text = AtomicReference("")
+  var text: String
+    set(value) {
+      _text.set(value)
+    }
+    get() = _text.get()
 
   fun updateCode(project: Project, code: String, language: String?) {
     updateLanguage(language)
@@ -42,6 +55,7 @@ class CodeEditorPart(val component: JComponent, private val editor: EditorEx) :
   }
 
   private fun updateText(project: Project, text: String) {
+    this.text = text
     WriteCommandAction.runWriteCommandAction(
         project, Computable { editor.document.replaceText(text, System.currentTimeMillis()) })
   }

src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagesPanel.kt

@@ -6,12 +6,13 @@ import com.intellij.openapi.ui.VerticalFlowLayout
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.protocol.ChatMessage
 import com.sourcegraph.cody.agent.protocol.Speaker
+import com.sourcegraph.cody.chat.ChatSession
 import com.sourcegraph.cody.chat.ChatUIConstants
 import com.sourcegraph.cody.vscode.CancellationToken
 import com.sourcegraph.common.CodyBundle
 import javax.swing.JPanel
 
-class MessagesPanel(private val project: Project) :
+class MessagesPanel(private val project: Project, private val chatSession: ChatSession) :
     JPanel(VerticalFlowLayout(VerticalFlowLayout.TOP, 0, 0, true, true)) {
   init {
     val welcomeText = CodyBundle.getString("messages-panel.welcome-text")
@@ -24,8 +25,7 @@ class MessagesPanel(private val project: Project) :
     removeBlinkingCursor()
 
     if (componentCount > 0) {
-      val lastPanel = components.last() as? JPanel
-      val lastMessage = lastPanel?.getComponent(0) as? SingleMessagePanel
+      val lastMessage = getLastMessage()
       if (message.id == lastMessage?.getMessageId()) {
         lastMessage.updateContentWith(message)
       } else {
@@ -50,7 +50,10 @@ class MessagesPanel(private val project: Project) :
 
   fun registerCancellationToken(cancellationToken: CancellationToken) {
     cancellationToken.onFinished {
-      ApplicationManager.getApplication().invokeLater { removeBlinkingCursor() }
+      ApplicationManager.getApplication().invokeLater {
+        removeBlinkingCursor()
+        getLastMessage()?.onPartFinished()
+      }
     }
   }
 
@@ -68,6 +71,11 @@ class MessagesPanel(private val project: Project) :
   private fun addChatMessageAsComponent(message: ChatMessage) {
     addComponentToChat(
         SingleMessagePanel(
-            message, project, this, ChatUIConstants.ASSISTANT_MESSAGE_GRADIENT_WIDTH))
+            message, project, this, ChatUIConstants.ASSISTANT_MESSAGE_GRADIENT_WIDTH, chatSession))
+  }
+
+  private fun getLastMessage(): SingleMessagePanel? {
+    val lastPanel = components.last() as? JPanel
+    return lastPanel?.getComponent(0) as? SingleMessagePanel
   }
 }