Add ExtendMsg to session for dynamic screen lists

Amolith created

The restore flow needs screens built dynamically after the command and
preset are known (snapshot selection depends on resolved config, file
picker depends on the chosen snapshot).

Add ExtendMsg which appends screens after the current cursor, replacing
any previously-queued future screens. This lets a single tea.Program
drive the entire interactive flow without splitting into multiple
sessions.

Key behaviours:
- Empty Screens slice is a no-op
- Screens beyond the cursor are truncated before appending
- Extended screens participate in normal advance/back navigation
- Window size is replayed to extended screens on activation

Change summary

internal/ui/screen.go       |  12 ++
internal/ui/session.go      |  17 ++++
internal/ui/session_test.go | 164 +++++++++++++++++++++++++++++++++++++++
3 files changed, 193 insertions(+)

Detailed changes

internal/ui/screen.go 🔗

@@ -79,3 +79,15 @@ type DoneMsg struct{}
 func DoneCmd() tea.Msg {
 	return DoneMsg{}
 }
+
+// ExtendMsg asks the session to append screens after the current
+// cursor position, replacing any screens that were already queued
+// beyond it. This allows a screen (or external code via [tea.Cmd])
+// to build the next part of the flow dynamically — for example,
+// after resolving a config to determine which command-specific
+// screens are needed.
+//
+// If Screens is empty, the message is a no-op.
+type ExtendMsg struct {
+	Screens []Screen
+}

internal/ui/session.go 🔗

@@ -144,6 +144,9 @@ func (s Session) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 
 	case DoneMsg:
 		return s.advance()
+
+	case ExtendMsg:
+		return s.extend(msg)
 	}
 
 	return s.forwardToScreen(msg)
@@ -230,6 +233,20 @@ func (s Session) advance() (tea.Model, tea.Cmd) {
 	return s, s.activateScreen()
 }
 
+// extend appends screens after the current cursor position, replacing
+// any screens that were already queued beyond it. This lets the flow
+// be built dynamically — e.g. after resolving which command-specific
+// screens are needed.
+func (s Session) extend(msg ExtendMsg) (tea.Model, tea.Cmd) {
+	if len(msg.Screens) == 0 {
+		return s, nil
+	}
+	// Truncate to current position + 1 (keep the active screen and
+	// all completed screens before it), then append the new screens.
+	s.screens = append(s.screens[:s.cursor+1], msg.Screens...)
+	return s, nil
+}
+
 // activateScreen initialises the screen at the current cursor and
 // sends it the adjusted window size so it has correct dimensions
 // immediately. Used by both navigateBack and advance.

internal/ui/session_test.go 🔗

@@ -611,6 +611,170 @@ func TestBackgroundColorForwardedToScreen(t *testing.T) {
 	}
 }
 
+func TestExtendAppendsScreens(t *testing.T) {
+	t.Parallel()
+
+	a := newSelecting("first")
+	s := New([]Screen{a}, defaultStyles())
+	s.Init()
+
+	// Extend the session with two more screens before the current
+	// screen completes.
+	b := newFake("second")
+	c := newFake("third")
+	result, _ := s.Update(ExtendMsg{Screens: []Screen{b, c}})
+	sess := result.(Session)
+
+	if len(sess.screens) != 3 {
+		t.Fatalf("screens = %d, want 3 after extend", len(sess.screens))
+	}
+
+	// The cursor should still be on the first screen.
+	if sess.cursor != 0 {
+		t.Errorf("cursor = %d, want 0 (extend should not advance)", sess.cursor)
+	}
+}
+
+func TestExtendedScreensReachableViaAdvance(t *testing.T) {
+	t.Parallel()
+
+	a := newSelecting("first")
+	s := New([]Screen{a}, defaultStyles())
+	s.Init()
+
+	// Extend before completing the first screen.
+	b := newFake("second")
+	result, _ := s.Update(ExtendMsg{Screens: []Screen{b}})
+	sess := result.(Session)
+
+	// Complete the first screen. Without extend, this would finish
+	// the session. With extend, it should advance to b.
+	result, _ = sess.Update(enterMsg())
+	result, _ = result.(Session).Update(DoneMsg{})
+	sess = result.(Session)
+
+	if sess.cursor != 1 {
+		t.Errorf("cursor = %d, want 1 after advance to extended screen", sess.cursor)
+	}
+	if sess.done {
+		t.Error("session should not be done — extended screen is still pending")
+	}
+}
+
+func TestExtendedScreensGetSizeOnActivation(t *testing.T) {
+	t.Parallel()
+
+	var receivedHeight int
+	a := newSelecting("first")
+	b := &sizeCaptureScreen{
+		fakeScreen: fakeScreen{title: "extended"},
+		onSize:     func(h int) { receivedHeight = h },
+	}
+
+	s := New([]Screen{a}, defaultStyles())
+	s.Init()
+
+	// Set terminal size, then extend, then advance.
+	result, _ := s.Update(tea.WindowSizeMsg{Width: 80, Height: 30})
+	sess := result.(Session)
+
+	result, _ = sess.Update(ExtendMsg{Screens: []Screen{b}})
+	sess = result.(Session)
+
+	result, _ = sess.Update(enterMsg())
+	result.(Session).Update(DoneMsg{})
+
+	wantHeight := 30 - chromeLines
+	if receivedHeight != wantHeight {
+		t.Errorf("extended screen received height %d, want %d", receivedHeight, wantHeight)
+	}
+}
+
+func TestExtendEmptySliceIsNoop(t *testing.T) {
+	t.Parallel()
+
+	a := newFake("only")
+	s := New([]Screen{a}, defaultStyles())
+	s.Init()
+
+	result, _ := s.Update(ExtendMsg{Screens: nil})
+	sess := result.(Session)
+
+	if len(sess.screens) != 1 {
+		t.Errorf("screens = %d, want 1 after empty extend", len(sess.screens))
+	}
+}
+
+func TestExtendBackNavThroughExtendedScreens(t *testing.T) {
+	t.Parallel()
+
+	a := newSelecting("first")
+	b := newFake("extended")
+
+	s := New([]Screen{a}, defaultStyles())
+	s.Init()
+
+	// Extend, advance, then back-nav.
+	result, _ := s.Update(ExtendMsg{Screens: []Screen{b}})
+	sess := result.(Session)
+
+	// Advance past first screen.
+	result, _ = sess.Update(enterMsg())
+	result, _ = result.(Session).Update(DoneMsg{})
+	sess = result.(Session)
+
+	if sess.cursor != 1 {
+		t.Fatalf("cursor = %d, want 1 before back nav", sess.cursor)
+	}
+
+	// Back nav should return to the first screen.
+	result, _ = sess.Update(escMsg())
+	result, _ = result.(Session).Update(BackMsg{})
+	sess = result.(Session)
+
+	if sess.cursor != 0 {
+		t.Errorf("cursor = %d, want 0 after back through extended screen", sess.cursor)
+	}
+}
+
+func TestExtendTruncatesScreensBeyondCursor(t *testing.T) {
+	t.Parallel()
+
+	// If the user backs up and a new ExtendMsg arrives, screens after
+	// the cursor should be replaced by the new ones (the old future
+	// is no longer valid).
+	a := newSelecting("cmd")
+	b := newFake("old-future")
+
+	s := New([]Screen{a, b}, defaultStyles())
+	s.Init()
+
+	// Advance to b, then back to a.
+	result, _ := s.Update(enterMsg())
+	result, _ = result.(Session).Update(DoneMsg{})
+	sess := result.(Session)
+
+	result, _ = sess.Update(escMsg())
+	result, _ = result.(Session).Update(BackMsg{})
+	sess = result.(Session)
+
+	// Now extend with new screens. The old b should be replaced.
+	c := newFake("new-future-1")
+	d := newFake("new-future-2")
+	result, _ = sess.Update(ExtendMsg{Screens: []Screen{c, d}})
+	sess = result.(Session)
+
+	// Should have a + c + d = 3 screens, not a + b + c + d.
+	if len(sess.screens) != 3 {
+		t.Errorf("screens = %d, want 3 (old future should be replaced)", len(sess.screens))
+	}
+
+	// The screen at index 1 should be c, not b.
+	if sess.screens[1].Title() != "new-future-1" {
+		t.Errorf("screen[1].Title() = %q, want %q", sess.screens[1].Title(), "new-future-1")
+	}
+}
+
 func TestSizeReplayedOnAdvance(t *testing.T) {
 	t.Parallel()