Isolate chats between accounts (#503)

Fixes #374 ## Test Plan 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. 4. Send a message. These two chats should be isolated between different accounts. Both accounts should have one conversation each. It doesn't matter if you're on FREE or Enterprise: every chat now is assigned to account. You should also be able to switch between accounts while tokens are still being generated. Note about migration: Old chats will preserved in the new schema. This works so that orphaned chats will be assigned to the currently active account.
sourcegraph · Feb 7, 2024 · 1204f1c · 1204f1c
1 parent 982f53d
commit 1204f1c
Show file tree

Hide file tree

Showing 9 changed files with 82 additions and 37 deletions.
diff --git a/TESTING.md b/TESTING.md
@@ -18,6 +18,7 @@
     - [ ] [Autoscroll to latest message](#autoscroll-to-latest-message)
     - [ ] [Read chat history without interruptions](#read-chat-history-without-interruptions)
     - [ ] [Organize multiple chats](#organize-multiple-chats)
+    - [ ] [Isolate multiple chats](#isolate-multiple-chats)
 - Sourcegraph Code Search
     - [ ] [Find with Sourcegraph...](#find-with-sourcegraph)
     - [ ] [Search Selection on Sourcegraph Web](#search-selection-on-sourcegraph-web)
@@ -263,6 +264,19 @@ Test ideas:
 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. 
+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.
+
 ## 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**.

diff --git a/src/main/kotlin/com/sourcegraph/cody/CodyToolWindowContent.kt b/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)
 
@@ -104,6 +104,8 @@ class CodyToolWindowContent(private val project: Project) {
     }
   }
 
+  @RequiresEdt fun refreshHistoryTree() = historyTree.rebuildTree(project)
+
   @RequiresEdt
   fun refreshPanelsVisibility() {
     val codyAuthenticationManager = CodyAuthenticationManager.instance

diff --git a/src/main/kotlin/com/sourcegraph/cody/chat/AgentChatSession.kt b/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().update(project, internalId, messages)
   }
 
   @RequiresEdt

diff --git a/src/main/kotlin/com/sourcegraph/cody/config/SettingsMigration.kt b/src/main/kotlin/com/sourcegraph/cody/config/SettingsMigration.kt
@@ -11,46 +11,47 @@ import com.intellij.openapi.progress.EmptyProgressIndicator
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.wm.ToolWindowManager
+import com.sourcegraph.cody.CodyToolWindowContent
 import com.sourcegraph.cody.CodyToolWindowFactory
 import com.sourcegraph.cody.api.SourcegraphApiRequestExecutor
 import com.sourcegraph.cody.api.SourcegraphApiRequests
+import com.sourcegraph.cody.history.HistoryService
 import com.sourcegraph.cody.initialization.Activity
-import com.sourcegraph.config.AccessTokenStorage
-import com.sourcegraph.config.CodyApplicationService
-import com.sourcegraph.config.CodyProjectService
-import com.sourcegraph.config.ConfigUtil
-import com.sourcegraph.config.UserLevelConfig
+import com.sourcegraph.config.*
 import java.util.concurrent.CompletableFuture
 
-private const val RUN_ONCE_CODY_ACCOUNTS_IDS_REFRESH = "CodyAccountsIdsRefresh"
-private const val RUN_ONCE_TOGGLE_CODY_TOOL_WINDOW_AFTER_MIGRATION =
-    "ToggleCodyToolWindowAfterMigration"
-private const val RUN_ONCE_CODY_APPLICATION_SETTINGS_MIGRATION = "CodyApplicationSettingsMigration"
-private const val RUN_ONCE_CODY_PROJECT_SETTINGS_MIGRATION = "CodyProjectSettingsMigration"
-
 class SettingsMigration : Activity {
 
   private val codyAuthenticationManager = CodyAuthenticationManager.instance
 
   override fun runActivity(project: Project) {
-    RunOnceUtil.runOnceForProject(project, RUN_ONCE_CODY_PROJECT_SETTINGS_MIGRATION) {
-      val customRequestHeaders = extractCustomRequestHeaders(project)
+    RunOnceUtil.runOnceForProject(project, "CodyProjectSettingsMigration") {
       migrateProjectSettings(project)
-      migrateAccounts(project, customRequestHeaders)
-    }
-    RunOnceUtil.runOnceForApp(RUN_ONCE_CODY_APPLICATION_SETTINGS_MIGRATION) {
-      migrateApplicationSettings()
+      migrateAccounts(project)
     }
-    RunOnceUtil.runOnceForApp(RUN_ONCE_TOGGLE_CODY_TOOL_WINDOW_AFTER_MIGRATION) {
-      ApplicationManager.getApplication().invokeLater { toggleCodyToolbarWindow(project) }
+    RunOnceUtil.runOnceForApp("CodyApplicationSettingsMigration") { migrateApplicationSettings() }
+    RunOnceUtil.runOnceForApp("ToggleCodyToolWindowAfterMigration") {
+      toggleCodyToolbarWindow(project)
     }
-    RunOnceUtil.runOnceForApp(RUN_ONCE_CODY_ACCOUNTS_IDS_REFRESH) {
-      val customRequestHeaders = extractCustomRequestHeaders(project)
-      refreshAccountsIds(customRequestHeaders)
+    RunOnceUtil.runOnceForApp("CodyAccountsIdsRefresh") { refreshAccountsIds(project) }
+    RunOnceUtil.runOnceForApp("CodyAssignOrphanedChatsToActiveAccount") {
+      migrateOrphanedChatsToActiveAccount(project)
     }
   }
 
-  private fun refreshAccountsIds(customRequestHeaders: String) {
+  private fun migrateOrphanedChatsToActiveAccount(project: Project) {
+    val activeAccountId = CodyAuthenticationManager.instance.getActiveAccount(project)?.id
+    HistoryService.getInstance()
+        .state
+        .chats
+        .filter { it.accountId == null }
+        .forEach { it.accountId = activeAccountId }
+    // required because this activity is executed later than tool window creation
+    CodyToolWindowContent.executeOnInstanceIfNotDisposed(project) { refreshHistoryTree() }
+  }
+
+  private fun refreshAccountsIds(project: Project) {
+    val customRequestHeaders = extractCustomRequestHeaders(project)
     codyAuthenticationManager.getAccounts().forEach { codyAccount ->
       val server = SourcegraphServerPath.from(codyAccount.server.url, customRequestHeaders)
       val token = codyAuthenticationManager.getTokenForAccount(codyAccount)
@@ -68,12 +69,15 @@ class SettingsMigration : Activity {
   }
 
   private fun toggleCodyToolbarWindow(project: Project) {
-    val toolWindowManager = ToolWindowManager.getInstance(project)
-    val toolWindow = toolWindowManager.getToolWindow(CodyToolWindowFactory.TOOL_WINDOW_ID)
-    toolWindow?.setAvailable(CodyApplicationSettings.instance.isCodyEnabled, null)
+    ApplicationManager.getApplication().invokeLater {
+      val toolWindowManager = ToolWindowManager.getInstance(project)
+      val toolWindow = toolWindowManager.getToolWindow(CodyToolWindowFactory.TOOL_WINDOW_ID)
+      toolWindow?.setAvailable(CodyApplicationSettings.instance.isCodyEnabled, null)
+    }
   }
 
-  private fun migrateAccounts(project: Project, customRequestHeaders: String) {
+  private fun migrateAccounts(project: Project) {
+    val customRequestHeaders = extractCustomRequestHeaders(project)
     val requestExecutorFactory = SourcegraphApiRequestExecutor.Factory.instance
     migrateDotcomAccount(project, requestExecutorFactory, customRequestHeaders)
     migrateEnterpriseAccount(project, requestExecutorFactory, customRequestHeaders)

diff --git a/src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt b/src/main/kotlin/com/sourcegraph/cody/config/notification/AccountSettingChangeListener.kt
@@ -40,6 +40,7 @@ class AccountSettingChangeListener(project: Project) : ChangeListener(project) {
               CodyToolWindowContent.executeOnInstanceIfNotDisposed(project) {
                 refreshPanelsVisibility()
                 refreshMyAccountTab()
+                refreshHistoryTree()
               }
             }
 

diff --git a/src/main/kotlin/com/sourcegraph/cody/history/HistoryService.kt b/src/main/kotlin/com/sourcegraph/cody/history/HistoryService.kt
@@ -1,8 +1,10 @@
 package com.sourcegraph.cody.history
 
 import com.intellij.openapi.components.*
+import com.intellij.openapi.project.Project
 import com.sourcegraph.cody.agent.protocol.ChatMessage
 import com.sourcegraph.cody.agent.protocol.Speaker
+import com.sourcegraph.cody.config.CodyAuthenticationManager
 import com.sourcegraph.cody.history.state.ChatState
 import com.sourcegraph.cody.history.state.HistoryState
 import com.sourcegraph.cody.history.state.MessageState
@@ -18,8 +20,8 @@ class HistoryService : SimplePersistentStateComponent<HistoryState>(HistoryState
     synchronized(listeners) { listeners += listener }
   }
 
-  fun update(internalId: String, chatMessages: List<ChatMessage>) {
-    val found = getChatOrCreate(internalId)
+  fun update(project: Project, internalId: String, chatMessages: List<ChatMessage>) {
+    val found = getChatOrCreate(project, internalId)
     found.messages = chatMessages.map(::convertToMessageState).toMutableList()
     if (chatMessages.lastOrNull()?.speaker == Speaker.HUMAN) {
       found.setUpdatedTimeAt(LocalDateTime.now())
@@ -46,10 +48,11 @@ class HistoryService : SimplePersistentStateComponent<HistoryState>(HistoryState
     return message
   }
 
-  private fun getChatOrCreate(internalId: String): ChatState {
+  private fun getChatOrCreate(project: Project, internalId: String): ChatState {
     val found = state.chats.find { it.internalId == internalId }
     if (found != null) return found
-    val newChat = ChatState().also { it.internalId = internalId }
+    val activeAccountId = CodyAuthenticationManager.instance.getActiveAccount(project)?.id
+    val newChat = ChatState.create(activeAccountId, internalId)
     state.chats += newChat
     return newChat
   }

diff --git a/src/main/kotlin/com/sourcegraph/cody/history/HistoryTree.kt b/src/main/kotlin/com/sourcegraph/cody/history/HistoryTree.kt
@@ -4,11 +4,13 @@ import com.intellij.icons.AllIcons
 import com.intellij.openapi.actionSystem.AnAction
 import com.intellij.openapi.actionSystem.AnActionEvent
 import com.intellij.openapi.actionSystem.DefaultActionGroup
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.SimpleToolWindowPanel
 import com.intellij.ui.PopupHandler
 import com.intellij.ui.ScrollPaneFactory
 import com.intellij.ui.treeStructure.SimpleTree
 import com.intellij.util.EditSourceOnDoubleClickHandler
+import com.sourcegraph.cody.config.CodyAuthenticationManager
 import com.sourcegraph.cody.history.node.LeafNode
 import com.sourcegraph.cody.history.node.PeriodNode
 import com.sourcegraph.cody.history.node.RootNode
@@ -27,12 +29,13 @@ import javax.swing.tree.TreePath
 import javax.swing.tree.TreeSelectionModel
 
 class HistoryTree(
+    project: Project,
     private val onSelect: (ChatState) -> Unit,
     private val onRemove: (ChatState) -> Unit,
     private val onRemoveAll: () -> Unit
 ) : SimpleToolWindowPanel(true, true) {
 
-  private val model = DefaultTreeModel(buildTree())
+  private val model = DefaultTreeModel(buildTree(project))
   private val root
     get() = model.root as RootNode
 
@@ -72,6 +75,10 @@ class HistoryTree(
     HistoryService.getInstance().listenOnUpdate(::updatePresentation)
   }
 
+  fun rebuildTree(project: Project) {
+    model.setRoot(buildTree(project))
+  }
+
   private fun updatePresentation(chat: ChatState) {
     val leafNotInTree =
         root.periods().flatMap { it.leafs() }.none { it.chat.internalId == chat.internalId }
@@ -157,9 +164,9 @@ class HistoryTree(
     model.setRoot(RootNode())
   }
 
-  private fun buildTree(): DefaultMutableTreeNode {
+  private fun buildTree(project: Project): DefaultMutableTreeNode {
     val root = RootNode()
-    for ((period, chats) in getChatsGroupedByPeriod()) {
+    for ((period, chats) in getChatsGroupedByPeriod(project)) {
       val periodNode = PeriodNode(period)
       for (chat in chats) {
         periodNode.add(LeafNode(chat))
@@ -169,10 +176,13 @@ class HistoryTree(
     return root
   }
 
-  private fun getChatsGroupedByPeriod(): Map<String, List<ChatState>> =
+  private fun getChatsGroupedByPeriod(project: Project): Map<String, List<ChatState>> =
       HistoryService.getInstance()
           .state
           .chats
+          .filter {
+            it.accountId == CodyAuthenticationManager.instance.getActiveAccount(project)?.id
+          }
           .sortedByDescending { chat -> chat.getUpdatedTimeAt() }
           .groupBy { chat -> DurationGroupFormatter.format(chat.getUpdatedTimeAt()) }
 

diff --git a/src/main/kotlin/com/sourcegraph/cody/history/state/ChatState.kt b/src/main/kotlin/com/sourcegraph/cody/history/state/ChatState.kt
@@ -16,6 +16,8 @@ class ChatState : BaseState() {
 
   @get:OptionTag(tag = "updatedAt", nameAttribute = "") var updatedAt: String? by string()
 
+  @get:OptionTag(tag = "accountId", nameAttribute = "") var accountId: String? by string()
+
   fun title(): String? = messages.firstOrNull()?.text
 
   fun setUpdatedTimeAt(date: LocalDateTime) {
@@ -30,5 +32,12 @@ class ChatState : BaseState() {
   companion object {
 
     private val DATE_FORMAT = DateTimeFormatter.ISO_LOCAL_DATE_TIME
+
+    fun create(accountId: String?, internalId: String): ChatState {
+      val chat = ChatState()
+      chat.accountId = accountId
+      chat.internalId = internalId
+      return chat
+    }
   }
 }
diff --git a/src/test/kotlin/com/sourcegraph/cody/history/state/HistoryStateTest.kt b/src/test/kotlin/com/sourcegraph/cody/history/state/HistoryStateTest.kt
@@ -14,6 +14,7 @@ class HistoryStateTest : TestCase() {
         HistoryState().apply {
           chats +=
               ChatState().apply {
+                accountId = "VXNlkjoxEFU3NjE="
                 internalId = "0f8b7034-9fa8-488a-a13e-09c52677008a"
                 updatedAt = "2024-01-31T01:06:18.524621"
                 messages +=
@@ -37,6 +38,7 @@ class HistoryStateTest : TestCase() {
         <chats>
           <list>
             <chat>
+              <accountId value="VXNlkjoxEFU3NjE=" />
               <internalId value="0f8b7034-9fa8-488a-a13e-09c52677008a" />
               <messages>
                 <list>