nub/internal/agent/compact_test.go

package agent

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

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

func TestCompactionCutIndex_KeepsLastNTurnsRaw(t *testing.T) {
	nodes := []*session.Node{
		{ID: "u1", Message: textMsgNode(llm.RoleUser, "u1")},
		{ID: "a1", Message: textMsgNode(llm.RoleAssistant, "a1")},
		{ID: "u2", Message: textMsgNode(llm.RoleUser, "u2")},
		{ID: "a2", Message: textMsgNode(llm.RoleAssistant, "a2")},
		{ID: "u3", Message: textMsgNode(llm.RoleUser, "u3")},
	}
	if cut := compactionCutIndex(nodes, 2); cut != 2 {
		t.Errorf("cut = %d, want 2 (start of u2)", cut)
	}
	if cut := compactionCutIndex(nodes, 10); cut != 0 {
		t.Errorf("cut = %d, want 0 when keepTurns exceeds available turns", cut)
	}
}

func TestIsUserTurnBoundary_IgnoresToolResultOnlyMessages(t *testing.T) {
	toolResultMsg := llm.Message{
		Role: llm.RoleUser,
		Content: []llm.Block{
			{Kind: llm.KindToolResult, ToolUseID: "x", Result: []llm.Block{{Kind: llm.KindText, Text: "ok"}}},
		},
	}
	if isUserTurnBoundary(&session.Node{Message: toolResultMsg}) {
		t.Error("tool-result-only message should not count as a turn boundary")
	}

	realUserMsg := textMsgNode(llm.RoleUser, "hello")
	if !isUserTurnBoundary(&session.Node{Message: realUserMsg}) {
		t.Error("a real user text message should count as a turn boundary")
	}
}

func textMsgNode(role llm.Role, text string) llm.Message {
	return llm.Message{Role: role, Content: []llm.Block{{Kind: llm.KindText, Text: text}}}
}

// scriptedProvider unterscheidet Compaction-Zusammenfassungs-Requests
// (erkennbar am System-Prompt) von normalen Turn-Requests, damit die
// Aufrufreihenfolge im Test nicht von der genauen Compaction-Logik abhängt.
type scriptedProvider struct {
	caps    llm.Caps
	batches [][]llm.Event
	idx     int
}

func (p *scriptedProvider) Name() string   { return "scripted" }
func (p *scriptedProvider) Caps() llm.Caps { return p.caps }

func (p *scriptedProvider) Stream(ctx context.Context, req llm.Request) (iter.Seq2[llm.Event, error], error) {
	if len(req.System) > 0 && strings.Contains(req.System[0].Text, "Zusammenfassung") {
		return func(yield func(llm.Event, error) bool) {
			yield(llm.BlockStart{Index: 0, Block: llm.Block{Kind: llm.KindText}}, nil)
			yield(llm.BlockDelta{Index: 0, Text: "SUMMARY: files touched, decisions made, todo remains"}, nil)
			yield(llm.BlockStop{Index: 0}, nil)
			yield(llm.Done{Stop: llm.StopEnd}, nil)
		}, nil
	}

	batch := p.batches[p.idx]
	p.idx++
	return func(yield func(llm.Event, error) bool) {
		for _, ev := range batch {
			if !yield(ev, nil) {
				return
			}
		}
	}, nil
}

// TestCompaction_AutoTriggersAndPreservesRewind ist das M4-Fertig-Kriterium
// aus Abschnitt 6: eine künstlich verlängerte Session compactet automatisch,
// läuft weiter, und ein Rewind auf einen Knoten *vor* der Compaction
// funktioniert weiterhin.
func TestCompaction_AutoTriggersAndPreservesRewind(t *testing.T) {
	provider := &scriptedProvider{
		caps: llm.Caps{MaxContext: 10}, // winzig -> Budget nach wenigen Turns überschritten
		batches: [][]llm.Event{
			textEvents("resp-1"),
			textEvents("resp-2"),
			textEvents("resp-3"),
		},
	}
	sess := session.New("test-compaction")
	loop := &Loop{
		Provider:  provider,
		Tools:     tool.NewRegistry(),
		Model:     "test-model",
		Env:       tool.Env{Cwd: ".", RepoRoot: "."},
		Session:   sess,
		KeepTurns: 2,
		CompactAt: 0.75,
	}

	in := make(chan Input)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	out := loop.Run(ctx, in)

	done := make(chan struct{})
	go func() {
		defer close(done)
		for range out {
		}
	}()

	for _, text := range []string{"user one", "user two", "user three"} {
		select {
		case in <- Input{Text: text}:
		case <-ctx.Done():
			t.Fatal("timed out sending input")
		}
	}
	close(in)
	<-done

	var summary *session.Node
	var rootID string
	for id, n := range sess.Nodes {
		if n.Kind == session.NodeSummary {
			summary = n
		}
		if n.ParentID == "" && n.Kind == session.NodeMessage {
			rootID = id
		}
	}
	if summary == nil {
		t.Fatal("expected a summary node after auto-compaction")
	}
	if rootID == "" {
		t.Fatal("could not find original root node")
	}

	path := sess.PathToHead()
	if len(path) == 0 || path[0].Content[0].Text != summary.Message.Content[0].Text {
		t.Errorf("expected current path to start at the summary, got: %+v", path)
	}

	// Rewind auf den allerersten (nie kompaktierten) Node muss weiterhin
	// funktionieren.
	if err := sess.Branch(rootID); err != nil {
		t.Fatal(err)
	}
	rewound := sess.PathToHead()
	if len(rewound) != 1 || rewound[0].Content[0].Text != "user one" {
		t.Errorf("rewind to pre-compaction root failed: %+v", rewound)
	}
}