nub/internal/ui/tui/session_test.go

package tui

import (
	"context"
	"iter"
	"testing"
	"time"

	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textarea"
	"github.com/charmbracelet/bubbles/viewport"

	"nub/internal/agent"
	"nub/internal/llm"
	"nub/internal/session"
	"nub/internal/tool"
)

// idleProvider wird in diesen Tests nie tatsächlich gestreamt — der Loop
// läuft nur, um startSession realistisch zu testen (in/out-Channel-Handling),
// ohne dass eine Nachricht gesendet wird.
type idleProvider struct{}

func (idleProvider) Name() string   { return "idle" }
func (idleProvider) Caps() llm.Caps { return llm.Caps{} }
func (idleProvider) Stream(ctx context.Context, req llm.Request) (iter.Seq2[llm.Event, error], error) {
	return func(yield func(llm.Event, error) bool) {}, nil
}

func newSwitchableModel(t *testing.T, dir string, store *session.Store) *Model {
	t.Helper()
	loop := &agent.Loop{
		Provider: idleProvider{},
		Tools:    tool.NewRegistry(),
		Model:    "test-model",
		Env:      tool.Env{Cwd: dir, RepoRoot: dir},
	}
	return &Model{
		ctx:       context.Background(),
		loop:      loop,
		textarea:  textarea.New(),
		viewport:  viewport.New(80, 20),
		spinner:   spinner.New(),
		modelName: loop.Model,
	}
}

func TestStartSession_ReplaysHistoryAndBumpsGeneration(t *testing.T) {
	dir := t.TempDir()
	store, err := session.Create(dir)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { store.Close() })

	n1 := session.NewNode("", textMsg(llm.RoleUser, "hello from history"), session.NodeMeta{})
	if err := store.Append(n1); err != nil {
		t.Fatal(err)
	}
	n2 := session.NewNode(n1.ID, textMsg(llm.RoleAssistant, "hi there"), session.NodeMeta{})
	if err := store.Append(n2); err != nil {
		t.Fatal(err)
	}

	m := newSwitchableModel(t, dir, store)
	if m.sessionGen != 0 {
		t.Fatalf("sessionGen = %d before any startSession call, want 0", m.sessionGen)
	}

	cmd := m.startSession(store)
	t.Cleanup(func() {
		if m.sessionCancel != nil {
			m.sessionCancel()
		}
	})

	if m.sessionGen != 1 {
		t.Errorf("sessionGen = %d, want 1 after first startSession", m.sessionGen)
	}
	if cmd == nil {
		t.Fatal("expected startSession to return a listen cmd")
	}
	if len(m.entries) != 2 {
		t.Fatalf("entries = %+v, want 2 replayed messages", m.entries)
	}
	if m.entries[0].kind != entryUser || m.entries[0].text != "hello from history" {
		t.Errorf("entries[0] = %+v", m.entries[0])
	}
	if m.entries[1].kind != entryAssistant || m.entries[1].text != "hi there" {
		t.Errorf("entries[1] = %+v", m.entries[1])
	}

	// Der zurückgegebene Cmd muss tatsächlich auf dem neuen out-Channel
	// lauschen (gen=1) — daran hängt die Stale-Event-Erkennung beim Wechsel.
	// cmd() blockiert, bis out etwas liefert oder schließt; also im
	// Hintergrund aufrufen und den Kontext canceln, um out zu schließen.
	resultCh := make(chan interface{}, 1)
	go func() { resultCh <- cmd() }()
	m.sessionCancel()

	select {
	case msg := <-resultCh:
		evMsg, ok := msg.(eventsClosedMsg)
		if !ok {
			t.Fatalf("expected eventsClosedMsg, got %T", msg)
		}
		if evMsg.gen != 1 {
			t.Errorf("listen cmd bound to gen %d, want 1", evMsg.gen)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for the listen cmd to observe channel closure")
	}
}

func TestStartSession_RestoresHistoricalTokenUsage(t *testing.T) {
	dir := t.TempDir()
	store, err := session.Create(dir)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { store.Close() })

	n1 := session.NewNode("", textMsg(llm.RoleUser, "hi"), session.NodeMeta{})
	if err := store.Append(n1); err != nil {
		t.Fatal(err)
	}
	n2 := session.NewNode(n1.ID, textMsg(llm.RoleAssistant, "hello"),
		session.NodeMeta{Usage: llm.Usage{InputTokens: 1200, OutputTokens: 30, CacheReadTokens: 900}})
	if err := store.Append(n2); err != nil {
		t.Fatal(err)
	}
	n3 := session.NewNode(n2.ID, textMsg(llm.RoleUser, "again"), session.NodeMeta{})
	if err := store.Append(n3); err != nil {
		t.Fatal(err)
	}
	n4 := session.NewNode(n3.ID, textMsg(llm.RoleAssistant, "again reply"),
		session.NodeMeta{Usage: llm.Usage{InputTokens: 1400, OutputTokens: 20, CacheReadTokens: 1300}})
	if err := store.Append(n4); err != nil {
		t.Fatal(err)
	}

	m := newSwitchableModel(t, dir, store)
	m.tokensUsed = 999999 // Altlast einer vorherigen Session, muss überschrieben werden
	m.cacheReadTokens = 999999

	m.startSession(store)
	t.Cleanup(func() {
		if m.sessionCancel != nil {
			m.sessionCancel()
		}
	})

	wantTokens := 1200 + 30 + 1400 + 20
	if m.tokensUsed != wantTokens {
		t.Errorf("tokensUsed = %d, want %d (sum of persisted NodeMeta.Usage)", m.tokensUsed, wantTokens)
	}
	wantCache := 900 + 1300
	if m.cacheReadTokens != wantCache {
		t.Errorf("cacheReadTokens = %d, want %d", m.cacheReadTokens, wantCache)
	}
}

func TestStartSession_SwitchingCancelsPreviousLoop(t *testing.T) {
	dir := t.TempDir()
	storeA, err := session.Create(dir)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { storeA.Close() })
	storeB, err := session.Create(dir)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { storeB.Close() })

	m := newSwitchableModel(t, dir, storeA)
	m.startSession(storeA)
	oldOut := m.out
	oldGen := m.sessionGen

	m.startSession(storeB)
	t.Cleanup(func() {
		if m.sessionCancel != nil {
			m.sessionCancel()
		}
	})

	if m.sessionGen != oldGen+1 {
		t.Errorf("sessionGen = %d, want %d after switching", m.sessionGen, oldGen+1)
	}
	if m.store != storeB {
		t.Error("expected store to point at the new session")
	}
	if m.loop.Session != storeB {
		t.Error("expected loop.Session to be rebound to the new store")
	}

	select {
	case _, ok := <-oldOut:
		if ok {
			t.Error("expected the old session's out channel to close, not emit an event")
		}
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for the previous loop to shut down after cancel")
	}
}

func TestUpdate_StaleSessionEventsAreIgnored(t *testing.T) {
	m := newTestModel(t)
	m.sessionGen = 5
	m.turnActive = true

	_, cmd := m.Update(uiEventMsg{gen: 3, ev: tool.TextDelta{Text: "should be ignored"}})
	if cmd != nil {
		t.Error("stale uiEventMsg should not schedule any follow-up cmd")
	}
	if m.hasLiveText {
		t.Error("stale uiEventMsg must not be applied to model state")
	}

	_, cmd = m.Update(eventsClosedMsg{gen: 3})
	if m.quitting {
		t.Error("stale eventsClosedMsg from a switched-away session must not quit the TUI")
	}
	_ = cmd
}