Add ability to cycle through conversation branches in tui

pkg/tui/chat.go

@@ -59,6 +59,7 @@ type chatModel struct {
 
 	// app state
 	conversation    *models.Conversation
+	rootMessages    []models.Message
 	messages        []models.Message
 	selectedMessage int
 	waitingForReply bool
@@ -275,6 +276,7 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		}
 	case msgConversationLoaded:
 		m.conversation = (*models.Conversation)(msg)
+		m.rootMessages, _ = m.ctx.Store.RootMessages(m.conversation.ID)
 		cmds = append(cmds, m.loadMessages(m.conversation))
 	case msgMessagesLoaded:
 		m.selectedMessage = len(msg) - 1
@@ -464,6 +466,40 @@ func (m *chatModel) handleMessagesKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 			scrollIntoView(&m.content, offset, m.content.Height/2)
 		}
 		return true, nil
+	case "ctrl+h", "ctrl+l":
+		dir := CyclePrev
+		if msg.String() == "ctrl+l" {
+			dir = CycleNext
+		}
+
+		var err error
+		var selected *models.Message
+		if m.selectedMessage == 0 {
+			selected, err = m.cycleSelectedRoot(m.conversation, dir)
+			if err != nil {
+				return true, wrapError(fmt.Errorf("Could not cycle conversation root: %v", err))
+			}
+		} else if m.selectedMessage > 0 {
+			selected, err = m.cycleSelectedReply(&m.messages[m.selectedMessage-1], dir)
+			if err != nil {
+				return true, wrapError(fmt.Errorf("Could not cycle reply: %v", err))
+			}
+		}
+
+		if selected == nil {
+			return false, nil
+		}
+
+		// Retrieve updated view at this point
+		newPath, err := m.ctx.Store.PathToLeaf(selected)
+		if err != nil {
+			m.err = fmt.Errorf("Could not fetch messages: %v", err)
+		}
+
+		m.messages = append(m.messages[:m.selectedMessage], newPath...)
+		m.rebuildMessageCache()
+		m.updateContent()
+		return true, nil
 	case "ctrl+r":
 		// resubmit the conversation with all messages up until and including the selected message
 		if m.waitingForReply || len(m.messages) == 0 {
@@ -479,6 +515,73 @@ func (m *chatModel) handleMessagesKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 	return false, nil
 }
 
+type CycleDirection int
+
+const (
+	CycleNext CycleDirection = 1
+	CyclePrev CycleDirection = -1
+)
+
+func cycleMessages(m *models.Message, msgs []models.Message, dir CycleDirection) (*models.Message, error) {
+	currentIndex := -1
+	for i, reply := range msgs {
+		if reply.ID == m.ID {
+			currentIndex = i
+			break
+		}
+	}
+
+	if currentIndex < 0 {
+		return nil, fmt.Errorf("message not found")
+	}
+
+	var next int
+	if dir == CyclePrev {
+		// Wrap around to the last reply if at the beginning
+		next = (currentIndex - 1 + len(msgs)) % len(msgs)
+	} else {
+		// Wrap around to the first reply if at the end
+		next = (currentIndex + 1) % len(msgs)
+	}
+	return &msgs[next], nil
+}
+
+func (m *chatModel) cycleSelectedRoot(conv *models.Conversation, dir CycleDirection) (*models.Message, error) {
+	if len(m.rootMessages) < 2 {
+		return nil, nil
+	}
+
+	nextRoot, err := cycleMessages(conv.SelectedRoot, m.rootMessages, dir)
+	if err != nil {
+		return nil, err
+	}
+
+	conv.SelectedRoot = nextRoot
+	err = m.ctx.Store.UpdateConversation(conv)
+	if err != nil {
+		return nil, fmt.Errorf("Could not update conversation: %v", err)
+	}
+	return nextRoot, nil
+}
+
+func (m *chatModel) cycleSelectedReply(message *models.Message, dir CycleDirection) (*models.Message, error) {
+	if len(message.Replies) < 2 {
+		return nil, nil
+	}
+
+	nextReply, err := cycleMessages(message.SelectedReply, message.Replies, dir)
+	if err != nil {
+		return nil, err
+	}
+
+	message.SelectedReply = nextReply
+	err = m.ctx.Store.UpdateMessage(message)
+	if err != nil {
+		return nil, fmt.Errorf("Could not update message: %v", err)
+	}
+	return nextReply, nil
+}
+
 func (m *chatModel) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 	switch msg.String() {
 	case "esc":
@@ -558,6 +661,29 @@ func (m *chatModel) renderMessageHeading(i int, message *models.Message) string
 	var suffix string
 
 	faint := lipgloss.NewStyle().Faint(true)
+
+	if i == 0 && len(m.rootMessages) > 0 {
+		selectedRootIndex := 0
+		for j, reply := range m.rootMessages {
+			if reply.ID == *m.conversation.SelectedRootID {
+				selectedRootIndex = j
+				break
+			}
+		}
+		suffix += faint.Render(fmt.Sprintf(" <%d/%d>", selectedRootIndex+1, len(m.rootMessages)))
+	}
+	if i > 0 && len(m.messages[i-1].Replies) > 1 {
+		// Find the selected reply index
+		selectedReplyIndex := 0
+		for j, reply := range m.messages[i-1].Replies {
+			if reply.ID == *m.messages[i-1].SelectedReplyID {
+				selectedReplyIndex = j
+				break
+			}
+		}
+		suffix += faint.Render(fmt.Sprintf(" <%d/%d>", selectedReplyIndex+1, len(m.messages[i-1].Replies)))
+	}
+
 	if m.focus == focusMessages {
 		if i == m.selectedMessage {
 			prefix = "> "
@@ -853,6 +979,8 @@ func (m *chatModel) persistConversation() error {
 			if err != nil {
 				return err
 			}
+			// add this message as a reply to the previous
+			m.messages[i-1].Replies = append(m.messages[i-1].Replies, saved[0])
 			m.messages[i] = saved[0]
 		} else {
 			// message has no id and no previous messages to add it to