fix: Missing context files

src/main/java/com/sourcegraph/cody/attribution/AttributionSearchCommand.kt

@@ -6,8 +6,6 @@ import com.sourcegraph.cody.agent.CodyAgentService
 import com.sourcegraph.cody.agent.CurrentConfigFeatures
 import com.sourcegraph.cody.agent.protocol.AttributionSearchParams
 import com.sourcegraph.cody.agent.protocol.AttributionSearchResponse
-import com.sourcegraph.cody.chat.AgentChatSession
-import com.sourcegraph.cody.chat.AgentChatSessionService
 import com.sourcegraph.cody.chat.SessionId
 import com.sourcegraph.cody.chat.ui.CodeEditorPart
 import java.util.*
@@ -51,7 +49,4 @@ class AttributionSearchCommand(private val project: Project) {
 
   private fun attributionEnabled(): Boolean =
       project.getService(CurrentConfigFeatures::class.java).get().attribution
-
-  private fun findChatSessionFor(messageId: UUID): AgentChatSession? =
-      AgentChatSessionService.getInstance(project).findByMessage(messageId)
 }

src/main/java/com/sourcegraph/cody/vscode/Completion.java

@@ -1,15 +1,15 @@
 package com.sourcegraph.cody.vscode;
 
-import com.sourcegraph.cody.agent.protocol.Message;
+import com.sourcegraph.cody.agent.protocol.ChatMessage;
 import java.util.List;
 
 public class Completion {
   public final String prefix;
-  public final List<Message> messages;
+  public final List<ChatMessage> messages;
   public final String content;
   public final String stopReason;
 
-  public Completion(String prefix, List<Message> messages, String content, String stopReason) {
+  public Completion(String prefix, List<ChatMessage> messages, String content, String stopReason) {
     this.prefix = prefix;
     this.messages = messages;
     this.content = content;
@@ -20,7 +20,7 @@ public Completion withPrefix(String newPrefix) {
     return new Completion(newPrefix, this.messages, this.content, this.stopReason);
   }
 
-  public Completion withMessages(List<Message> newMessages) {
+  public Completion withMessages(List<ChatMessage> newMessages) {
     return new Completion(this.prefix, newMessages, this.content, this.stopReason);
   }
 

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

@@ -186,7 +186,8 @@ private constructor(
                 .serializeNulls()
                 .registerTypeAdapter(CompletionItemID::class.java, CompletionItemIDSerializer)
                 .registerTypeAdapter(ContextFile::class.java, contextFileDeserializer)
-                .registerTypeAdapter(Speaker::class.java, SpeakerSerializer)
+                .registerTypeAdapter(Speaker::class.java, speakerDeserializer)
+                .registerTypeAdapter(Speaker::class.java, speakerSerializer)
                 .registerTypeAdapter(URI::class.java, uriDeserializer)
                 .registerTypeAdapter(URI::class.java, uriSerializer)
           }

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

@@ -37,10 +37,6 @@ class CodyAgentService(project: Project) : Disposable {
               ?.receiveWebviewExtensionMessage(params.message)
         }
       }
-
-      if (!project.isDisposed) {
-        AgentChatSessionService.getInstance(project).restoreAllSessions(agent)
-      }
     }
   }
 

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

@@ -1,7 +1,6 @@
 package com.sourcegraph.cody.agent.protocol
 
 import java.time.OffsetDateTime
-import java.util.*
 
 data class ChatError(
     val kind: String? = null,
@@ -27,14 +26,11 @@ data class ChatError(
 }
 
 data class ChatMessage(
-    override val speaker: Speaker,
-    override val text: String?,
+    val speaker: Speaker,
+    val text: String?,
     val displayText: String? = null,
     val contextFiles: List<ContextFile>? = null,
-    val error: ChatError? = null,
-    // Internal ID used for identifying updates of the message
-    // All partial messages which are part of the same response are required to have the same ID
-    val id: UUID = UUID.randomUUID()
-) : Message {
+    val error: ChatError? = null
+) {
   fun actualMessage(): String = displayText ?: text ?: ""
 }

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

@@ -1,3 +1,7 @@
 package com.sourcegraph.cody.agent.protocol
 
-data class ChatRestoreParams(val modelID: String, val messages: List<Message>, val chatID: String)
+data class ChatRestoreParams(
+    val modelID: String,
+    val messages: List<ChatMessage>,
+    val chatID: String
+)

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

@@ -4,7 +4,7 @@ import com.google.gson.*
 import java.lang.reflect.Type
 import java.net.URI
 
-sealed class ContextFile() {
+sealed class ContextFile {
   abstract val type: String
   abstract val uri: URI
   abstract val repoName: String?
@@ -68,14 +68,3 @@ val uriSerializer = JsonSerializer { uri: URI, type, context ->
     obj
   }
 }
-
-fun contextFilesFromList(list: List<Any>): List<String> {
-  val contextFiles = ArrayList<String>()
-  for (item in list) {
-    if (item is Map<*, *> && item.get("type") == "file") {
-      val path = (item.get("uri") as Map<*, *>).get("path")
-      contextFiles.add(path as String)
-    }
-  }
-  return contextFiles
-}

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

@@ -1,5 +1,6 @@
 package com.sourcegraph.cody.agent.protocol
 
+import com.google.gson.JsonDeserializer
 import com.google.gson.JsonPrimitive
 import com.google.gson.JsonSerializer
 import com.google.gson.annotations.SerializedName
@@ -9,4 +10,8 @@ enum class Speaker(val speaker: String) {
   @SerializedName("assistant") ASSISTANT("assistant")
 }
 
-val SpeakerSerializer = JsonSerializer<Speaker> { src, _, _ -> JsonPrimitive(src.speaker) }
+val speakerSerializer = JsonSerializer { speaker: Speaker, _, _ -> JsonPrimitive(speaker.speaker) }
+
+val speakerDeserializer = JsonDeserializer { src, _, _ ->
+  Speaker.values().find { it.speaker == src.asString }
+}

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

@@ -3,6 +3,7 @@ package com.sourcegraph.cody.chat
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.project.Project
 import com.intellij.util.concurrency.annotations.RequiresEdt
+import com.intellij.util.ui.UIUtil
 import com.intellij.xml.util.XmlStringUtil
 import com.jetbrains.rd.util.AtomicReference
 import com.sourcegraph.cody.agent.*
@@ -15,8 +16,7 @@ import com.sourcegraph.cody.history.state.ChatState
 import com.sourcegraph.cody.history.state.MessageState
 import com.sourcegraph.cody.vscode.CancellationToken
 import com.sourcegraph.common.CodyBundle
-import com.sourcegraph.common.CodyBundle.fmt
-import com.sourcegraph.common.UpgradeToCodyProNotification
+import com.sourcegraph.common.UpgradeToCodyProNotification.Companion.isCodyProJetbrains
 import com.sourcegraph.telemetry.GraphQlLogger
 import java.util.*
 import java.util.concurrent.CompletableFuture
@@ -27,7 +27,7 @@ class AgentChatSession
 private constructor(
     private val project: Project,
     newSessionId: CompletableFuture<SessionId>,
-    private val internalId: String = UUID.randomUUID().toString(),
+    private val internalId: String = UUID.randomUUID().toString()
 ) : ChatSession {
 
   /**
@@ -66,17 +66,29 @@ private constructor(
 
   fun hasSessionId(thatSessionId: SessionId): Boolean = getSessionId() == thatSessionId
 
-  fun hasMessageId(messageId: UUID): Boolean = messages.any { it.id == messageId }
-
   override fun getInternalId(): String = internalId
 
   override fun getCancellationToken(): CancellationToken = cancellationToken.get()
 
+  private val logger = LoggerFactory.getLogger(ChatSession::class.java)
+
   @RequiresEdt
   override fun sendMessage(text: String, contextFiles: List<ContextFile>) {
     val displayText = XmlStringUtil.escapeString(text)
-    val humanMessage = ChatMessage(Speaker.HUMAN, text, displayText)
-    addMessage(humanMessage)
+    val humanMessage =
+        ChatMessage(
+            Speaker.HUMAN,
+            text,
+            displayText,
+        )
+    addMessageAtIndex(humanMessage, index = messages.count())
+    val responsePlaceholder =
+        ChatMessage(
+            Speaker.ASSISTANT,
+            text = "",
+            displayText = "",
+        )
+    addMessageAtIndex(responsePlaceholder, index = messages.count(), shouldAddBlinkingCursor = true)
 
     CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
       val message =
@@ -102,34 +114,26 @@ private constructor(
 
   @Throws(ExecutionException::class, InterruptedException::class)
   override fun receiveMessage(extensionMessage: ExtensionMessage) {
-    fun addAssistantResponseToChat(text: String, displayText: String? = null) {
-      // Updates of the given message will always have the same UUID
-      val messageId =
-          UUID.nameUUIDFromBytes(extensionMessage.messages?.count().toString().toByteArray())
-      ApplicationManager.getApplication().invokeLater {
-        addMessage(ChatMessage(Speaker.ASSISTANT, text, displayText, id = messageId))
-      }
-    }
 
     try {
       val lastMessage = extensionMessage.messages?.lastOrNull()
+      val prevLastMessage = extensionMessage.messages?.dropLast(1)?.lastOrNull()
 
       if (lastMessage?.error != null && extensionMessage.isMessageInProgress == false) {
 
         getCancellationToken().dispose()
         val rateLimitError = lastMessage.error.toRateLimitError()
         if (rateLimitError != null) {
           RateLimitStateManager.reportForChat(project, rateLimitError)
-          UpgradeToCodyProNotification.isCodyProJetbrains(project).thenApply { isCodyPro ->
+          isCodyProJetbrains(project).thenApply { isCodyPro ->
             val text =
                 when {
                   rateLimitError.upgradeIsAvailable && isCodyPro ->
                       CodyBundle.getString("chat.rate-limit-error.upgrade")
-                          .fmt(rateLimitError.limit?.let { " $it" } ?: "")
                   else -> CodyBundle.getString("chat.rate-limit-error.explain")
                 }
 
-            addAssistantResponseToChat(text)
+            addErrorMessageAsAssistantMessage(text, index = extensionMessage.messages.count() - 1)
           }
         } else {
           // Currently we ignore other kind of errors like context window limit reached
@@ -140,15 +144,35 @@ private constructor(
             extensionMessage.isMessageInProgress == false) {
           getCancellationToken().dispose()
         } else {
-          if (lastMessage?.text != null && extensionMessage.chatID != null) {
-            addAssistantResponseToChat(lastMessage.text, lastMessage.displayText)
+
+          if (extensionMessage.chatID != null) {
+            if (prevLastMessage != null) {
+              if (lastMessage?.contextFiles != messages.lastOrNull()?.contextFiles) {
+                val index = extensionMessage.messages.count() - 2
+                ApplicationManager.getApplication().invokeLater {
+                  addMessageAtIndex(prevLastMessage, index)
+                }
+              }
+            }
+
+            if (lastMessage?.text != null) {
+              val index = extensionMessage.messages.count() - 1
+              ApplicationManager.getApplication().invokeLater {
+                addMessageAtIndex(lastMessage, index)
+              }
+            }
           }
         }
       }
     } catch (error: Exception) {
-      getCancellationToken().dispose()
+      getCancellationToken().abort()
       logger.error(CodyBundle.getString("chat-session.error-title"), error)
-      addAssistantResponseToChat(CodyBundle.getString("chat-session.error-title"))
+    }
+  }
+
+  private fun addErrorMessageAsAssistantMessage(stringMessage: String, index: Int) {
+    UIUtil.invokeLaterIfNeeded {
+      addMessageAtIndex(ChatMessage(Speaker.ASSISTANT, stringMessage), index)
     }
   }
 
@@ -173,12 +197,21 @@ private constructor(
   }
 
   @RequiresEdt
-  private fun addMessage(message: ChatMessage) {
-    if (messages.lastOrNull()?.id == message.id) {
-      messages.removeLast()
+  private fun addMessageAtIndex(
+      message: ChatMessage,
+      index: Int,
+      shouldAddBlinkingCursor: Boolean? = null
+  ) {
+    val messageToUpdate = messages.getOrNull(index)
+    if (messageToUpdate != null) {
+      messages[index] = message
+    } else {
+      messages.add(message)
     }
-    messages.add(message)
-    chatPanel.addOrUpdateMessage(message)
+    chatPanel.addOrUpdateMessage(
+        message,
+        index,
+        shouldAddBlinkingCursor = shouldAddBlinkingCursor ?: message.actualMessage().isBlank())
     HistoryService.getInstance(project).updateChatMessages(internalId, messages)
   }
 
@@ -225,7 +258,19 @@ private constructor(
             GraphQlLogger.logCodyEvent(project, "command:${commandId.displayName}", "executed")
           })
 
-      chatSession.addMessage(ChatMessage(Speaker.HUMAN, commandId.displayName))
+      chatSession.addMessageAtIndex(
+          ChatMessage(
+              Speaker.HUMAN,
+              commandId.displayName,
+          ),
+          chatSession.messages.count())
+      chatSession.addMessageAtIndex(
+          ChatMessage(
+              Speaker.ASSISTANT,
+              text = "",
+              displayText = "",
+          ),
+          chatSession.messages.count())
       AgentChatSessionService.getInstance(project).addSession(chatSession)
       return chatSession
     }
@@ -234,16 +279,17 @@ private constructor(
     fun createFromState(project: Project, state: ChatState): AgentChatSession {
       val sessionId = createNewPanel(project) { it.server.chatNew() }
       val chatSession = AgentChatSession(project, sessionId, state.internalId!!)
-      for (message in state.messages) {
+      state.messages.forEachIndexed { index, message ->
         val parsed =
             when (val speaker = message.speaker) {
               MessageState.SpeakerState.HUMAN -> Speaker.HUMAN
               MessageState.SpeakerState.ASSISTANT -> Speaker.ASSISTANT
               else -> error("unrecognized speaker $speaker")
             }
-        val chatMessage = ChatMessage(parsed, message.text)
+        val chatMessage = ChatMessage(speaker = parsed, message.text)
         chatSession.messages.add(chatMessage)
-        chatSession.chatPanel.addOrUpdateMessage(chatMessage, shouldAddBlinkingCursor = false)
+        chatSession.chatPanel.addOrUpdateMessage(
+            chatMessage, index, shouldAddBlinkingCursor = false)
       }
       CodyAgentService.applyAgentOnBackgroundThread(project) { agent ->
         chatSession.restoreAgentSession(agent)

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

@@ -6,7 +6,6 @@ import com.intellij.openapi.project.Project
 import com.intellij.util.concurrency.annotations.RequiresEdt
 import com.sourcegraph.cody.agent.CodyAgent
 import com.sourcegraph.cody.history.state.ChatState
-import java.util.UUID
 import java.util.concurrent.ConcurrentLinkedQueue
 
 @Service(Service.Level.PROJECT)
@@ -35,9 +34,6 @@ class AgentChatSessionService(private val project: Project) {
   fun getSession(sessionId: SessionId): AgentChatSession? =
       chatSessions.find { it.hasSessionId(sessionId) }
 
-  fun findByMessage(messageId: UUID): AgentChatSession? =
-      chatSessions.find { it.hasMessageId(messageId) }
-
   fun restoreAllSessions(agent: CodyAgent) {
     chatSessions.forEach { it.restoreAgentSession(agent) }
   }