Add multi-repo enterprise context

src/main/java/com/sourcegraph/cody/PromptPanel.kt

@@ -5,6 +5,7 @@ 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.openapi.project.Project
 import com.intellij.ui.DocumentAdapter
 import com.intellij.ui.components.JBList
 import com.intellij.ui.components.JBScrollPane
@@ -27,9 +28,7 @@ import javax.swing.event.AncestorEvent
 import javax.swing.event.AncestorListener
 import javax.swing.event.DocumentEvent
 
-class PromptPanel(
-    private val chatSession: ChatSession,
-) : JLayeredPane() {
+class PromptPanel(project: Project, private val chatSession: ChatSession) : JLayeredPane() {
 
   /** View components */
   private val autoGrowingTextArea = AutoGrowingTextArea(5, 9, this)
@@ -45,7 +44,7 @@ class PromptPanel(
 
   /** Related components */
   private val promptMessageHistory =
-      CodyChatMessageHistory(CHAT_MESSAGE_HISTORY_CAPACITY, chatSession)
+      CodyChatMessageHistory(project, CHAT_MESSAGE_HISTORY_CAPACITY, chatSession)
 
   init {
     /** Initialize view */

src/main/java/com/sourcegraph/cody/agent/WebviewPostMessageParams.kt

@@ -3,8 +3,7 @@ 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.*
+import com.sourcegraph.cody.agent.protocol.Repo
 
 /**
  * A message sent from the webview to the extension host. See vscode/src/chat/protocol.ts for the
@@ -18,6 +17,8 @@ data class WebviewMessage(
     val contextFiles: List<ContextFile>? = null,
     val error: ChatError? = null,
     val query: String? = null,
+    val explicitRepos: List<Repo>? = null,
+    val repoId: String? = null
 )
 
 data class WebviewReceiveMessageParams(val id: String, val message: WebviewMessage)

src/main/java/com/sourcegraph/vcs/RepoUtil.java

@@ -159,7 +159,7 @@ private static String doReplacements(
                 .thenCompose(
                     agent ->
                         agent.getServer().convertGitCloneURLToCodebaseName(new CloneURL(cloneURL)))
-                .completeOnTimeout(null, 4, TimeUnit.SECONDS)
+                .completeOnTimeout(null, 15, TimeUnit.SECONDS)
                 .get();
 
         if (codebaseName == null) {

src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt

@@ -32,7 +32,7 @@ class CodyToolWindowContent(private val project: Project) {
 
   private var codyOnboardingGuidancePanel: CodyOnboardingGuidancePanel? = null
   private val signInWithSourcegraphPanel = SignInWithSourcegraphPanel(project)
-  private val historyTree = HistoryTree(::selectChat, ::removeChat, ::removeAllChats)
+  private val historyTree = HistoryTree(project, ::selectChat, ::removeChat, ::removeAllChats)
   private val tabbedPane = JBTabbedPane()
   private val currentChatSession: AtomicReference<AgentChatSession?> = AtomicReference(null)
 
@@ -137,7 +137,7 @@ class CodyToolWindowContent(private val project: Project) {
   }
 
   private fun removeChat(state: ChatState) {
-    HistoryService.getInstance().remove(state.internalId)
+    HistoryService.getInstance(project).remove(state.internalId)
     if (AgentChatSessionService.getInstance(project).removeSession(state)) {
       val isVisible = currentChatSession.get()?.getInternalId() == state.internalId
       if (isVisible) {
@@ -148,7 +148,7 @@ class CodyToolWindowContent(private val project: Project) {
 
   private fun removeAllChats() {
     AgentChatSessionService.getInstance(project).removeAllSessions()
-    HistoryService.getInstance().removeAll()
+    HistoryService.getInstance(project).removeAll()
     switchToChatSession(AgentChatSession.createNew(project))
   }
 

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

@@ -8,22 +8,28 @@ import com.intellij.openapi.vfs.VirtualFile
 import com.sourcegraph.cody.config.CodyProjectSettings
 import com.sourcegraph.config.ConfigUtil
 import com.sourcegraph.vcs.RepoUtil
+import java.util.concurrent.CompletableFuture
 
 @Service(Service.Level.PROJECT)
 class CodyAgentCodebase(val project: Project) {
 
   // TODO: Support list of repository names instead of just one.
-  private val application = ApplicationManager.getApplication()
   private val settings = CodyProjectSettings.getInstance(project)
-  private var inferredUrl: String? = null
+  private var inferredUrl: CompletableFuture<String> = CompletableFuture()
 
-  fun getUrl(): String? = settings.remoteUrl ?: inferredUrl
+  init {
+    onFileOpened(project, null)
+  }
+
+  fun getUrl(): CompletableFuture<String> =
+      if (settings.remoteUrl != null) CompletableFuture.completedFuture(settings.remoteUrl)
+      else inferredUrl
 
   fun onFileOpened(project: Project, file: VirtualFile?) {
-    application.executeOnPooledThread {
+    ApplicationManager.getApplication().executeOnPooledThread {
       val repositoryName = RepoUtil.findRepositoryName(project, file)
-      if (repositoryName != null && inferredUrl != repositoryName) {
-        inferredUrl = repositoryName
+      if (repositoryName != null && inferredUrl.getNow(null) != repositoryName) {
+        inferredUrl.complete(repositoryName)
         CodyAgentService.applyAgentOnBackgroundThread(project) {
           it.server.configurationDidChange(ConfigUtil.getAgentConfiguration(project))
         }

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

@@ -1,6 +1,7 @@
 package com.sourcegraph.cody.agent
 
 import com.sourcegraph.cody.agent.protocol.*
+import com.sourcegraph.cody.agent.protocol.util.ChatRemoteReposResponse
 import java.util.concurrent.CompletableFuture
 import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
 import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
@@ -35,6 +36,9 @@ interface CodyAgentServer {
   @JsonRequest("graphql/getRepoId")
   fun getRepoId(repoName: GetRepoIDResponse): CompletableFuture<String?>
 
+  @JsonRequest("graphql/getRepoIds")
+  fun getRepoIds(repoName: GetRepoIdsParam): CompletableFuture<GetRepoIdsResponse>
+
   @JsonRequest("git/codebaseName")
   fun convertGitCloneURLToCodebaseName(cloneURL: CloneURL): CompletableFuture<String?>
 
@@ -105,4 +109,6 @@ interface CodyAgentServer {
   fun attributionSearch(
       params: AttributionSearchParams
   ): CompletableFuture<AttributionSearchResponse>
+
+  @JsonRequest("chat/remoteRepos") fun chatRemoteRepos(): CompletableFuture<ChatRemoteReposResponse>
 }

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

@@ -0,0 +1,3 @@
+package com.sourcegraph.cody.agent.protocol
+
+data class GetRepoIdsParam(val names: List<String>, val first: Int)

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

@@ -0,0 +1,3 @@
+package com.sourcegraph.cody.agent.protocol
+
+data class GetRepoIdsResponse(val repos: List<Repo>)

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

@@ -0,0 +1,3 @@
+package com.sourcegraph.cody.agent.protocol
+
+data class Repo(val name: String, val id: String)

src/main/kotlin/com/sourcegraph/cody/agent/protocol/util/ChatRemoteReposResponse.kt

@@ -0,0 +1,5 @@
+package com.sourcegraph.cody.agent.protocol.util
+
+import com.sourcegraph.cody.agent.protocol.Repo
+
+data class ChatRemoteReposResponse(val remoteRepos: List<Repo>)

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

@@ -179,7 +179,7 @@ private constructor(
     }
     messages.add(message)
     chatPanel.addOrUpdateMessage(message)
-    HistoryService.getInstance().update(internalId, messages)
+    HistoryService.getInstance(project).updateChatMessages(internalId, messages)
   }
 
   @RequiresEdt
@@ -220,12 +220,7 @@ private constructor(
       val chatSession = AgentChatSession(project, sessionId)
 
       chatSession.createCancellationToken(
-          onCancel = {
-            CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
-              agent.server.webviewReceiveMessage(
-                  WebviewReceiveMessageParams(sessionId.get(), WebviewMessage(command = "abort")))
-            }
-          },
+          onCancel = { chatSession.sendWebviewMessage(WebviewMessage(command = "abort")) },
           onFinish = {
             GraphQlLogger.logCodyEvent(project, "command:${commandId.displayName}", "executed")
           })

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

@@ -1,11 +1,16 @@
 package com.sourcegraph.cody.chat
 
+import com.intellij.openapi.project.Project
 import com.intellij.ui.components.JBTextArea
 import com.sourcegraph.cody.history.HistoryService
 import com.sourcegraph.cody.history.state.MessageState
 import java.util.*
 
-class CodyChatMessageHistory(private val capacity: Int, chatSession: ChatSession) {
+class CodyChatMessageHistory(
+    private val project: Project,
+    private val capacity: Int,
+    chatSession: ChatSession
+) {
   var currentValue: String = ""
   private var upperStack: Stack<String> = Stack<String>()
   private var lowerStack: Stack<String> = Stack<String>()
@@ -53,8 +58,8 @@ class CodyChatMessageHistory(private val capacity: Int, chatSession: ChatSession
   }
 
   private fun preloadHistoricalMessages(chatSession: ChatSession) {
-    HistoryService.getInstance()
-        .state
+    HistoryService.getInstance(project)
+        .getHistoryReadOnly()
         .chats
         .find { it.internalId == chatSession.getInternalId() }
         ?.messages

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

@@ -18,12 +18,12 @@ import javax.swing.BorderFactory
 import javax.swing.JButton
 import javax.swing.JPanel
 
-class ChatPanel(project: Project, val chatSession: ChatSession) :
+class ChatPanel(project: Project, chatSession: ChatSession) :
     JPanel(VerticalFlowLayout(VerticalFlowLayout.CENTER, 0, 0, true, false)) {
   private val messagesPanel = MessagesPanel(project, chatSession)
   private val chatPanel = ChatScrollPane(messagesPanel)
-  val promptPanel: PromptPanel = PromptPanel(chatSession)
-  private val contextView: EnhancedContextPanel = EnhancedContextPanel(project)
+  val promptPanel: PromptPanel = PromptPanel(project, chatSession)
+  private val contextView: EnhancedContextPanel = EnhancedContextPanel(project, chatSession)
 
   private val stopGeneratingButton =
       object : JButton("Stop generating", IconUtil.desaturate(AllIcons.Actions.Suspend)) {

src/main/kotlin/com/sourcegraph/cody/config/CodyAccountDetailsProvider.kt

@@ -22,12 +22,15 @@ class CodyAccountDetailsProvider(
       account: CodyAccount,
       indicator: ProgressIndicator
   ): CompletableFuture<DetailsLoadingResult<CodyAccountDetails>> {
-    val token =
-        accountsModel.newCredentials.getOrElse(account) { accountManager.findCredentials(account) }
-            ?: return CompletableFuture.completedFuture(noToken())
-    val executor = service<SourcegraphApiRequestExecutor.Factory>().create(token)
+
     return ProgressManager.getInstance()
         .submitIOTask(indicator) { indicator ->
+          val token =
+              accountsModel.newCredentials.getOrElse(account) {
+                accountManager.findCredentials(account)
+              } ?: return@submitIOTask noToken()
+          val executor = service<SourcegraphApiRequestExecutor.Factory>().create(token)
+
           if (account.isCodyApp()) {
             val details = CodyAccountDetails(account.id, account.name, account.name, null)
             DetailsLoadingResult(details, IconUtil.toBufferedImage(defaultIcon), null, false)

src/main/kotlin/com/sourcegraph/cody/config/DialogValidationUtils.kt

@@ -2,6 +2,7 @@ package com.sourcegraph.cody.config
 
 import com.intellij.openapi.ui.ValidationInfo
 import com.intellij.openapi.util.NlsContexts
+import javax.swing.JComponent
 import javax.swing.JTextField
 
 object DialogValidationUtils {
@@ -12,11 +13,11 @@ object DialogValidationUtils {
 
   /** Returns [ValidationInfo] with [message] if [isValid] returns false */
   fun custom(
-      textField: JTextField,
+      component: JComponent,
       @NlsContexts.DialogMessage message: String,
       isValid: () -> Boolean
   ): ValidationInfo? {
-    return if (!isValid()) ValidationInfo(message, textField) else null
+    return if (!isValid()) ValidationInfo(message, component) else null
   }
 
   /**

src/main/kotlin/com/sourcegraph/cody/context/RemoteRepoUtils.kt

@@ -0,0 +1,23 @@
+package com.sourcegraph.cody.context
+
+import com.intellij.openapi.project.Project
+import com.sourcegraph.cody.agent.CodyAgentService
+import com.sourcegraph.cody.agent.protocol.GetRepoIdsParam
+import com.sourcegraph.cody.agent.protocol.Repo
+import java.util.concurrent.CompletableFuture
+
+object RemoteRepoUtils {
+  fun getRepository(project: Project, url: String): CompletableFuture<Repo?> {
+    val result = CompletableFuture<List<Repo>>()
+    CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
+      try {
+        agent.server.getRepoIds(GetRepoIdsParam(listOf(url), 1)).thenApply {
+          result.complete(it?.repos)
+        }
+      } catch (e: Exception) {
+        result.complete(null)
+      }
+    }
+    return result.thenApply { it?.firstOrNull() }
+  }
+}

src/main/kotlin/com/sourcegraph/cody/context/ui/AddRepositoryDialog.kt

@@ -0,0 +1,94 @@
+package com.sourcegraph.cody.context.ui
+
+import com.intellij.openapi.editor.event.DocumentListener
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.DialogWrapper
+import com.intellij.openapi.ui.ValidationInfo
+import com.intellij.ui.TextFieldWithAutoCompletion
+import com.intellij.ui.TextFieldWithAutoCompletionListProvider
+import com.intellij.util.Alarm
+import com.sourcegraph.cody.config.DialogValidationUtils
+import com.sourcegraph.cody.context.RemoteRepoUtils
+import java.net.URL
+import java.util.concurrent.TimeUnit
+import javax.swing.JComponent
+import javax.swing.JLabel
+import javax.swing.JPanel
+import org.jetbrains.annotations.NotNull
+
+class AddRepositoryDialog(private val project: Project, private val addAction: (String) -> Unit) :
+    DialogWrapper(project) {
+
+  private val repoUrlInputField = TextFieldWithAutoCompletion.create(project, listOf(), false, null)
+
+  init {
+    init()
+    title = "Add Remote Repository"
+    setOKButtonText("Add")
+    setValidationDelay(100)
+  }
+
+  override fun doValidateAll(): List<ValidationInfo> {
+    fun validateNonEmpty() =
+        DialogValidationUtils.custom(repoUrlInputField, "Remote repository URL cannot be empty") {
+          repoUrlInputField.text.isNotBlank()
+        }
+
+    fun validateValidUrl() =
+        DialogValidationUtils.custom(repoUrlInputField, "Remote repository URL must be valid") {
+          val url =
+              if (repoUrlInputField.text.startsWith("http")) repoUrlInputField.text
+              else "http://" + repoUrlInputField.text
+          runCatching { URL(url) }.isSuccess
+        }
+
+    fun validateRepoExists() =
+        DialogValidationUtils.custom(
+            repoUrlInputField, "Remote repository not found on the server") {
+              val repo =
+                  RemoteRepoUtils.getRepository(project, repoUrlInputField.text)
+                      .completeOnTimeout(null, 2, TimeUnit.SECONDS)
+                      .get()
+              repo != null
+            }
+
+    return listOfNotNull(validateNonEmpty() ?: validateValidUrl() ?: validateRepoExists())
+  }
+
+  override fun getValidationThreadToUse(): Alarm.ThreadToUse {
+    return Alarm.ThreadToUse.POOLED_THREAD
+  }
+
+  override fun doOKAction() {
+    addAction(repoUrlInputField.text)
+    close(OK_EXIT_CODE, true)
+  }
+
+  override fun createCenterPanel(): JComponent {
+    val panel = JPanel()
+    val label = JLabel("Repository URL: ")
+    panel.add(label)
+
+    // TODO: we can provide repository suggestions using `provider.setItems` method
+    val completionProvider: TextFieldWithAutoCompletionListProvider<String> =
+        object : TextFieldWithAutoCompletionListProvider<String>(listOf()) {
+          @NotNull
+          override fun getLookupString(@NotNull s: String): String {
+            return s
+          }
+        }
+
+    repoUrlInputField.setPreferredWidth(300)
+    repoUrlInputField.installProvider(completionProvider)
+    repoUrlInputField.addDocumentListener(
+        object : DocumentListener {
+          override fun documentChanged(event: com.intellij.openapi.editor.event.DocumentEvent) {
+            initValidation()
+          }
+        })
+
+    panel.add((repoUrlInputField))
+
+    return panel
+  }
+}