Fix real live flamegraph key handling and startup viewport sync

7 files changed, 273 insertions, 27 deletions

diff --git a/internal/flamegraph/livetrie.go b/internal/flamegraph/livetrie.go
index 0d42b6b..13d7de9 100644
--- a/internal/flamegraph/livetrie.go
+++ b/internal/flamegraph/livetrie.go
@@ -12,7 +12,11 @@ import (
 	"ior/internal/event"
 )
 
-const liveTrieMinFraction = 0.001
+const (
+	liveTrieMinFraction                     = 0.001
+	liveTrieMinVisibleChildrenWhenPruned    = 8
+	liveTrieVisibleChildrenFallbackMaxDepth = 1
+)
 
 type trieSnapshot struct {
 	Name     string          `json:"n"`
@@ -244,29 +248,45 @@ func subtreeTotal(node *trieNode) uint64 {
 }
 
 func buildSnapshot(node *trieNode, depth int, minFraction float64, rootTotal uint64) *trieSnapshot {
-	snapshot, _ := buildSnapshotWithTotal(node, depth, minFraction, rootTotal)
+	snapshot, _ := buildSnapshotWithTotal(node, depth, minFraction, rootTotal, false)
 	return snapshot
 }
 
-func buildSnapshotWithTotal(node *trieNode, depth int, minFraction float64, rootTotal uint64) (*trieSnapshot, uint64) {
+type childSnapshotState struct {
+	node     *trieNode
+	snapshot *trieSnapshot
+	total    uint64
+}
+
+func buildSnapshotWithTotal(node *trieNode, depth int, minFraction float64, rootTotal uint64, forceKeep bool) (*trieSnapshot, uint64) {
 	total := node.value
 	children := slices.Clone(node.children)
 	sort.Slice(children, func(i, j int) bool {
 		return children[i].name < children[j].name
 	})
 
-	childSnapshots := make([]*trieSnapshot, 0, len(children))
+	childStates := make([]childSnapshotState, 0, len(children))
 	for _, child := range children {
-		childSnapshot, childTotal := buildSnapshotWithTotal(child, depth+1, minFraction, rootTotal)
+		childSnapshot, childTotal := buildSnapshotWithTotal(child, depth+1, minFraction, rootTotal, false)
 		total += childTotal
-		if childSnapshot != nil {
-			childSnapshots = append(childSnapshots, childSnapshot)
-		}
+		childStates = append(childStates, childSnapshotState{
+			node:     child,
+			snapshot: childSnapshot,
+			total:    childTotal,
+		})
 	}
 
-	if depth > 0 && rootTotal > 0 && float64(total)/float64(rootTotal) < minFraction {
+	if !forceKeep && depth > 0 && rootTotal > 0 && float64(total)/float64(rootTotal) < minFraction {
 		return nil, total
 	}
+	ensureFallbackVisibleChildren(childStates, depth, minFraction, rootTotal)
+
+	childSnapshots := make([]*trieSnapshot, 0, len(childStates))
+	for _, child := range childStates {
+		if child.snapshot != nil {
+			childSnapshots = append(childSnapshots, child.snapshot)
+		}
+	}
 
 	snapshot := &trieSnapshot{
 		Name:  node.name,
@@ -278,3 +298,43 @@ func buildSnapshotWithTotal(node *trieNode, depth int, minFraction float64, root
 	}
 	return snapshot, total
 }
+
+func ensureFallbackVisibleChildren(children []childSnapshotState, depth int, minFraction float64, rootTotal uint64) {
+	if depth > liveTrieVisibleChildrenFallbackMaxDepth {
+		return
+	}
+	visible := 0
+	for _, child := range children {
+		if child.snapshot != nil {
+			visible++
+		}
+	}
+	if visible > 0 {
+		return
+	}
+
+	candidates := make([]int, 0, len(children))
+	for idx, child := range children {
+		if child.total > 0 {
+			candidates = append(candidates, idx)
+		}
+	}
+	sort.Slice(candidates, func(i, j int) bool {
+		left := children[candidates[i]]
+		right := children[candidates[j]]
+		if left.total == right.total {
+			return left.node.name < right.node.name
+		}
+		return left.total > right.total
+	})
+
+	limit := liveTrieMinVisibleChildrenWhenPruned
+	if len(candidates) < limit {
+		limit = len(candidates)
+	}
+	for i := 0; i < limit; i++ {
+		idx := candidates[i]
+		forced, _ := buildSnapshotWithTotal(children[idx].node, depth+1, minFraction, rootTotal, true)
+		children[idx].snapshot = forced
+	}
+}
diff --git a/internal/flamegraph/livetrie_test.go b/internal/flamegraph/livetrie_test.go
index 632f668..c5ed32c 100644
--- a/internal/flamegraph/livetrie_test.go
+++ b/internal/flamegraph/livetrie_test.go
@@ -221,6 +221,41 @@ func TestLiveTrieSnapshotJSONPrunesTinyNodes(t *testing.T) {
 	}
 }
 
+func TestLiveTrieSnapshotJSONKeepsFallbackChildrenWhenAllAreTinyAtRoot(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm"}, "count")
+	const total = 6000
+	for i := 0; i < total; i++ {
+		comm := fmt.Sprintf("svc-%04d", i)
+		lt.Ingest(newTestPair(comm, 42, uint32(100000+i), "/tmp/a", 1, 1, 1))
+	}
+
+	snap := decodeLiveSnapshot(t, lt)
+	if len(snap.Children) == 0 {
+		t.Fatalf("expected fallback root children when pruning would hide every branch")
+	}
+	if got, want := len(snap.Children), liveTrieMinVisibleChildrenWhenPruned; got != want {
+		t.Fatalf("expected fallback to keep %d root children, got %d", want, got)
+	}
+}
+
+func TestLiveTrieSnapshotJSONKeepsFallbackChildrenAtDepthOne(t *testing.T) {
+	lt := NewLiveTrie([]string{"comm", "pid"}, "count")
+	const total = 6000
+	for i := 0; i < total; i++ {
+		pid := uint32(100000 + i)
+		lt.Ingest(newTestPair("svc", pid, pid, "/tmp/a", 1, 1, 1))
+	}
+
+	snap := decodeLiveSnapshot(t, lt)
+	commNode := findSnapshotPath(t, &snap, "svc")
+	if len(commNode.Children) == 0 {
+		t.Fatalf("expected fallback depth-one children for pid branches")
+	}
+	if got, want := len(commNode.Children), liveTrieMinVisibleChildrenWhenPruned; got != want {
+		t.Fatalf("expected fallback to keep %d depth-one children, got %d", want, got)
+	}
+}
+
 func TestLiveTrieConcurrentIngestAndSnapshot(t *testing.T) {
 	lt := NewLiveTrie([]string{"comm", "pid"}, "count")
 
diff --git a/internal/tui/dashboard/model_test.go b/internal/tui/dashboard/model_test.go
index 8904a2f..d5b78e0 100644
--- a/internal/tui/dashboard/model_test.go
+++ b/internal/tui/dashboard/model_test.go
@@ -215,7 +215,7 @@ func TestSetLiveTriePreloadsInitialSnapshotWithoutVersionChange(t *testing.T) {
 	}
 }
 
-func TestFlameTickPausedContinuesBootstrapRefresh(t *testing.T) {
+func TestFlameTickPausedFreezesAfterInitialSnapshot(t *testing.T) {
 	liveTrie := coreflamegraph.NewLiveTrie([]string{"comm", "path"}, "count")
 	m := NewModelWithConfig(nil, nil, 250, common.DefaultKeyMap())
 	m.SetLiveTrie(liveTrie)
@@ -235,8 +235,8 @@ func TestFlameTickPausedContinuesBootstrapRefresh(t *testing.T) {
 
 	next, _ = model.Update(flameTickMsg{})
 	model = next.(Model)
-	if got, want := model.flamegraphModel.LastVersion(), liveTrie.Version(); got != want {
-		t.Fatalf("expected paused flame tick bootstrap to refresh version, got %d want %d", got, want)
+	if got, want := model.flamegraphModel.LastVersion(), initialVersion; got != want {
+		t.Fatalf("expected paused flame tick to freeze version at %d, got %d", want, got)
 	}
 }
 
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"}}
diff --git a/internal/tui/tui.go b/internal/tui/tui.go
index 0381784..19e164f 100644
--- a/internal/tui/tui.go
+++ b/internal/tui/tui.go
@@ -367,7 +367,10 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.attaching = false
 		m.dashboard.SetStreamSource(m.runtime.eventStreamSource())
 		m.dashboard.SetLiveTrie(m.runtime.liveTrie())
-		return m, m.dashboard.Init()
+		width, height := common.EffectiveViewport(m.width, m.height)
+		next, sizeCmd := m.dashboard.Update(tea.WindowSizeMsg{Width: width, Height: height})
+		m.dashboard = next.(dashboardui.Model)
+		return m, tea.Batch(sizeCmd, m.dashboard.Init())
 	case TracingErrorMsg:
 		m.attaching = false
 		m.lastErr = msg.Err
@@ -430,7 +433,9 @@ func (m *Model) normalizeKeyEvent(msg tea.Msg) (tea.Msg, bool) {
 			return nil, false
 		}
 		// Fallback: treat release as press for terminals that only emit release events.
-		m.armPressSuppression(keyID)
+		if shouldSuppressMatchingPressAfterRelease(pressMsg) {
+			m.armPressSuppression(keyID)
+		}
 		m.recordKeyEvent(pressMsg, false)
 		return pressMsg, true
 	default:
@@ -474,18 +479,24 @@ func (m *Model) recordKeyEvent(msg tea.KeyPressMsg, wasPress bool) {
 }
 
 func keyEventID(msg tea.KeyPressMsg) string {
-	return fmt.Sprintf("code:%d/mod:%d", msg.Code, msg.Mod)
+	return fmt.Sprintf("code:%d/mod:%d/key:%q/text:%q", msg.Code, msg.Mod, msg.String(), msg.Text)
 }
 
 func releaseHasIdentity(msg tea.KeyPressMsg) bool {
-	if msg.Code != 0 {
+	if msg.Text != "" {
 		return true
 	}
-	if msg.Text != "" {
+	keyStr := msg.String()
+	if keyStr != "" && keyStr != "\x00" {
 		return true
 	}
+	// Some terminals emit release-only space events without text identity.
+	return msg.Code == tea.KeySpace
+}
+
+func shouldSuppressMatchingPressAfterRelease(msg tea.KeyPressMsg) bool {
 	keyStr := msg.String()
-	return keyStr != "" && keyStr != "\x00"
+	return msg.Code == tea.KeySpace || keyStr == " " || keyStr == "space" || msg.Text == " "
 }
 
 func (m Model) updateActiveModel(msg tea.Msg) (tea.Model, tea.Cmd) {
diff --git a/internal/tui/tui_test.go b/internal/tui/tui_test.go
index 6cdc427..946d0e3 100644
--- a/internal/tui/tui_test.go
+++ b/internal/tui/tui_test.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"os"
 	"path/filepath"
+	"regexp"
 	"strings"
 	"testing"
 	"time"
@@ -281,6 +282,69 @@ func TestTracingStartedRebindsEventStreamSource(t *testing.T) {
 	}
 }
 
+func TestTracingStartedUsesCurrentViewportForFlameNavigationWithoutResize(t *testing.T) {
+	trie := coreflamegraph.NewLiveTrie([]string{"comm", "path", "tracepoint"}, "count")
+	coreflamegraph.SeedTestFlameData(trie)
+
+	m := NewModel(-1, func(context.Context) error { return nil })
+	m.screen = ScreenDashboard
+	m.attaching = true
+	m.width = 120
+	m.height = 30
+	m.runtime.SetLiveTrie(trie)
+
+	next, _ := m.Update(TracingStartedMsg{})
+	m = next.(Model)
+
+	if strings.Contains(m.View().Content, "sel:none") {
+		t.Fatalf("expected flamegraph selection to be available immediately after tracing start")
+	}
+
+	selectedLabel := func(view string) string {
+		re := regexp.MustCompile(`sel:[0-9]+/[0-9]+ ([^|]+) \|`)
+		match := re.FindStringSubmatch(view)
+		if len(match) != 2 {
+			return ""
+		}
+		return strings.TrimSpace(match[1])
+	}
+
+	moved := false
+	before := selectedLabel(m.View().Content)
+	for i := 0; i < 12 && !moved; i++ {
+		next, _ = m.Update(tea.KeyPressMsg{Code: tea.KeyRight})
+		m = next.(Model)
+		after := selectedLabel(m.View().Content)
+		if after != "" && after != before {
+			moved = true
+			break
+		}
+	}
+	if !moved {
+		t.Fatalf("expected arrow navigation to move selection without requiring resize, view=%q", m.View().Content)
+	}
+}
+
+func TestTracingStartedAppliesViewportWhenModelSizeIsUnset(t *testing.T) {
+	trie := coreflamegraph.NewLiveTrie([]string{"comm", "path", "tracepoint"}, "count")
+	coreflamegraph.SeedTestFlameData(trie)
+
+	m := NewModel(-1, func(context.Context) error { return nil })
+	m.screen = ScreenDashboard
+	m.attaching = true
+	m.runtime.SetLiveTrie(trie)
+	m.width = 0
+	m.height = 0
+
+	next, _ := m.Update(TracingStartedMsg{})
+	m = next.(Model)
+
+	view := m.View().Content
+	if strings.Contains(view, "sel:none") {
+		t.Fatalf("expected tracing start to apply an effective viewport even when width/height are unset")
+	}
+}
+
 func TestExportKeyOpensModalOnDashboard(t *testing.T) {
 	flags.SetTUIExportEnable(true)
 	t.Cleanup(func() { flags.SetTUIExportEnable(true) })
@@ -407,6 +471,26 @@ func TestNormalizeKeyEventIgnoresUnidentifiedRelease(t *testing.T) {
 	}
 }
 
+func TestNormalizeKeyEventReleaseFallbackDoesNotSuppressArrowPress(t *testing.T) {
+	m := NewModel(-1, func(context.Context) error { return nil })
+
+	normalized, ok := m.normalizeKeyEvent(tea.KeyReleaseMsg{Code: tea.KeyRight})
+	if !ok {
+		t.Fatalf("expected right release fallback to be handled")
+	}
+	if _, isPress := normalized.(tea.KeyPressMsg); !isPress {
+		t.Fatalf("expected release fallback to normalize to KeyPressMsg, got %T", normalized)
+	}
+
+	normalized, ok = m.normalizeKeyEvent(tea.KeyPressMsg{Code: tea.KeyRight})
+	if !ok {
+		t.Fatalf("expected right key press to be accepted after release fallback")
+	}
+	if _, isPress := normalized.(tea.KeyPressMsg); !isPress {
+		t.Fatalf("expected normalized message to be KeyPressMsg, got %T", normalized)
+	}
+}
+
 func TestFlameOrderKeyDoesNotOpenProbeModal(t *testing.T) {
 	m := NewModel(-1, func(context.Context) error { return nil })
 	m.screen = ScreenDashboard