Switch to chatID-agnostic chat/models API

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

@@ -1,3 +1,10 @@
 package com.sourcegraph.cody.agent.protocol
 
-data class ChatModelsParams(val id: String)
+import com.google.gson.annotations.SerializedName
+
+data class ChatModelsParams(val modelUsage: ModelUsage)
+
+enum class ModelUsage {
+  @SerializedName("chat") chat,
+  @SerializedName("edit") edit
+}

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

@@ -4,7 +4,6 @@ import com.sourcegraph.cody.Icons
 import javax.swing.Icon
 
 data class ChatModelsResponse(val models: List<ChatModelProvider>) {
-
   data class ChatModelProvider(
       val default: Boolean,
       val codyProOnly: Boolean,

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

@@ -15,7 +15,6 @@ import com.sourcegraph.cody.agent.WebviewMessage
 import com.sourcegraph.cody.agent.WebviewReceiveMessageParams
 import com.sourcegraph.cody.agent.protocol.ChatError
 import com.sourcegraph.cody.agent.protocol.ChatMessage
-import com.sourcegraph.cody.agent.protocol.ChatModelsParams
 import com.sourcegraph.cody.agent.protocol.ChatModelsResponse
 import com.sourcegraph.cody.agent.protocol.ChatRestoreParams
 import com.sourcegraph.cody.agent.protocol.ChatSubmitMessageParams
@@ -59,7 +58,6 @@ private constructor(
 
   init {
     cancellationToken.get().dispose()
-    updateLlmDropdownModels()
   }
 
   fun getPanel(): ChatPanel = chatPanel
@@ -243,23 +241,21 @@ private constructor(
     chatPanel.registerCancellationToken(newCancellationToken)
   }
 
-  private fun updateLlmDropdownModels() {
-    if (chatModelProviderFromState != null) {
-      return
-    }
-
-    CodyAgentService.withAgent(project) { agent ->
-      val chatModels = agent.server.chatModels(ChatModelsParams(connectionId.get().get()))
-      val response =
-          chatModels.completeOnTimeout(null, 10, TimeUnit.SECONDS).get() ?: return@withAgent
-      chatPanel.updateLlmDropdownModels(response.models)
-    }
-  }
-
   fun updateFromState(agent: CodyAgent, state: ChatState) {
-    val newConnectionId = restoreChatSession(agent, state)
+    val chatMessages =
+        state.messages.map { message ->
+          val parsed =
+              when (val speaker = message.speaker) {
+                MessageState.SpeakerState.HUMAN -> Speaker.HUMAN
+                MessageState.SpeakerState.ASSISTANT -> Speaker.ASSISTANT
+                else -> error("unrecognized speaker $speaker")
+              }
+
+          ChatMessage(speaker = parsed, message.text)
+        }
+    val newConnectionId =
+        restoreChatSession(agent, chatMessages, chatModelProviderFromState, state.internalId!!)
     connectionId.getAndSet(newConnectionId)
-    updateLlmDropdownModels()
   }
 
   companion object {
@@ -317,29 +313,10 @@ private constructor(
 
     fun restoreChatSession(
         agent: CodyAgent,
-        chatState: ChatState
+        chatMessages: List<ChatMessage>,
+        chatModelProvider: ChatModelsResponse.ChatModelProvider?,
+        internalId: String
     ): CompletableFuture<ConnectionId> {
-      val modelFromState =
-          chatState.llm?.let {
-            ChatModelsResponse.ChatModelProvider(
-                default = it.model == null,
-                codyProOnly = false,
-                provider = it.provider,
-                title = it.title,
-                model = it.model ?: "")
-          }
-
-      val chatMessages =
-          chatState.messages.map { message ->
-            val parsed =
-                when (val speaker = message.speaker) {
-                  MessageState.SpeakerState.HUMAN -> Speaker.HUMAN
-                  MessageState.SpeakerState.ASSISTANT -> Speaker.ASSISTANT
-                  else -> error("unrecognized speaker $speaker")
-                }
-
-            ChatMessage(speaker = parsed, message.text)
-          }
 
       val messages =
           chatMessages
@@ -348,16 +325,27 @@ private constructor(
                 if (acc.lastOrNull()?.speaker == msg.speaker) acc else acc.plus(msg)
               }
 
-      val restoreParams = ChatRestoreParams(modelFromState?.model, messages, chatState.internalId!!)
+      val restoreParams = ChatRestoreParams(chatModelProvider?.model, messages, internalId)
       return agent.server.chatRestore(restoreParams)
     }
 
     fun createFromState(project: Project, state: ChatState): AgentChatSession {
-
       val agentChatSession = CompletableFuture<AgentChatSession>()
+
+      val chatModelProvider =
+          state.llm?.let {
+            ChatModelsResponse.ChatModelProvider(
+                default = it.model == null,
+                codyProOnly = false,
+                provider = it.provider,
+                title = it.title,
+                model = it.model ?: "")
+          }
+
       createNewPanel(project) { codyAgent ->
         val dummyConnectionId = codyAgent.server.chatNew()
-        val chatSession = AgentChatSession(project, dummyConnectionId, state.internalId!!)
+        val chatSession =
+            AgentChatSession(project, dummyConnectionId, state.internalId!!, chatModelProvider)
 
         chatSession.updateFromState(codyAgent, state)
         AgentChatSessionService.getInstance(project).addSession(chatSession)

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

@@ -1,7 +1,6 @@
 package com.sourcegraph.cody.chat.ui
 
 import com.intellij.icons.AllIcons
-import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.VerticalFlowLayout
 import com.intellij.util.IconUtil
@@ -10,6 +9,7 @@ import com.sourcegraph.cody.PromptPanel
 import com.sourcegraph.cody.agent.WebviewMessage
 import com.sourcegraph.cody.agent.protocol.ChatMessage
 import com.sourcegraph.cody.agent.protocol.ChatModelsResponse
+import com.sourcegraph.cody.agent.protocol.ModelUsage
 import com.sourcegraph.cody.chat.ChatSession
 import com.sourcegraph.cody.config.CodyAuthenticationManager
 import com.sourcegraph.cody.context.ui.EnhancedContextPanel
@@ -31,7 +31,11 @@ class ChatPanel(
 
   val promptPanel: PromptPanel = PromptPanel(project, chatSession)
   private val llmDropdown =
-      LlmDropdown(project, onSetSelectedItem = ::setLlmForAgentSession, chatModelProviderFromState)
+      LlmDropdown(
+          ModelUsage.chat,
+          project,
+          onSetSelectedItem = ::setLlmForAgentSession,
+          chatModelProviderFromState)
   private val messagesPanel = MessagesPanel(project, chatSession)
   private val chatPanel = ChatScrollPane(messagesPanel)
 
@@ -81,15 +85,6 @@ class ChatPanel(
     messagesPanel.addOrUpdateMessage(message, index)
   }
 
-  @RequiresEdt
-  fun addAllMessages(messages: List<ChatMessage>) {
-    if (messages.isNotEmpty()) {
-      llmDropdown.updateAfterFirstMessage()
-      promptPanel.updateEmptyTextAfterFirstMessage()
-    }
-    messages.forEach(messagesPanel::addChatMessageAsComponent)
-  }
-
   @RequiresEdt
   fun registerCancellationToken(cancellationToken: CancellationToken) {
     messagesPanel.registerCancellationToken(cancellationToken)
@@ -104,10 +99,6 @@ class ChatPanel(
     stopGeneratingButton.addActionListener { cancellationToken.abort() }
   }
 
-  fun updateLlmDropdownModels(models: List<ChatModelsResponse.ChatModelProvider>) {
-    ApplicationManager.getApplication().invokeLater { llmDropdown.updateModels(models) }
-  }
-
   private fun setLlmForAgentSession(chatModelProvider: ChatModelsResponse.ChatModelProvider) {
     val activeAccountType = CodyAuthenticationManager.getInstance(project).getActiveAccount()
     if (activeAccountType?.isEnterpriseAccount() == true) {

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

@@ -1,17 +1,23 @@
 package com.sourcegraph.cody.chat.ui
 
+import com.intellij.openapi.application.invokeLater
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.ComboBox
 import com.intellij.ui.MutableCollectionComboBoxModel
 import com.intellij.util.concurrency.annotations.RequiresEdt
+import com.sourcegraph.cody.agent.CodyAgentService
+import com.sourcegraph.cody.agent.protocol.ChatModelsParams
 import com.sourcegraph.cody.agent.protocol.ChatModelsResponse
+import com.sourcegraph.cody.agent.protocol.ModelUsage
 import com.sourcegraph.cody.config.AccountTier
 import com.sourcegraph.cody.config.CodyAuthenticationManager
 import com.sourcegraph.cody.ui.LlmComboBoxRenderer
 import com.sourcegraph.common.BrowserOpener
 import com.sourcegraph.common.CodyBundle
+import java.util.concurrent.TimeUnit
 
 class LlmDropdown(
+    private val modelUsage: ModelUsage,
     private val project: Project,
     private val onSetSelectedItem: (ChatModelsResponse.ChatModelProvider) -> Unit,
     private val chatModelProviderFromState: ChatModelsResponse.ChatModelProvider?,
@@ -27,17 +33,28 @@ class LlmDropdown(
     }
 
     isEnabled = false
+    updateModels()
   }
 
-  @RequiresEdt
-  fun updateModels(models: List<ChatModelsResponse.ChatModelProvider>) {
+  private fun updateModels() {
     if (chatModelProviderFromState != null) {
       return
     }
 
+    CodyAgentService.withAgent(project) { agent ->
+      val chatModels = agent.server.chatModels(ChatModelsParams(modelUsage))
+      val response =
+          chatModels.completeOnTimeout(null, 10, TimeUnit.SECONDS).get() ?: return@withAgent
+
+      invokeLater { updateModelsInUI(response) }
+    }
+  }
+
+  @RequiresEdt
+  private fun updateModelsInUI(response: ChatModelsResponse) {
     removeAllItems()
-    models.forEach(::addItem)
-    models.find { it.default }?.let { this.selectedItem = it }
+    response.models.forEach(::addItem)
+    response.models.find { it.default }?.let { this.selectedItem = it }
 
     CodyAuthenticationManager.getInstance(project).getActiveAccountTier().thenApply { accountTier ->
       isCurrentUserFree = accountTier == AccountTier.DOTCOM_FREE