Fix real live flamegraph key handling and startup viewport sync

2 files changed, 65 insertions, 9 deletions

diff --git a/internal/tui/flamegraph/model.go b/internal/tui/flamegraph/model.go
index b205d33..66fefc9 100644
--- a/internal/tui/flamegraph/model.go
+++ b/internal/tui/flamegraph/model.go
@@ -87,8 +87,7 @@ type Model struct {
 	animation AnimationState
 	animating bool
 	paused    bool
-	// hasNavigableSnapshot flips once we have at least one selectable non-root
-	// frame. Paused mode can still bootstrap snapshots until then.
+	// hasNavigableSnapshot flips once we have at least one selectable non-root frame.
 	hasNavigableSnapshot bool
 	isDark               bool
 	keys                 flameKeyMap
@@ -308,8 +307,9 @@ func (m *Model) RefreshFromLiveTrie() bool {
 	if m.liveTrie == nil {
 		return false
 	}
-	// Keep bootstrapping while paused until we have a navigable snapshot.
-	if m.paused && m.snapshot != nil && m.hasNavigableSnapshot {
+	// Once a snapshot exists, paused mode must freeze it regardless of current
+	// navigability so selection and percentages remain stable.
+	if m.paused && m.snapshot != nil {
 		return false
 	}
 	version := m.liveTrie.Version()
@@ -371,6 +371,11 @@ func (m *Model) SetDarkMode(isDark bool) {
 }
 
 func (m *Model) rebuildFrames(animate bool) {
+	prevPath := ""
+	if len(m.frames) > 0 && m.selectedIdx >= 0 && m.selectedIdx < len(m.frames) {
+		prevPath = m.frames[m.selectedIdx].Path
+	}
+
 	var root *snapshotNode
 	rootPath := ""
 	if m.zoomRoot != nil {
@@ -391,6 +396,7 @@ func (m *Model) rebuildFrames(animate bool) {
 	if len(m.frames) > 1 {
 		m.hasNavigableSnapshot = true
 	}
+	m.restoreSelectionByPath(prevPath)
 	m.clampSelection()
 	m.recomputeFilterState()
 	m.ensureSelectionNavigable()
@@ -398,6 +404,31 @@ func (m *Model) rebuildFrames(animate bool) {
 	m.subtreeSet = computeSubtreeSetInto(m.frames, m.selectedIdx, m.subtreeSet)
 }
 
+func (m *Model) restoreSelectionByPath(path string) {
+	if path == "" || len(m.frames) == 0 {
+		return
+	}
+	if idx := m.frameIndexByPath(path); idx >= 0 {
+		m.selectedIdx = idx
+		return
+	}
+	for idx, frame := range m.frames {
+		if hasPathBoundaryPrefix(path, frame.Path) || hasPathBoundaryPrefix(frame.Path, path) {
+			m.selectedIdx = idx
+			return
+		}
+	}
+}
+
+func (m Model) frameIndexByPath(path string) int {
+	for idx, frame := range m.frames {
+		if frame.Path == path {
+			return idx
+		}
+	}
+	return -1
+}
+
 func (m *Model) zoomIn() {
 	if len(m.frames) == 0 || m.snapshot == nil {
 		m.statusMessage = "Zoom unavailable: no frame selected"
diff --git a/internal/tui/flamegraph/model_test.go b/internal/tui/flamegraph/model_test.go
index 7387ac6..e98d936 100644
--- a/internal/tui/flamegraph/model_test.go
+++ b/internal/tui/flamegraph/model_test.go
@@ -90,7 +90,7 @@ func TestRefreshFromLiveTriePausedBlocksAfterNavigableSnapshot(t *testing.T) {
 	}
 }
 
-func TestRefreshFromLiveTriePausedKeepsBootstrappingWithoutNavigableSnapshot(t *testing.T) {
+func TestRefreshFromLiveTriePausedBlocksAfterAnySnapshot(t *testing.T) {
 	trie := coreflamegraph.NewLiveTrie([]string{"comm", "path"}, "count")
 	m := NewModel(trie)
 	m.paused = true
@@ -99,11 +99,11 @@ func TestRefreshFromLiveTriePausedKeepsBootstrappingWithoutNavigableSnapshot(t *
 	m.hasNavigableSnapshot = false
 	m.lastVersion = 1
 
-	if changed := m.RefreshFromLiveTrie(); !changed {
-		t.Fatalf("expected paused refresh to continue bootstrapping before navigation is possible")
+	if changed := m.RefreshFromLiveTrie(); changed {
+		t.Fatalf("expected paused refresh to freeze after first snapshot even when non-navigable")
 	}
-	if got, want := m.lastVersion, trie.Version(); got != want {
-		t.Fatalf("expected paused bootstrap refresh to track trie version, got %d want %d", got, want)
+	if got, want := m.lastVersion, uint64(1); got != want {
+		t.Fatalf("expected paused refresh to keep existing snapshot version, got %d want %d", got, want)
 	}
 }
 
@@ -295,6 +295,31 @@ func TestLiveFixtureArrowTraversalWhileStreamingVisitsAllFrames(t *testing.T) {
 	}
 }
 
+func TestSelectionRestoresByPathAcrossLiveRefresh(t *testing.T) {
+	trie := coreflamegraph.NewLiveTrie([]string{"comm", "path", "tracepoint"}, "count")
+	coreflamegraph.SeedTestLiveFlameData(trie, 0)
+
+	m := NewModel(trie)
+	m.SetViewport(180, 40)
+	if changed := m.RefreshFromLiveTrie(); !changed {
+		t.Fatalf("expected initial refresh")
+	}
+	m = pressFlameKey(t, m, tea.KeyPressMsg{Code: tea.KeyRight})
+	selected := m.frames[m.selectedIdx].Path
+	if selected == "" || selected == "root" {
+		t.Fatalf("expected selection to move off root, got %q", selected)
+	}
+
+	trie.Reset()
+	coreflamegraph.SeedTestLiveFlameData(trie, 2)
+	if changed := m.RefreshFromLiveTrie(); !changed {
+		t.Fatalf("expected refresh after live update")
+	}
+	if got := m.frames[m.selectedIdx].Path; got != selected {
+		t.Fatalf("expected selection path to persist across refresh, got %q want %q", got, selected)
+	}
+}
+
 func TestKeyboardNavigationSingleNodeClamped(t *testing.T) {
 	m := NewModel(nil)
 	m.frames = []tuiFrame{{Name: "root", Depth: 0, Col: 0, Path: "root"}}