feat(TUI): improves UX with message navigation modal to jump and restore to specific messages (#1969)

2 files changed, 326 insertions, 10 deletions

diff --git a/packages/tui/internal/components/chat/messages.go b/packages/tui/internal/components/chat/messages.go
index 0be2f462c..1ee72b9e5 100644
--- a/packages/tui/internal/components/chat/messages.go
+++ b/packages/tui/internal/components/chat/messages.go
@@ -39,6 +39,7 @@ type MessagesComponent interface {
 	CopyLastMessage() (tea.Model, tea.Cmd)
 	UndoLastMessage() (tea.Model, tea.Cmd)
 	RedoLastMessage() (tea.Model, tea.Cmd)
+	ScrollToMessage(messageID string) (tea.Model, tea.Cmd)
 }
 
 type messagesComponent struct {
@@ -57,6 +58,7 @@ type messagesComponent struct {
 	partCount          int
 	lineCount          int
 	selection          *selection
+	messagePositions   map[string]int // map message ID to line position
 }
 
 type selection struct {
@@ -228,6 +230,7 @@ func (m *messagesComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.rendering = false
 		m.clipboard = msg.clipboard
 		m.loading = false
+		m.messagePositions = msg.messagePositions
 		m.tail = m.viewport.AtBottom()
 
 		// Preserve scroll across reflow
@@ -256,11 +259,12 @@ func (m *messagesComponent) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 }
 
 type renderCompleteMsg struct {
-	viewport  viewport.Model
-	clipboard []string
-	header    string
-	partCount int
-	lineCount int
+	viewport         viewport.Model
+	clipboard        []string
+	header           string
+	partCount        int
+	lineCount        int
+	messagePositions map[string]int
 }
 
 func (m *messagesComponent) renderView() tea.Cmd {
@@ -286,6 +290,7 @@ func (m *messagesComponent) renderView() tea.Cmd {
 		blocks := make([]string, 0)
 		partCount := 0
 		lineCount := 0
+		messagePositions := make(map[string]int) // Track message ID to line position
 
 		orphanedToolCalls := make([]opencode.ToolPart, 0)
 
@@ -308,6 +313,9 @@ func (m *messagesComponent) renderView() tea.Cmd {
 
 			switch casted := message.Info.(type) {
 			case opencode.UserMessage:
+				// Track the position of this user message
+				messagePositions[casted.ID] = lineCount
+
 				if casted.ID == m.app.Session.Revert.MessageID {
 					reverted = true
 					revertedMessageCount = 1
@@ -767,11 +775,12 @@ func (m *messagesComponent) renderView() tea.Cmd {
 		}
 
 		return renderCompleteMsg{
-			header:    header,
-			clipboard: clipboard,
-			viewport:  viewport,
-			partCount: partCount,
-			lineCount: lineCount,
+			header:           header,
+			clipboard:        clipboard,
+			viewport:         viewport,
+			partCount:        partCount,
+			lineCount:        lineCount,
+			messagePositions: messagePositions,
 		}
 	}
 }
@@ -1190,6 +1199,18 @@ func (m *messagesComponent) RedoLastMessage() (tea.Model, tea.Cmd) {
 	}
 }
 
+func (m *messagesComponent) ScrollToMessage(messageID string) (tea.Model, tea.Cmd) {
+	if m.messagePositions == nil {
+		return m, nil
+	}
+
+	if position, exists := m.messagePositions[messageID]; exists {
+		m.viewport.SetYOffset(position)
+		m.tail = false // Stop auto-scrolling to bottom when manually navigating
+	}
+	return m, nil
+}
+
 func NewMessagesComponent(app *app.App) MessagesComponent {
 	vp := viewport.New()
 	vp.KeyMap = viewport.KeyMap{}
@@ -1214,5 +1235,6 @@ func NewMessagesComponent(app *app.App) MessagesComponent {
 		showThinkingBlocks: showThinkingBlocks,
 		cache:              NewPartCache(),
 		tail:               true,
+		messagePositions:   make(map[string]int),
 	}
 }
diff --git a/packages/tui/internal/components/dialog/navigation.go b/packages/tui/internal/components/dialog/navigation.go
new file mode 100644
index 000000000..2f8bb5626
--- /dev/null
+++ b/packages/tui/internal/components/dialog/navigation.go
@@ -0,0 +1,294 @@
+package dialog
+
+import (
+	"fmt"
+	"strings"
+	"time"
+
+	tea "github.com/charmbracelet/bubbletea/v2"
+	"github.com/muesli/reflow/truncate"
+	"github.com/sst/opencode-sdk-go"
+	"github.com/sst/opencode/internal/app"
+	"github.com/sst/opencode/internal/components/list"
+	"github.com/sst/opencode/internal/components/modal"
+	"github.com/sst/opencode/internal/layout"
+	"github.com/sst/opencode/internal/styles"
+	"github.com/sst/opencode/internal/theme"
+	"github.com/sst/opencode/internal/util"
+)
+
+// NavigationDialog interface for the session navigation dialog
+type NavigationDialog interface {
+	layout.Modal
+}
+
+// ScrollToMessageMsg is sent when a message should be scrolled to
+type ScrollToMessageMsg struct {
+	MessageID string
+}
+
+// RestoreToMessageMsg is sent when conversation should be restored to a specific message
+type RestoreToMessageMsg struct {
+	MessageID string
+	Index     int
+}
+
+// navigationItem represents a user message in the navigation list
+type navigationItem struct {
+	messageID string
+	content   string
+	timestamp time.Time
+	index     int // Index in the full message list
+	toolCount int // Number of tools used in this message
+}
+
+func (n navigationItem) Render(
+	selected bool,
+	width int,
+	isFirstInViewport bool,
+	baseStyle styles.Style,
+) string {
+	t := theme.CurrentTheme()
+
+	// Format timestamp - only apply color when not selected
+	var timeStr string
+	var timeVisualLen int
+	if selected {
+		timeStr = n.timestamp.Format("15:04")
+		timeVisualLen = len(timeStr)
+	} else {
+		infoStyle := styles.NewStyle().Foreground(t.Info()).Render
+		timeStr = infoStyle(n.timestamp.Format("15:04"))
+		timeVisualLen = len(n.timestamp.Format("15:04")) // Visual length without color codes
+	}
+
+	// Tool count display (fixed width for alignment) - only apply color when not selected
+	toolInfo := ""
+	toolInfoVisualLen := 0
+	if n.toolCount > 0 {
+		toolInfoText := fmt.Sprintf("(%d tools)", n.toolCount)
+		if selected {
+			toolInfo = toolInfoText
+		} else {
+			infoStyle := styles.NewStyle().Foreground(t.Info()).Render
+			toolInfo = infoStyle(toolInfoText)
+		}
+		toolInfoVisualLen = len(toolInfoText) // Use the visual length, not the styled length
+	}
+
+	// Calculate available space for content
+	// Reserve space for: timestamp + space + toolInfo + padding + some buffer
+	reservedSpace := timeVisualLen + 1 + toolInfoVisualLen + 4
+	contentWidth := width - reservedSpace
+	if contentWidth < 8 {
+		contentWidth = 8
+	}
+
+	truncatedContent := truncate.StringWithTail(n.content, uint(contentWidth), "...")
+
+	// Apply normal text color to content for non-selected items
+	var styledContent string
+	if selected {
+		styledContent = truncatedContent
+	} else {
+		textStyle := styles.NewStyle().Foreground(t.Text()).Render
+		styledContent = textStyle(truncatedContent)
+	}
+
+	// Create the line with proper spacing - content left-aligned, tools right-aligned
+	var text string
+	if toolInfo != "" {
+		// Calculate spacing to right-align the tool count
+		contentPart := fmt.Sprintf("%s %s", timeStr, styledContent)
+		totalContentLen := timeVisualLen + 1 + len(truncatedContent) // Use visual length for content
+		availableWidth := width - 2                                  // Account for padding
+		spacingNeeded := availableWidth - totalContentLen - toolInfoVisualLen
+		if spacingNeeded < 1 {
+			spacingNeeded = 1
+		}
+		text = fmt.Sprintf("%s%s%s", contentPart, strings.Repeat(" ", spacingNeeded), toolInfo)
+	} else {
+		text = fmt.Sprintf("%s %s", timeStr, styledContent)
+	}
+
+	var itemStyle styles.Style
+	if selected {
+		itemStyle = baseStyle.
+			Background(t.Primary()).
+			Foreground(t.BackgroundElement()).
+			Width(width).
+			PaddingLeft(1)
+	} else {
+		itemStyle = baseStyle.
+			PaddingLeft(1)
+	}
+
+	return itemStyle.Render(text)
+}
+
+func (n navigationItem) Selectable() bool {
+	return true
+}
+
+type navigationDialog struct {
+	width  int
+	height int
+	modal  *modal.Modal
+	list   list.List[navigationItem]
+	app    *app.App
+}
+
+func (n *navigationDialog) Init() tea.Cmd {
+	return nil
+}
+
+func (n *navigationDialog) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	switch msg := msg.(type) {
+	case tea.WindowSizeMsg:
+		n.width = msg.Width
+		n.height = msg.Height
+		n.list.SetMaxWidth(layout.Current.Container.Width - 12)
+	case tea.KeyPressMsg:
+		switch msg.String() {
+		case "up", "down":
+			// Handle navigation and immediately scroll to selected message
+			var cmd tea.Cmd
+			listModel, cmd := n.list.Update(msg)
+			n.list = listModel.(list.List[navigationItem])
+
+			// Get the newly selected item and scroll to it immediately
+			if item, idx := n.list.GetSelectedItem(); idx >= 0 {
+				return n, tea.Sequence(
+					cmd,
+					util.CmdHandler(ScrollToMessageMsg{MessageID: item.messageID}),
+				)
+			}
+			return n, cmd
+		case "r":
+			// Restore conversation to selected message
+			if item, idx := n.list.GetSelectedItem(); idx >= 0 {
+				return n, tea.Sequence(
+					util.CmdHandler(RestoreToMessageMsg{MessageID: item.messageID, Index: item.index}),
+					util.CmdHandler(modal.CloseModalMsg{}),
+				)
+			}
+		case "enter":
+			// Keep Enter functionality for closing the modal
+			if _, idx := n.list.GetSelectedItem(); idx >= 0 {
+				return n, util.CmdHandler(modal.CloseModalMsg{})
+			}
+		}
+	}
+
+	var cmd tea.Cmd
+	listModel, cmd := n.list.Update(msg)
+	n.list = listModel.(list.List[navigationItem])
+	return n, cmd
+}
+
+func (n *navigationDialog) Render(background string) string {
+	listView := n.list.View()
+
+	t := theme.CurrentTheme()
+	keyStyle := styles.NewStyle().Foreground(t.Warning()).Background(t.BackgroundPanel()).Render
+	mutedStyle := styles.NewStyle().Foreground(t.TextMuted()).Background(t.BackgroundPanel()).Render
+
+	helpText := keyStyle("↑/↓") + mutedStyle(" jump") + "   " + keyStyle("r") + mutedStyle(" restore")
+
+	bgColor := t.BackgroundPanel()
+	helpView := styles.NewStyle().
+		Background(bgColor).
+		Width(layout.Current.Container.Width - 14).
+		PaddingLeft(1).
+		PaddingTop(1).
+		Render(helpText)
+
+	content := strings.Join([]string{listView, helpView}, "\n")
+
+	return n.modal.Render(content, background)
+}
+
+func (n *navigationDialog) Close() tea.Cmd {
+	return nil
+}
+
+// extractMessagePreview extracts a preview from message parts
+func extractMessagePreview(parts []opencode.PartUnion) string {
+	for _, part := range parts {
+		switch casted := part.(type) {
+		case opencode.TextPart:
+			text := strings.TrimSpace(casted.Text)
+			if text != "" {
+				return text
+			}
+		}
+	}
+	return "No text content"
+}
+
+// countToolsInResponse counts tools in the assistant's response to a user message
+func countToolsInResponse(messages []app.Message, userMessageIndex int) int {
+	count := 0
+	// Look at subsequent messages to find the assistant's response
+	for i := userMessageIndex + 1; i < len(messages); i++ {
+		message := messages[i]
+		// If we hit another user message, stop looking
+		if _, isUser := message.Info.(opencode.UserMessage); isUser {
+			break
+		}
+		// Count tools in this assistant message
+		for _, part := range message.Parts {
+			switch part.(type) {
+			case opencode.ToolPart:
+				count++
+			}
+		}
+	}
+	return count
+}
+
+// NewNavigationDialog creates a new session navigation dialog
+func NewNavigationDialog(app *app.App) NavigationDialog {
+	var items []navigationItem
+
+	// Filter to only user messages and extract relevant info
+	for i, message := range app.Messages {
+		if userMsg, ok := message.Info.(opencode.UserMessage); ok {
+			preview := extractMessagePreview(message.Parts)
+			toolCount := countToolsInResponse(app.Messages, i)
+
+			items = append(items, navigationItem{
+				messageID: userMsg.ID,
+				content:   preview,
+				timestamp: time.UnixMilli(int64(userMsg.Time.Created)),
+				index:     i,
+				toolCount: toolCount,
+			})
+		}
+	}
+
+	listComponent := list.NewListComponent(
+		list.WithItems(items),
+		list.WithMaxVisibleHeight[navigationItem](12),
+		list.WithFallbackMessage[navigationItem]("No user messages in this session"),
+		list.WithAlphaNumericKeys[navigationItem](true),
+		list.WithRenderFunc(
+			func(item navigationItem, selected bool, width int, baseStyle styles.Style) string {
+				return item.Render(selected, width, false, baseStyle)
+			},
+		),
+		list.WithSelectableFunc(func(item navigationItem) bool {
+			return true
+		}),
+	)
+	listComponent.SetMaxWidth(layout.Current.Container.Width - 12)
+
+	return &navigationDialog{
+		list: listComponent,
+		app:  app,
+		modal: modal.New(
+			modal.WithTitle("Jump to Message"),
+			modal.WithMaxWidth(layout.Current.Container.Width-8),
+		),
+	}
+}