From a22119f738fbbb9930c66308d64f7c35b4b15a72 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sun, 9 Jun 2024 18:42:17 +0000
Subject: [PATCH] Better handling of newly saved conversations

When a new conversation is created in the chat view's
`persistConversation`, we now set `rootMessages` appropriately.
---
 pkg/tui/views/chat/chat.go         | 1 +
 pkg/tui/views/chat/conversation.go | 6 +++---
 pkg/tui/views/chat/update.go       | 3 +++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/pkg/tui/views/chat/chat.go b/pkg/tui/views/chat/chat.go
index cbace1e..f18a577 100644
--- a/pkg/tui/views/chat/chat.go
+++ b/pkg/tui/views/chat/chat.go
@@ -35,6 +35,7 @@ type (
 	msgMessagesLoaded []models.Message
 	// sent when the conversation has been persisted, triggers a reload of contents
 	msgConversationPersisted struct {
+		isNew        bool
 		conversation *models.Conversation
 		messages     []models.Message
 	}
diff --git a/pkg/tui/views/chat/conversation.go b/pkg/tui/views/chat/conversation.go
index f53e717..deade01 100644
--- a/pkg/tui/views/chat/conversation.go
+++ b/pkg/tui/views/chat/conversation.go
@@ -203,14 +203,14 @@ func (m *Model) persistConversation() tea.Cmd {
 	messages := m.messages
 
 	var err error
-	if m.conversation.ID == 0 {
+	if conversation.ID == 0 {
 		return func() tea.Msg {
 			// Start a new conversation with all messages so far
 			conversation, messages, err = m.Shared.Ctx.Store.StartConversation(messages...)
 			if err != nil {
 				return shared.MsgError(fmt.Errorf("Could not start new conversation: %v", err))
 			}
-			return msgConversationPersisted{conversation, messages}
+			return msgConversationPersisted{true, conversation, messages}
 		}
 	}
 
@@ -239,7 +239,7 @@ func (m *Model) persistConversation() tea.Cmd {
 				return fmt.Errorf("Error: no messages to reply to")
 			}
 		}
-		return msgConversationPersisted{conversation, messages}
+		return msgConversationPersisted{false, conversation, messages}
 	}
 }
 
diff --git a/pkg/tui/views/chat/update.go b/pkg/tui/views/chat/update.go
index 19ee29f..6d40ef6 100644
--- a/pkg/tui/views/chat/update.go
+++ b/pkg/tui/views/chat/update.go
@@ -166,6 +166,9 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
 	case msgConversationPersisted:
 		m.conversation = msg.conversation
 		m.messages = msg.messages
+		if msg.isNew {
+			m.rootMessages = []models.Message{m.messages[0]}
+		}
 		m.rebuildMessageCache()
 		m.updateContent()
 	case msgMessageCloned: