diff --git a/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java b/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java
index b58484ba94..1830eb3f70 100644
--- a/src/main/java/com/sourcegraph/cody/agent/CodyAgentClient.java
+++ b/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;
     }
 
diff --git a/src/main/java/com/sourcegraph/cody/attribution/AttributionListener.kt b/src/main/java/com/sourcegraph/cody/attribution/AttributionListener.kt
new file mode 100644
index 0000000000..128fc9fbd0
--- /dev/null
+++ b/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) }
+    }
+  }
+}
diff --git a/src/main/java/com/sourcegraph/cody/attribution/AttributionSearchCommand.kt b/src/main/java/com/sourcegraph/cody/attribution/AttributionSearchCommand.kt
new file mode 100644
index 0000000000..52bcfba7e6
--- /dev/null
+++ b/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)
+}
diff --git a/src/main/java/com/sourcegraph/cody/chat/CodeEditorFactory.java b/src/main/java/com/sourcegraph/cody/chat/CodeEditorFactory.java
index 4c0c8dacbf..1ea9ba0d1b 100644
--- a/src/main/java/com/sourcegraph/cody/chat/CodeEditorFactory.java
+++ b/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;
   }
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt b/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
index 87fdd4e04f..8a4fa21fe4 100644
--- a/src/main/kotlin/com/sourcegraph/cody/agent/CodyAgentServer.kt
+++ b/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>
 }
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/AttributionSearchParams.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/AttributionSearchParams.kt
new file mode 100644
index 0000000000..6cd2407024
--- /dev/null
+++ b/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,
+)
diff --git a/src/main/kotlin/com/sourcegraph/cody/agent/protocol/AttributionSearchResponse.kt b/src/main/kotlin/com/sourcegraph/cody/agent/protocol/AttributionSearchResponse.kt
new file mode 100644
index 0000000000..ba059c854a
--- /dev/null
+++ b/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,
+)
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt b/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt
index 48ed1da480..9d3cf04e6a 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt
+++ b/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
 
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSessionService.kt b/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSessionService.kt
index 0b3a556d2a..6659faa84f 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSessionService.kt
+++ b/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) }
   }
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt
index 15d26b5fab..e23edf1c78 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/ChatSession.kt
+++ b/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>)
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 5ec6ecd571..f8c4ccf1be 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/ui/ChatPanel.kt
+++ b/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)
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ui/CodeEditorButtons.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ui/CodeEditorButtons.kt
index e1f964f5be..33d570ee39 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/ui/CodeEditorButtons.kt
+++ b/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
+      }
     }
   }
 
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagePart.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagePart.kt
index 78b3da002b..f05ec426b6 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagePart.kt
+++ b/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()) })
   }
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagesPanel.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagesPanel.kt
index 2ea0fbdf03..06bf276354 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/ui/MessagesPanel.kt
+++ b/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
   }
 }
diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/ui/SingleMessagePanel.kt b/src/main/kotlin/com/sourcegraph/cody/chat/ui/SingleMessagePanel.kt
index 5533c75065..5e0529f2a8 100644
--- a/src/main/kotlin/com/sourcegraph/cody/chat/ui/SingleMessagePanel.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/chat/ui/SingleMessagePanel.kt
@@ -6,6 +6,8 @@ import com.intellij.util.ui.SwingHelper
 import com.intellij.util.ui.UIUtil
 import com.sourcegraph.cody.agent.protocol.ChatMessage
 import com.sourcegraph.cody.agent.protocol.Speaker
+import com.sourcegraph.cody.attribution.AttributionListener
+import com.sourcegraph.cody.attribution.AttributionSearchCommand
 import com.sourcegraph.cody.chat.*
 import com.sourcegraph.cody.ui.HtmlViewer.createHtmlViewer
 import java.awt.Color
@@ -22,6 +24,7 @@ class SingleMessagePanel(
     private val project: Project,
     private val parentPanel: JPanel,
     private val gradientWidth: Int,
+    private val chatSession: ChatSession,
 ) : PanelWithGradientBorder(gradientWidth, chatMessage.speaker) {
   private var lastMessagePart: MessagePart? = null
 
@@ -50,6 +53,11 @@ class SingleMessagePanel(
     if (lastPart is CodeEditorPart) {
       lastPart.updateCode(project, code, language)
     } else {
+      // For completeness of [onPartFinished] semantics.
+      // At this point the implementation only considers
+      // lastMessagePart if it is CodeEditorPart, so this
+      // is always no-op.
+      onPartFinished()
       addAsNewCodeComponent(code, language)
     }
   }
@@ -66,6 +74,7 @@ class SingleMessagePanel(
     if (lastPart is TextPart) {
       lastPart.updateText(text)
     } else {
+      onPartFinished()
       addAsNewTextComponent(text)
     }
   }
@@ -83,6 +92,28 @@ class SingleMessagePanel(
     else ColorUtil.brighter(UIUtil.getPanelBackground(), 3)
   }
 
+  /**
+   * Trigger attribution search if the part that finished is a code snippet.
+   *
+   * Call sites should include:
+   * - including new text component after writing a code snippet (triggers attribution search
+   *   mid-chat message).
+   * - including new code component after writing a text snippet (no-op because the implementation
+   *   only considers [CodeEditorPart] [lastMessagePart], but added for completeness of
+   *   [onPartFinished] semantics.
+   * - in a cancellation token callback in [MessagesPanel] (triggering attribution search if code
+   *   snippet is the final part as well as if Cody's typing is cancelled.
+   */
+  fun onPartFinished() {
+    val lastPart = lastMessagePart
+    if (lastPart is CodeEditorPart) {
+      chatSession.getSessionId()?.let { sessionId ->
+        val listener = AttributionListener.UiThreadDecorator(lastPart.attribution)
+        AttributionSearchCommand(project).onSnippetFinished(lastPart.text, sessionId, listener)
+      }
+    }
+  }
+
   companion object {
     private val extensions = listOf(TablesExtension.create())
 
diff --git a/src/main/kotlin/com/sourcegraph/cody/ui/AttributionButtonController.kt b/src/main/kotlin/com/sourcegraph/cody/ui/AttributionButtonController.kt
new file mode 100644
index 0000000000..57c3d33105
--- /dev/null
+++ b/src/main/kotlin/com/sourcegraph/cody/ui/AttributionButtonController.kt
@@ -0,0 +1,61 @@
+package com.sourcegraph.cody.ui
+
+import com.intellij.openapi.project.Project
+import com.intellij.util.concurrency.annotations.RequiresEdt
+import com.sourcegraph.cody.agent.CurrentConfigFeatures
+import com.sourcegraph.cody.agent.protocol.AttributionSearchResponse
+import com.sourcegraph.cody.attribution.AttributionListener
+import com.sourcegraph.common.CodyBundle
+import com.sourcegraph.common.CodyBundle.fmt
+
+class AttributionButtonController(val button: ConditionalVisibilityButton) : AttributionListener {
+
+  private val extraUpdates: MutableList<Runnable> = ArrayList()
+
+  companion object {
+    fun setup(project: Project): AttributionButtonController {
+      val button =
+          ConditionalVisibilityButton(CodyBundle.getString("chat.attribution.searching.label"))
+      button.isEnabled = false // non-clickable
+      val currentConfigFeatures: CurrentConfigFeatures =
+          project.getService(CurrentConfigFeatures::class.java)
+      // Only display the button if attribution is enabled.
+      button.visibilityAllowed = currentConfigFeatures.get().attribution
+      return AttributionButtonController(button)
+    }
+  }
+
+  @RequiresEdt
+  override fun onAttributionSearchStart() {
+    button.toolTipText = CodyBundle.getString("chat.attribution.searching.tooltip")
+  }
+
+  @RequiresEdt
+  override fun updateAttribution(attribution: AttributionSearchResponse) {
+    if (attribution.error != null) {
+      button.text = CodyBundle.getString("chat.attribution.error.label")
+      button.toolTipText =
+          CodyBundle.getString("chat.attribution.error.tooltip").fmt(attribution.error)
+    } else if (attribution.repoNames.isEmpty()) {
+      button.text = CodyBundle.getString("chat.attribution.success.label")
+      button.toolTipText = CodyBundle.getString("chat.attribution.success.tooltip")
+    } else {
+      val count = "${attribution.repoNames.size}" + if (attribution.limitHit) "+" else ""
+      val repoNames =
+          attribution.repoNames.joinToString(
+              prefix = "<ul><li>", separator = "</li><li>", postfix = "</li></ul>")
+      button.text = CodyBundle.getString("chat.attribution.failure.label")
+      button.toolTipText =
+          CodyBundle.getString("chat.attribution.failure.tooltip").fmt(count, repoNames)
+    }
+    button.updatePreferredSize()
+    for (action in extraUpdates) {
+      action.run()
+    }
+  }
+
+  /** Run extra actions on button update, like resizing components. */
+  fun onUpdate(action: Runnable) {
+    extraUpdates += action
+  }
+}
diff --git a/src/main/kotlin/com/sourcegraph/cody/ui/ConditionalVisibilityButton.kt b/src/main/kotlin/com/sourcegraph/cody/ui/ConditionalVisibilityButton.kt
new file mode 100644
index 0000000000..b2e386c0d3
--- /dev/null
+++ b/src/main/kotlin/com/sourcegraph/cody/ui/ConditionalVisibilityButton.kt
@@ -0,0 +1,31 @@
+package com.sourcegraph.cody.ui
+
+import java.awt.Dimension
+
+/**
+ * [ConditionalVisibilityButton] is only made visible if visibility is allowed.
+ *
+ * This is to implement a hover visibility that is conditional on another factor, like enabling
+ * attribution setting.
+ */
+class ConditionalVisibilityButton(text: String) : TransparentButton(text) {
+
+  var visibilityAllowed: Boolean = true
+    set(value) {
+      field = value
+      if (!value) {
+        super.setVisible(false)
+      }
+    }
+
+  override fun setVisible(value: Boolean) {
+    if ((value && visibilityAllowed) // either make visible if visibility allowed
+    || (!value) // or make invisible
+    ) {
+      super.setVisible(value)
+    }
+  }
+
+  override fun getPreferredSize(): Dimension =
+      if (visibilityAllowed) super.getPreferredSize() else Dimension(0, 0)
+}
diff --git a/src/main/kotlin/com/sourcegraph/cody/ui/TransparentButton.kt b/src/main/kotlin/com/sourcegraph/cody/ui/TransparentButton.kt
index 2d96bfcc32..4b10af3b28 100644
--- a/src/main/kotlin/com/sourcegraph/cody/ui/TransparentButton.kt
+++ b/src/main/kotlin/com/sourcegraph/cody/ui/TransparentButton.kt
@@ -6,8 +6,9 @@ import com.intellij.util.ui.UIUtil
 import java.awt.*
 import javax.swing.JButton
 
-class TransparentButton(text: String) : JButton(text) {
+open class TransparentButton(text: String) : JButton(text) {
   private val cornerRadius = 5
+  private val fontMetric: FontMetrics
 
   init {
     isContentAreaFilled = false
@@ -15,12 +16,16 @@ class TransparentButton(text: String) : JButton(text) {
     isBorderPainted = false
     isVisible = false
 
-    // Calculate the preferred size based on the size of the text
-    val fm = getFontMetrics(font)
+    this.fontMetric = getFontMetrics(font)
+    updatePreferredSize()
+  }
+
+  /** Calculate the preferred size based on the size of the text. */
+  fun updatePreferredSize() {
     val horizontalPadding = 10
     val verticalPadding = 5
-    val width = fm.stringWidth(getText()) + horizontalPadding * 2
-    val height = fm.height + verticalPadding * 2
+    val width = fontMetric.stringWidth(getText()) + horizontalPadding * 2
+    val height = fontMetric.height + verticalPadding * 2
     preferredSize = Dimension(width, height)
   }
 
diff --git a/src/main/resources/CodyBundle.properties b/src/main/resources/CodyBundle.properties
index 1e8190b9cd..b969676112 100644
--- a/src/main/resources/CodyBundle.properties
+++ b/src/main/resources/CodyBundle.properties
@@ -43,6 +43,14 @@ chat.rate-limit-error.explain=\
     To ensure that Cody can stay operational for all Cody users, please come back tomorrow for more chats, commands, and autocompletes. \
     <a href=\"https://sourcegraph.com/docs/cody/core-concepts/cody-gateway#rate-limits-and-quotas\">Learn more.</a>\
   </html>
+chat.attribution.searching.label=Attribution search
+chat.attribution.searching.tooltip=Guard Rails: Running Code Attribution Check...
+chat.attribution.error.label=Guard Rails API Error
+chat.attribution.error.tooltip=Guard Rails API Error: {0}
+chat.attribution.success.label=Guard Rails Check Passed
+chat.attribution.success.tooltip=Snippet not found on Sourcegraph.com
+chat.attribution.failure.label=Guard Rails Check Failed
+chat.attribution.failure.tooltip=<html>Guard Rails Check Failed. Code found in {0} repositories: {1}</html>
 subscription-tab.chat-rate-limit-error=You've used all your chat messages and commands for the month. Upgrade to Pro for unlimited usage.
 subscription-tab.autocomplete-rate-limit-error=You've used all your autocompletions for the month. Upgrade to Pro for unlimited usage.
 subscription-tab.chat-and-autocomplete-rate-limit-error=You've used all your autocompletions and chats for this month. Upgrade to Cody Pro to get unlimited interactions.