Add multi-repo enterprise context (#510)

## Test plan

1. Open `sourcegraph/cody` project with non-enterprise account.
- Open new chat and ask question about current repo (e.g. some class) -
assistant should know the answer
- Open new chat and ask question about squirrel - assistant should
describe you an animal
- Open new chat and disable local context. Ask about current repo (e.g.
some class) - assistant should not have a context
- Save current context as default. Close the IDE. Reopen the IDE. 
- Go to Chat History tab and open previous chats one by one. Check if
both history and context settings are properly preserved.
- Open new chat and check if context is disabled. It should be, as we
previously set that as new default. Enable it again and set as default.
  
2. Open `sourcegraph/cody` project with enterprise account.
- All checks from point 1) apply there as well.
- Click [+] button in the context panel and type sourcegraph repo url
(`github.com/sourcegraph/sourcegraph`)
  - Check if validator blocks from entering incomplete or invalid URL
  - Add the `sourcegraph/sourcegraph` repo by hitting Add button
- Open new chat and ask question about squirrel - assistant should
describe you an HTTP server, **NOT** animal.
- Set new configuration as default and open new chat. It should use
previous context configuration.
- Disable `sourcegraph/sourcegraph` remote repo context.
- Ask question about squirrel. It should be back to describing you an
animal or having no context.
- Save current context as default. Close the IDE. Reopen the IDE. 
- Go to Chat History tab and open previous chats one by one. Check if
both history and context settings are properly preserved.
- Open new chat and check if `sourcegraph/sourcegraph` is disabled. It
should be, as we previously set that as new default. Enable it again and
set as default.
- Open new chat and check if `sourcegraph/sourcegraph` is enabled, it
should be.
- Remove `sourcegraph/sourcegraph` repo by clicking on it to select it,
and then clicking [-] button.
- Ask question about squirrel. It should be back to describing you an
animal or having no context.

## Changes

* Added multi-repo enterprise context
* Improved look of the local context panel
* Fixed several bugs like:
* history service being instantiated as application-wide service instead
of project-wide
  * current repo url parsing and timing issues
 

![image](https://github.com/sourcegraph/jetbrains/assets/1519649/91a014db-6a64-4418-aee6-96400c25374f)

TESTING.md

@@ -19,6 +19,9 @@
     - [ ] [Read chat history without interruptions](#read-chat-history-without-interruptions)
     - [ ] [Organize multiple chats](#organize-multiple-chats)
     - [ ] [Isolate multiple chats](#isolate-multiple-chats)
+- Multi-repo context
+    - [] [Free/pro accounts:](#freepro-accounts)
+    - [] [Enterprise accounts:](#enterprise-accounts)
 - Sourcegraph Code Search
     - [ ] [Find with Sourcegraph...](#find-with-sourcegraph)
     - [ ] [Search Selection on Sourcegraph Web](#search-selection-on-sourcegraph-web)
@@ -52,7 +55,8 @@ Verify the remaining SSO methods by performing the same steps for `Sign in with
 
 ### Remove all accounts
 
-Prerequisite: You have to be **signed in**. This is important because we expect certain components to be refreshed automatically.
+Prerequisite: You have to be **signed in**. This is important because we expect certain components to be refreshed
+automatically.
 
 1. Navigate to `Settings` > `Sourcegraph & Cody`.
 2. Remove all accounts and apply settings.
@@ -130,9 +134,9 @@ Prerequisite: You have to be **signed in**. This is important because we expect
 
 ### General commands availability from keyboard shortcuts
 
-| Command       | Windows / Linux                                      | MacOs                                                |
-|---------------|------------------------------------------------------|------------------------------------------------------|
-| Explain Code  | <kbd>Alt</kbd> + <kbd>Shift</kbd> + <kbd>1</kbd>     | <kbd>control</kbd> + <kbd>Shift</kbd> + <kbd>E</kbd> |
+| Command       | Windows / Linux                                  | MacOs                                                |
+|---------------|--------------------------------------------------|------------------------------------------------------|
+| Explain Code  | <kbd>Alt</kbd> + <kbd>Shift</kbd> + <kbd>1</kbd> | <kbd>control</kbd> + <kbd>Shift</kbd> + <kbd>E</kbd> |
 | Smell Code    | <kbd>Alt</kbd> + <kbd>Shift</kbd> + <kbd>2</kbd> | <kbd>control</kbd> + <kbd>Shift</kbd> + <kbd>S</kbd> |
 | Generate Test | <kbd>Alt</kbd> + <kbd>Shift</kbd> + <kbd>3</kbd> | <kbd>control</kbd> + <kbd>Shift</kbd> + <kbd>T</kbd> |
 
@@ -248,43 +252,86 @@ Useful tips:
 
 Test ideas:
 
-1. Delete "active" chat. You should be able to delete the currently opened chat. Messages should be removed from Chat tab.
+1. Delete "active" chat. You should be able to delete the currently opened chat. Messages should be removed from Chat
+   tab.
 2. Restore historical chat, focus on chat input field and use UP/DOWN keys to cycle between previous questions.
 3. Press "new chat" as fast as you can. Especially during the IDE startup.
 4. Switch between chats as fast as you can.
-5. Press "new chat" while being inside `My Account` tab or something other than Chat tab. Tabs should switch automatically.
+5. Press "new chat" while being inside `My Account` tab or something other than Chat tab. Tabs should switch
+   automatically.
 6. Use commands/recipes inside empty, new chat. Verify serialization/deserialization.
-7. Ask about codebase to force response with listed context files and verify if everything is correctly serialized/deserialized. Links to context files should be clickable.
-8. Remove all chats using history UI. Tree presentation is empty and branches like "Today" are removed from panel. File with transcripts should also disappear.
+7. Ask about codebase to force response with listed context files and verify if everything is correctly
+   serialized/deserialized. Links to context files should be clickable.
+8. Remove all chats using history UI. Tree presentation is empty and branches like "Today" are removed from panel. File
+   with transcripts should also disappear.
 9. Use only the keyboard. For example, navigate transcripts with arrows, delete, enter.
 10. Start typing while being focused on Chat History to perform search-by-title.
 11. Open multiple chats and ask few simultaneous questions in several sessions at once.
 12. Open new chat with <kbd>Alt</kbd> + <kbd>=</kbd> shortcut (or <kbd>Option</kbd> + <kbd>=</kbd> on Mac).
-13. Open existing chat with shortcut <kbd>Alt</kbd> + <kbd>-</kbd> (or <kbd>Option</kbd> + <kbd>-</kbd> on Mac) and start typing question. Tab should be switched automatically on Chat.
-14. Second click <kbd>Alt</kbd> + <kbd-</kbd> should hide tool window if focused (similar behavior as other tool windows).
-15. Click <kbd>Esc</kbd> while being focused inside Cody tool window. You should be automatically focused on code.  
+13. Open existing chat with shortcut <kbd>Alt</kbd> + <kbd>-</kbd> (or <kbd>Option</kbd> + <kbd>-</kbd> on Mac) and
+    start typing question. Tab should be switched automatically on Chat.
+14. Second click <kbd>Alt</kbd> + <kbd-</kbd> should hide tool window if focused (similar behavior as other tool
+    windows).
+15. Click <kbd>Esc</kbd> while being focused inside Cody tool window. You should be automatically focused on code.
 
 #### Isolate multiple chats
 
 Prerequisite: You need two working accounts. Preferably one Free, and one Enterprise.
 
 1. Switch to first account.
 2. Send a message.
-3. Switch to second account. 
+3. Switch to second account.
 4. Send a message.
 
 These two chats should be isolated between different accounts. Both accounts should have one conversation each.
 
 You should also be able to switch between accounts while tokens are still being generated.
 
+## Multi-repo context
+
+### Free/pro accounts:
+
+1. Open `sourcegraph/cody` project with non-enterprise account.
+2. Open new chat and ask question about current repo (e.g. some class) - assistant should know the answer.
+3. Open new chat and ask question about squirrel - assistant should describe you an animal.
+4. Open new chat and disable local context. Ask about current repo (e.g. some class) - assistant should not have a
+   context.
+5. Save current context as default. Close the IDE. Reopen the IDE.
+    - Go to Chat History tab and open previous chats one by one. Both history and context settings are properly
+      preserved.
+    - Open new chat. Context should be disabled, as we previously set that as new default. Enable it again and set as
+      default.
+
+### Enterprise accounts:
+
+1. Open `sourcegraph/cody` project with enterprise account.
+2. Re-do all check from `Testing free/pro accounts` section but now with enterprise account.
+3. Click [+] button in the context panel and type sourcegraph repo url (`github.com/sourcegraph/sourcegraph`)
+    - Validator should block accepting incomplete or invalid URL.
+    - Add the `sourcegraph/sourcegraph` repo by hitting Add button.
+4. Open new chat and ask question about squirrel - assistant should describe you an HTTP server, **NOT** animal.
+5. Set current context as default and open new chat. It should use/display default context configuration.
+6. Disable `sourcegraph/sourcegraph` remote repo context.
+7. Ask question about squirrel. It should again describe you an animal or have no context.
+8. Save current context as default. Close the IDE. Reopen the IDE.
+    - Go to Chat History tab and open previous chats one by one. Check if both history and context settings are properly
+      preserved.
+    - Open new chat and check if `sourcegraph/sourcegraph` is disabled. It should be, as we previously set that as a
+      default. Enable it again and set as default.
+    - Open new chat and check if `sourcegraph/sourcegraph` is enabled, it should be.
+    - Remove `sourcegraph/sourcegraph` repo by selecting it, and then clicking [-] button.
+    - Ask question about squirrel. It should again describe you an animal or have no context.
+
 ## Code Search
 
-All `Code Search` actions are available under the same `Sourcegraph` right-click context menu, so for simplicity, we describe only the **Expected behaviours**.
+All `Code Search` actions are available under the same `Sourcegraph` right-click context menu, so for simplicity, we
+describe only the **Expected behaviours**.
 
 To open the context menu:
 
 1. Open a file in the repository that is indexed by Sourcegraph.
-2. Select the fragment of code you want to search for (for example: `System.out.println` or `println` may be the simplest candidate).
+2. Select the fragment of code you want to search for (for example: `System.out.println` or `println` may be the
+   simplest candidate).
 3. Right-click on selected fragment, navigate to the `Sourcegraph` sub-menu and choose one of the actions.
 
 ### Find with Sourcegraph...
@@ -303,7 +350,8 @@ To open the context menu:
 #### Expected behaviour:
 
 1. The browser is launched.
-2. The result is a list of fragments that are found **within the same repository** from which the searched fragment originates.
+2. The result is a list of fragments that are found **within the same repository** from which the searched fragment
+   originates.
 
 ### Open Selection on Sourcegraph Web
 
@@ -319,7 +367,8 @@ To open the context menu:
 
 1. A link is copied to the clipboard.
 2. Notification pops up with successful message.
-3. After pasting the link into the browser, the Code Search page opens with the file and the exact line from which the searched fragment originates.
+3. After pasting the link into the browser, the Code Search page opens with the file and the exact line from which the
+   searched fragment originates.
 
 ## [Product-led growth](https://handbook.sourcegraph.com/departments/data-analytics/product-led-growth/)
 
@@ -355,7 +404,8 @@ To open the context menu:
 
 1. Open project with enabled Git VCS. This repository must be publicly available on GitHub.
 2. Open to `Cody` tool window.
-3. Click on repository button to open `Context Selection` dialog. Button is placed inside `Cody` tool window on left, bottom
+3. Click on repository button to open `Context Selection` dialog. Button is placed inside `Cody` tool window on left,
+   bottom
    corner.
 
 #### Expected behaviour

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

@@ -104,7 +104,7 @@ class CodyToolWindowContent(private val project: Project) {
     }
   }
 
-  @RequiresEdt fun refreshHistoryTree() = historyTree.rebuildTree(project)
+  @RequiresEdt fun refreshHistoryTree() = historyTree.rebuildTree()
 
   @RequiresEdt
   fun refreshPanelsVisibility() {
@@ -144,7 +144,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) {
@@ -155,7 +155,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(project, 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

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