tui: add shortcut to reselect pid and restart tracing

4 files changed, 107 insertions, 36 deletions

diff --git a/internal/tui/common/keys.go b/internal/tui/common/keys.go
index 4945b3b..87c947c 100644
--- a/internal/tui/common/keys.go
+++ b/internal/tui/common/keys.go
@@ -4,23 +4,24 @@ import "github.com/charmbracelet/bubbles/key"
 
 // KeyMap groups all key bindings shared by TUI screens.
 type KeyMap struct {
-	Tab      key.Binding
-	ShiftTab key.Binding
-	One      key.Binding
-	Two      key.Binding
-	Three    key.Binding
-	Four     key.Binding
-	Five     key.Binding
-	Six      key.Binding
-	Seven    key.Binding
-	DirGroup key.Binding
-	Probes   key.Binding
-	Export   key.Binding
-	Quit     key.Binding
-	Help     key.Binding
-	Enter    key.Binding
-	Esc      key.Binding
-	Refresh  key.Binding
+	Tab       key.Binding
+	ShiftTab  key.Binding
+	One       key.Binding
+	Two       key.Binding
+	Three     key.Binding
+	Four      key.Binding
+	Five      key.Binding
+	Six       key.Binding
+	Seven     key.Binding
+	DirGroup  key.Binding
+	SelectPID key.Binding
+	Probes    key.Binding
+	Export    key.Binding
+	Quit      key.Binding
+	Help      key.Binding
+	Enter     key.Binding
+	Esc       key.Binding
+	Refresh   key.Binding
 }
 
 // Keys contains the default shared key map.
@@ -29,23 +30,24 @@ var Keys = DefaultKeyMap()
 // DefaultKeyMap builds the default key bindings used by models.
 func DefaultKeyMap() KeyMap {
 	return KeyMap{
-		Tab:      key.NewBinding(key.WithKeys("tab"), key.WithHelp("tab", "next tab")),
-		ShiftTab: key.NewBinding(key.WithKeys("shift+tab"), key.WithHelp("shift+tab", "prev tab")),
-		One:      key.NewBinding(key.WithKeys("1"), key.WithHelp("1", "overview")),
-		Two:      key.NewBinding(key.WithKeys("2"), key.WithHelp("2", "syscalls")),
-		Three:    key.NewBinding(key.WithKeys("3"), key.WithHelp("3", "files")),
-		Four:     key.NewBinding(key.WithKeys("4"), key.WithHelp("4", "processes")),
-		Five:     key.NewBinding(key.WithKeys("5"), key.WithHelp("5", "lat+gaps")),
-		Six:      key.NewBinding(key.WithKeys("6"), key.WithHelp("6", "stream")),
-		Seven:    key.NewBinding(key.WithKeys("7"), key.WithHelp("7", "stream")),
-		DirGroup: key.NewBinding(key.WithKeys("d"), key.WithHelp("d", "dir group")),
-		Probes:   key.NewBinding(key.WithKeys("p"), key.WithHelp("p", "probes")),
-		Export:   key.NewBinding(key.WithKeys("e"), key.WithHelp("e", "export")),
-		Quit:     key.NewBinding(key.WithKeys("q", "ctrl+c"), key.WithHelp("q", "quit")),
-		Help:     key.NewBinding(key.WithKeys("?"), key.WithHelp("?", "help")),
-		Enter:    key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "select")),
-		Esc:      key.NewBinding(key.WithKeys("esc"), key.WithHelp("esc", "back")),
-		Refresh:  key.NewBinding(key.WithKeys("r"), key.WithHelp("r", "refresh")),
+		Tab:       key.NewBinding(key.WithKeys("tab"), key.WithHelp("tab", "next tab")),
+		ShiftTab:  key.NewBinding(key.WithKeys("shift+tab"), key.WithHelp("shift+tab", "prev tab")),
+		One:       key.NewBinding(key.WithKeys("1"), key.WithHelp("1", "overview")),
+		Two:       key.NewBinding(key.WithKeys("2"), key.WithHelp("2", "syscalls")),
+		Three:     key.NewBinding(key.WithKeys("3"), key.WithHelp("3", "files")),
+		Four:      key.NewBinding(key.WithKeys("4"), key.WithHelp("4", "processes")),
+		Five:      key.NewBinding(key.WithKeys("5"), key.WithHelp("5", "lat+gaps")),
+		Six:       key.NewBinding(key.WithKeys("6"), key.WithHelp("6", "stream")),
+		Seven:     key.NewBinding(key.WithKeys("7"), key.WithHelp("7", "stream")),
+		DirGroup:  key.NewBinding(key.WithKeys("d"), key.WithHelp("d", "dir group")),
+		SelectPID: key.NewBinding(key.WithKeys("s"), key.WithHelp("s", "select pid")),
+		Probes:    key.NewBinding(key.WithKeys("p"), key.WithHelp("p", "probes")),
+		Export:    key.NewBinding(key.WithKeys("e"), key.WithHelp("e", "export")),
+		Quit:      key.NewBinding(key.WithKeys("q", "ctrl+c"), key.WithHelp("q", "quit")),
+		Help:      key.NewBinding(key.WithKeys("?"), key.WithHelp("?", "help")),
+		Enter:     key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "select")),
+		Esc:       key.NewBinding(key.WithKeys("esc"), key.WithHelp("esc", "back")),
+		Refresh:   key.NewBinding(key.WithKeys("r"), key.WithHelp("r", "refresh")),
 	}
 }
 
@@ -55,7 +57,7 @@ func (k KeyMap) DashboardShortHelp() []key.Binding {
 	if help := k.Export.Help(); help.Key != "" || help.Desc != "" {
 		bindings = append(bindings, k.Export)
 	}
-	bindings = append(bindings, k.Probes, k.Help, k.Quit)
+	bindings = append(bindings, k.SelectPID, k.Probes, k.Help, k.Quit)
 	return bindings
 }
 
@@ -65,7 +67,7 @@ func (k KeyMap) DashboardFullHelp() [][]key.Binding {
 	if help := k.Export.Help(); help.Key != "" || help.Desc != "" {
 		controls = append(controls, k.Export)
 	}
-	controls = append(controls, k.DirGroup, k.Probes, k.Refresh, k.Help, k.Quit)
+	controls = append(controls, k.DirGroup, k.SelectPID, k.Probes, k.Refresh, k.Help, k.Quit)
 
 	return [][]key.Binding{
 		{k.One, k.Two, k.Three, k.Four, k.Five, k.Six},
diff --git a/internal/tui/common/keys_test.go b/internal/tui/common/keys_test.go
index 000dc9c..3636107 100644
--- a/internal/tui/common/keys_test.go
+++ b/internal/tui/common/keys_test.go
@@ -13,6 +13,11 @@ func TestDefaultKeyMapIncludesDirGroupBinding(t *testing.T) {
 	if probesHelp.Key != "p" || probesHelp.Desc != "probes" {
 		t.Fatalf("unexpected probes binding help: key=%q desc=%q", probesHelp.Key, probesHelp.Desc)
 	}
+
+	selectHelp := keys.SelectPID.Help()
+	if selectHelp.Key != "s" || selectHelp.Desc != "select pid" {
+		t.Fatalf("unexpected select pid binding help: key=%q desc=%q", selectHelp.Key, selectHelp.Desc)
+	}
 }
 
 func TestDashboardFullHelpIncludesDirGroupBinding(t *testing.T) {
@@ -45,6 +50,18 @@ func TestDashboardFullHelpIncludesDirGroupBinding(t *testing.T) {
 	if !found {
 		t.Fatalf("expected probes binding in dashboard full help controls")
 	}
+
+	found = false
+	for _, binding := range groups[1] {
+		help := binding.Help()
+		if help.Key == "s" && help.Desc == "select pid" {
+			found = true
+			break
+		}
+	}
+	if !found {
+		t.Fatalf("expected select pid binding in dashboard full help controls")
+	}
 }
 
 func TestDashboardShortHelpIncludesProbesBinding(t *testing.T) {
diff --git a/internal/tui/tui.go b/internal/tui/tui.go
index db810b2..032a27a 100644
--- a/internal/tui/tui.go
+++ b/internal/tui/tui.go
@@ -236,6 +236,9 @@ func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			m.probeModal = probes.NewModel(getProbeManager()).Open()
 			return m, nil
 		}
+		if m.screen == ScreenDashboard && !m.attaching && m.lastErr == nil && key.Matches(msg, m.keys.SelectPID) && !m.exporter.Visible() && !m.probeModal.Visible() && !m.dashboard.BlocksGlobalShortcuts() {
+			return m.reselectPID()
+		}
 	case tuiexport.RequestMsg:
 		return m, runExportCmd(msg.Option, m.dashboard.LatestSnapshot())
 	case tuiexport.CompletedMsg:
@@ -325,6 +328,26 @@ func (m Model) handlePidSelected(msg PidSelectedMsg) (tea.Model, tea.Cmd) {
 	return m, tea.Batch(m.spin.Tick, m.beginTraceCmd())
 }
 
+func (m Model) reselectPID() (tea.Model, tea.Cmd) {
+	m.stopTrace()
+	m.screen = ScreenPIDPicker
+	m.attaching = false
+	m.lastErr = nil
+	m.showHelp = false
+	m.exporter = tuiexport.NewModel()
+	m.probeModal = probes.NewModel(getProbeManager())
+	m.pidPicker = pidpicker.New()
+
+	var sizeCmd tea.Cmd
+	if m.width > 0 && m.height > 0 {
+		msg := tea.WindowSizeMsg{Width: m.width, Height: m.height}
+		next, cmd := m.pidPicker.Update(msg)
+		m.pidPicker = next.(pidpicker.Model)
+		sizeCmd = cmd
+	}
+	return m, tea.Batch(sizeCmd, m.pidPicker.Init())
+}
+
 func selectedPIDFilter(pid int) int {
 	if pid <= 0 {
 		return -1
diff --git a/internal/tui/tui_test.go b/internal/tui/tui_test.go
index d9a69a5..761ac0f 100644
--- a/internal/tui/tui_test.go
+++ b/internal/tui/tui_test.go
@@ -281,6 +281,35 @@ func TestExportKeyOpensModalOnDashboard(t *testing.T) {
 	}
 }
 
+func TestSelectPIDKeyReturnsToFreshPickerAndStopsTrace(t *testing.T) {
+	m := NewModel(-1, func(context.Context) error { return nil })
+	m.screen = ScreenDashboard
+	m.attaching = false
+	m.width = 120
+	m.height = 30
+	stopped := false
+	m.traceStop = func() { stopped = true }
+
+	next, cmd := m.Update(tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'s'}})
+	updated := next.(Model)
+
+	if !stopped {
+		t.Fatalf("expected active tracing to be stopped before returning to picker")
+	}
+	if updated.screen != ScreenPIDPicker {
+		t.Fatalf("expected PID picker screen, got %v", updated.screen)
+	}
+	if updated.attaching {
+		t.Fatalf("expected attaching=false on picker screen")
+	}
+	if updated.traceStop != nil {
+		t.Fatalf("expected traceStop to be cleared after stopping")
+	}
+	if cmd == nil {
+		t.Fatalf("expected picker init command when returning to picker")
+	}
+}
+
 func TestExportKeyIgnoredWhenExportDisabled(t *testing.T) {
 	flags.SetTUIExportEnable(false)
 	t.Cleanup(func() { flags.SetTUIExportEnable(true) })