Fix crash on Edit Code

src/main/kotlin/com/sourcegraph/cody/auth/deprecated/DeprecatedCodyAccount.kt

@@ -20,8 +20,6 @@ data class DeprecatedCodyAccount(
     @Attribute("id") var id: String = generateId(),
 ) {
 
-  fun isDotcomAccount(): Boolean = server.url.lowercase().startsWith(ConfigUtil.DOTCOM_URL)
-
   fun credentialAttributes(): CredentialAttributes =
       CredentialAttributes(generateServiceName("Sourcegraph", id))
 

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

@@ -1,7 +1,6 @@
 package com.sourcegraph.cody.chat.ui
 
 import com.intellij.openapi.application.invokeLater
-import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.ui.MutableCollectionComboBoxModel
@@ -15,9 +14,9 @@ import com.sourcegraph.cody.agent.protocol_extensions.isDeprecated
 import com.sourcegraph.cody.agent.protocol_generated.Chat_ModelsParams
 import com.sourcegraph.cody.agent.protocol_generated.Model
 import com.sourcegraph.cody.agent.protocol_generated.ModelAvailabilityStatus
+import com.sourcegraph.cody.auth.CodyAccount
+import com.sourcegraph.cody.auth.SourcegraphServerPath
 import com.sourcegraph.cody.edit.EditCommandPrompt
-import com.sourcegraph.cody.history.HistoryService
-import com.sourcegraph.cody.history.state.LLMState
 import com.sourcegraph.cody.ui.LlmComboBoxRenderer
 import com.sourcegraph.common.BrowserOpener
 import java.util.concurrent.TimeUnit
@@ -27,10 +26,8 @@ class LlmDropdown(
     private val project: Project,
     private val onSetSelectedItem: (Model) -> Unit,
     val parentDialog: EditCommandPrompt?,
-    private val chatModelFromState: Model?,
-    private val model: String? = null
+    private val fixedModel: String? = null
 ) : ComboBox<ModelAvailabilityStatus>(MutableCollectionComboBoxModel()) {
-  private var hasServerSentModels = project.service<CurrentConfigFeatures>().get().serverSentModels
 
   init {
     renderer = LlmComboBoxRenderer(this)
@@ -58,7 +55,6 @@ class LlmDropdown(
   }
 
   private fun handleConfigUpdate(config: ConfigFeatures) {
-    hasServerSentModels = config.serverSentModels
     if (!isVisible && config.serverSentModels) {
       isVisible = true
       revalidate()
@@ -73,19 +69,13 @@ class LlmDropdown(
     val availableModels = models.map { it }.filterNot { it.model.isDeprecated() }
     availableModels.sortedBy { it.model.isCodyProOnly() }.forEach { addItem(it) }
 
-    val selectedFromChatState = chatModelFromState
-    val selectedFromHistory = HistoryService.getInstance(project).getDefaultLlm()
+    val defaultLlm = serverToRecentModel[CodyAccount.getActiveAccount()?.server]
 
     selectedItem =
-        availableModels.find {
-          it.model.id == model ||
-              it.model.id == selectedFromHistory?.model ||
-              it.model.id == selectedFromChatState?.id
-        } ?: models.firstOrNull()
-
-    // If the dropdown is already disabled, don't change it. It can happen
-    // in the case of the legacy commands (updateAfterFirstMessage happens before this call).
-    isEnabled = isEnabled && chatModelFromState == null
+        availableModels.find { it.model.id == fixedModel || it.model.id == defaultLlm?.id }
+            ?: models.firstOrNull()
+
+    isEnabled = fixedModel == null
     isVisible = selectedItem != null
     setMaximumRowCount(15)
 
@@ -106,10 +96,14 @@ class LlmDropdown(
         return
       }
 
-      HistoryService.getInstance(project).setDefaultLlm(LLMState.fromChatModel(modelProvider.model))
+      CodyAccount.getActiveAccount()?.server?.also { serverToRecentModel[it] = modelProvider.model }
 
       super.setSelectedItem(anObject)
       onSetSelectedItem(modelProvider.model)
     }
   }
+
+  companion object {
+    private val serverToRecentModel = HashMap<SourcegraphServerPath, Model>()
+  }
 }

src/main/kotlin/com/sourcegraph/cody/edit/EditCommandPrompt.kt

@@ -135,8 +135,7 @@ class EditCommandPrompt(
               project = project,
               onSetSelectedItem = { model = it.id },
               this,
-              chatModelFromState = null,
-              model = model)
+              fixedModel = model)
           .apply {
             foreground = boldLabelColor()
             background = textFieldBackground()

src/main/kotlin/com/sourcegraph/cody/history/HistoryService.kt

@@ -6,60 +6,20 @@ import com.intellij.openapi.components.State
 import com.intellij.openapi.components.Storage
 import com.intellij.openapi.components.service
 import com.intellij.openapi.project.Project
-import com.sourcegraph.cody.auth.deprecated.DeprecatedCodyAccountManager
 import com.sourcegraph.cody.history.state.AccountData
 import com.sourcegraph.cody.history.state.ChatState
 import com.sourcegraph.cody.history.state.HistoryState
-import com.sourcegraph.cody.history.state.LLMState
 
 @State(name = "ChatHistory", storages = [Storage("cody_history.xml")])
 @Service(Service.Level.PROJECT)
-class HistoryService(private val project: Project) :
-    SimplePersistentStateComponent<HistoryState>(HistoryState()) {
-
-  @Synchronized
-  fun getDefaultLlm(): LLMState? {
-    val account = DeprecatedCodyAccountManager.getInstance().account
-    val llm = account?.let { findEntry(it.id) }?.defaultLlm
-    if (llm == null) return null
-    return LLMState().also { it.copyFrom(llm) }
-  }
-
-  @Synchronized
-  fun setDefaultLlm(defaultLlm: LLMState) {
-    val newDefaultLlm = LLMState()
-    newDefaultLlm.copyFrom(defaultLlm)
-    getOrCreateActiveAccountEntry().defaultLlm = newDefaultLlm
-  }
-
-  @Synchronized
-  fun remove(internalId: String?) {
-    getOrCreateActiveAccountEntry().chats.removeIf { it.internalId == internalId }
-  }
-
-  @Synchronized
-  fun findActiveAccountChat(internalId: String): ChatState? =
-      getActiveAccountHistory()?.chats?.find { it.internalId == internalId }
+class HistoryService : SimplePersistentStateComponent<HistoryState>(HistoryState()) {
 
   @Synchronized
   fun getChatHistoryFor(accountId: String): List<ChatState>? = findEntry(accountId)?.chats
 
   private fun findEntry(accountId: String): AccountData? =
       state.accountData.find { it.accountId == accountId }
 
-  @Synchronized
-  fun getActiveAccountHistory(): AccountData? =
-      DeprecatedCodyAccountManager.getInstance().account?.let { findEntry(it.id) }
-
-  private fun getOrCreateActiveAccountEntry(): AccountData {
-    val activeAccount =
-        DeprecatedCodyAccountManager.getInstance().account
-            ?: throw IllegalStateException("No active account")
-
-    val existingEntry = findEntry(activeAccount.id)
-    return existingEntry ?: AccountData(activeAccount.id).also { state.accountData += it }
-  }
-
   companion object {
     @JvmStatic fun getInstance(project: Project): HistoryService = project.service<HistoryService>()
   }

src/test/kotlin/com/sourcegraph/cody/config/SettingsMigrationTest.kt

@@ -182,7 +182,7 @@ class SettingsMigrationTest : BasePlatformTestCase() {
                   ChatState("chat4"))
         }
     val project = myFixture.project
-    project.registerServiceInstance(HistoryService::class.java, HistoryService(project))
+    project.registerServiceInstance(HistoryService::class.java, HistoryService())
     HistoryService.getInstance(project)
         .loadState(HistoryState().also { it.copyFrom(originalHistory) })