rony-chat-bot/internal/agent/runner_compaction_test.go

package agent

import (
	"context"
	"errors"
	"iter"
	"strings"
	"testing"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	llmpersona "github.com/VictorVargas/rony-llm-agent/pkg/persona"
)

// compactionStub is an llm.LLMClient that:
//   - reports a configurable MaxContextWindow via Capabilities
//   - records Generate calls (used for summarization)
//   - returns a fixed Stream that also bumps the runner's r.usage via
//     the usage chunk (so subsequent compactions see realistic usage)
type compactionStub struct {
	window        int
	streamInput   int // input tokens reported on the streamed usage chunk
	generateCalls int
	generateErr   error
	summary       string
	lastGenerate  llm.CompletionRequest
}

func (s *compactionStub) Generate(_ context.Context, req llm.CompletionRequest) (llm.CompletionResponse, error) {
	s.generateCalls++
	s.lastGenerate = req
	if s.generateErr != nil {
		return llm.CompletionResponse{}, s.generateErr
	}
	return llm.CompletionResponse{
		Content:    s.summary,
		StopReason: llm.StopReasonEndTurn,
	}, nil
}

func (s *compactionStub) Stream(_ context.Context, _ llm.CompletionRequest) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		yield(llm.StreamChunk{Delta: "ok", FinishReason: "stop", Usage: llm.TokenUsage{InputTokens: s.streamInput, OutputTokens: 3}}, nil)
	}
}

func (s *compactionStub) Name() string { return "compaction-stub" }
func (s *compactionStub) Capabilities() llm.ProviderCapabilities {
	return llm.ProviderCapabilities{MaxContextWindow: s.window}
}

// pumpStream drives a single Stream call through the runner so the usage
// chunk updates r.usage — needed because the compaction trigger reads
// from r.usage, not from the request itself.
func pumpStream(t *testing.T, r *Runner, history []Message) {
	t.Helper()
	for chunk, err := range r.Stream(context.Background(), history) {
		if err != nil {
			t.Fatal(err)
		}
		_ = chunk
	}
}

// makeHistory builds a synthetic N-turn conversation (user + assistant
// pairs), each "long" enough to be obviously past any small threshold.
func makeHistory(turns int, contentLen int) []Message {
	out := make([]Message, 0, turns*2)
	payload := strings.Repeat("x", contentLen)
	for i := 0; i < turns; i++ {
		out = append(out,
			Message{Role: RoleUser, Content: payload + " q"},
			Message{Role: RoleAssistant, Content: payload + " a"},
		)
	}
	return out
}

func personaMinimal() llmpersona.Persona {
	return llmpersona.Persona{Name: "t", Tone: "concise", Language: "English"}
}

func TestCompactDisabledIsNoop(t *testing.T) {
	cli := &compactionStub{window: 100, streamInput: 999, summary: "ignored"}
	r := New(cli, personaMinimal(), "sys", nil, 5) // no WithCompaction
	pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
	got, err := r.Compact(context.Background(), makeHistory(8, 200))
	if err != nil {
		t.Fatal(err)
	}
	if cli.generateCalls != 0 {
		t.Errorf("Generate called %d times, want 0 (compaction disabled)", cli.generateCalls)
	}
	if len(got) != 16 {
		t.Errorf("len(history) = %d, want 16 (unchanged)", len(got))
	}
	stats := r.LastCompaction()
	if stats.Happened {
		t.Errorf("stats.Happened = true, want false")
	}
}

func TestCompactBelowThresholdIsNoop(t *testing.T) {
	// window = 1000, threshold = 0.5 → 500. Previous turn reported only
	// 50 input tokens — well below the trigger.
	cli := &compactionStub{window: 1000, streamInput: 50, summary: "ignored"}
	r := New(cli, personaMinimal(), "sys", nil, 5).
		WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 2})
	pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
	got, err := r.Compact(context.Background(), makeHistory(8, 40))
	if err != nil {
		t.Fatal(err)
	}
	if cli.generateCalls != 0 {
		t.Errorf("Generate called %d times, want 0 (below threshold)", cli.generateCalls)
	}
	if len(got) != 16 {
		t.Errorf("len(history) = %d, want 16 (unchanged)", len(got))
	}
}

func TestCompactUsesEstimatedInputBeforeFirstStream(t *testing.T) {
	cli := &compactionStub{window: 400, summary: "summary"}
	r := New(cli, personaMinimal(), "system prompt", nil, 5).
		WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 2})

	got, err := r.Compact(context.Background(), makeHistory(8, 200))
	if err != nil {
		t.Fatal(err)
	}
	if cli.generateCalls != 1 {
		t.Fatalf("Generate called %d times, want 1", cli.generateCalls)
	}
	if len(got) != 5 {
		t.Fatalf("len(compacted) = %d, want 5", len(got))
	}
}

func TestCompactShortHistoryIsNoop(t *testing.T) {
	// Above threshold but only 3 turns total (3 user, 3 assistant = 6
	// messages). KeepRecentTurns = 4, so splitByTurns has nothing older
	// to summarize.
	cli := &compactionStub{window: 100, streamInput: 999, summary: "ignored"}
	r := New(cli, personaMinimal(), "sys", nil, 5).
		WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 4})
	pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
	history := makeHistory(3, 200) // 6 messages, 3 user turns
	got, err := r.Compact(context.Background(), history)
	if err != nil {
		t.Fatal(err)
	}
	if cli.generateCalls != 0 {
		t.Errorf("Generate called %d times, want 0 (not enough turns)", cli.generateCalls)
	}
	if len(got) != len(history) {
		t.Errorf("len(history) = %d, want %d (unchanged)", len(got), len(history))
	}
}

func TestCompactUnknownWindowIsNoop(t *testing.T) {
	// Provider doesn't report a window (Capabilities returns 0).
	cli := &compactionStub{window: 0, streamInput: 999, summary: "ignored"}
	r := New(cli, personaMinimal(), "sys", nil, 5).
		WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.5, KeepRecentTurns: 2})
	pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
	got, err := r.Compact(context.Background(), makeHistory(8, 200))
	if err != nil {
		t.Fatal(err)
	}
	if cli.generateCalls != 0 {
		t.Errorf("Generate called %d times, want 0 (window unknown)", cli.generateCalls)
	}
	if len(got) != 16 {
		t.Errorf("len(history) = %d, want 16 (unchanged)", len(got))
	}
}

func TestCompactTriggersSummarize(t *testing.T) {
	// 8 user turns × ~200 chars each = ~1600 chars ≈ 400 tokens. With
	// window=400, threshold 0.5 = 200, and the stub reports 999 input
	// tokens on the previous turn, compaction fires.
	cli := &compactionStub{window: 400, streamInput: 999, summary: "user asked 8 things, here's the gist"}
	r := New(cli, personaMinimal(), "sys", nil, 5).
		WithCompaction(CompactionConfig{
			Enabled:         true,
			ThresholdRatio:  0.5,
			KeepRecentTurns: 2,
		})
	pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
	history := makeHistory(8, 200)
	got, err := r.Compact(context.Background(), history)
	if err != nil {
		t.Fatal(err)
	}
	if cli.generateCalls != 1 {
		t.Fatalf("Generate called %d times, want 1 (compaction fired)", cli.generateCalls)
	}
	// The summary call's messages should be [system, user(transcript)].
	if len(cli.lastGenerate.Messages) != 2 {
		t.Fatalf("summarize call has %d messages, want 2", len(cli.lastGenerate.Messages))
	}
	if cli.lastGenerate.Messages[0].Role != RoleSystem {
		t.Errorf("summarize call[0].Role = %q, want system", cli.lastGenerate.Messages[0].Role)
	}
	if cli.lastGenerate.Messages[1].Role != RoleUser {
		t.Errorf("summarize call[1].Role = %q, want user", cli.lastGenerate.Messages[1].Role)
	}
	if !strings.Contains(cli.lastGenerate.Messages[1].Content, "User:") {
		t.Errorf("summarize transcript missing 'User:' lines")
	}
	if got[0].Role != RoleSystem {
		t.Errorf("compacted[0].Role = %q, want system (summary)", got[0].Role)
	}
	if !strings.Contains(got[0].Content, "user asked 8 things") {
		t.Errorf("compacted[0] missing summary text: %q", got[0].Content)
	}
	// Recent tail: 2 turns = 4 messages, exactly the last 4 of the
	// original history.
	if len(got) != 1+4 {
		t.Fatalf("len(compacted) = %d, want 5 (summary + 4)", len(got))
	}
	for i := 0; i < 4; i++ {
		if got[1+i].Content != history[12+i].Content {
			t.Errorf("kept[%d] content differs from history[%d]", i, 12+i)
		}
	}
	stats := r.LastCompaction()
	if !stats.Happened {
		t.Fatal("stats.Happened = false, want true")
	}
	if stats.OlderTurns != 6 {
		t.Errorf("OlderTurns = %d, want 6", stats.OlderTurns)
	}
	if stats.KeptTurns != 2 {
		t.Errorf("KeptTurns = %d, want 2", stats.KeptTurns)
	}
	if stats.WindowTokens != 400 {
		t.Errorf("WindowTokens = %d, want 400", stats.WindowTokens)
	}
	wantUsed := estimatePromptTokens("sys") + totalPromptTokens(history)
	if stats.UsedTokens != wantUsed {
		t.Errorf("UsedTokens = %d, want %d", stats.UsedTokens, wantUsed)
	}
	if stats.SummaryTokens == 0 {
		t.Error("SummaryTokens = 0, want > 0")
	}
}

func TestCompactSummarizeFailsFallsBackToTruncation(t *testing.T) {
	cli := &compactionStub{
		window:      400,
		streamInput: 999,
		generateErr: errors.New("summarize boom"),
	}
	r := New(cli, personaMinimal(), "sys", nil, 5).
		WithCompaction(CompactionConfig{
			Enabled:         true,
			ThresholdRatio:  0.5,
			KeepRecentTurns: 2,
		})
	pumpStream(t, r, []Message{{Role: RoleUser, Content: "hi"}})
	// 7 completed turns + 1 trailing user message (the request's new
	// question). Real production requests always end on a user message.
	history := append(makeHistory(7, 200), Message{Role: RoleUser, Content: "new question"})
	got, err := r.Compact(context.Background(), history)
	if err != nil {
		t.Fatalf("Compact should swallow summarize failure, got err: %v", err)
	}
	if cli.generateCalls != 1 {
		t.Errorf("Generate called %d times, want 1 (tried once before fallback)", cli.generateCalls)
	}
	// Fallback should have dropped the oldest turns. The trailing user
	// message (the current question) must always be preserved.
	last := got[len(got)-1]
	if last.Role != RoleUser {
		t.Errorf("last message role = %q, want user", last.Role)
	}
	if last.Content != "new question" {
		t.Errorf("last message content = %q, want %q", last.Content, "new question")
	}
	stats := r.LastCompaction()
	if !stats.Happened {
		t.Error("stats.Happened = false, want true (fallback counts as compaction)")
	}
	if stats.OlderTurns == 0 {
		t.Error("OlderTurns = 0, want > 0 (fallback dropped turns)")
	}
}

func TestCompactEmptyHistoryIsNoop(t *testing.T) {
	cli := &compactionStub{window: 100, streamInput: 999}
	r := New(cli, personaMinimal(), "sys", nil, 5).
		WithCompaction(CompactionConfig{Enabled: true, ThresholdRatio: 0.1, KeepRecentTurns: 2})
	got, err := r.Compact(context.Background(), nil)
	if err != nil {
		t.Fatal(err)
	}
	if got != nil {
		t.Errorf("got = %v, want nil", got)
	}
	if cli.generateCalls != 0 {
		t.Errorf("Generate called %d times, want 0", cli.generateCalls)
	}
}

func TestFitHistoryDropsOldTurnsToFit(t *testing.T) {
	cli := &compactionStub{window: 400}
	r := New(cli, personaMinimal(), "sys", nil, 5)
	history := makeHistory(4, 400)

	got, err := r.fitHistory("sys", history)
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 2 {
		t.Fatalf("len(history) = %d, want 2", len(got))
	}
	if got[0].Role != RoleUser || got[1].Role != RoleAssistant {
		t.Errorf("remaining turn roles = %q, %q", got[0].Role, got[1].Role)
	}
}

func TestSplitByTurns(t *testing.T) {
	history := []Message{
		{Role: RoleUser, Content: "u1"},
		{Role: RoleAssistant, Content: "a1"},
		{Role: RoleUser, Content: "u2"},
		{Role: RoleAssistant, Content: "a2"},
		{Role: RoleUser, Content: "u3"},
		{Role: RoleAssistant, Content: "a3"},
		{Role: RoleUser, Content: "u4"},
		{Role: RoleAssistant, Content: "a4"},
	}
	older, recent := splitByTurns(history, 2)
	if len(older) != 4 {
		t.Errorf("len(older) = %d, want 4", len(older))
	}
	if len(recent) != 4 {
		t.Errorf("len(recent) = %d, want 4", len(recent))
	}
	if older[0].Content != "u1" || older[3].Content != "a2" {
		t.Errorf("older = %v, want [u1,a1,u2,a2]", older)
	}
	if recent[0].Content != "u3" || recent[3].Content != "a4" {
		t.Errorf("recent = %v, want [u3,a3,u4,a4]", recent)
	}
}

func TestSplitByTurnsNothingToCompact(t *testing.T) {
	history := []Message{
		{Role: RoleUser, Content: "u1"},
		{Role: RoleAssistant, Content: "a1"},
		{Role: RoleUser, Content: "u2"},
	}
	older, recent := splitByTurns(history, 4)
	if older != nil {
		t.Errorf("older = %v, want nil", older)
	}
	if len(recent) != 3 {
		t.Errorf("len(recent) = %d, want 3 (everything)", len(recent))
	}
}

func TestRenderTranscriptSkipsToolAndEmpty(t *testing.T) {
	msgs := []Message{
		{Role: RoleUser, Content: "hi"},
		{Role: RoleAssistant, Content: "hello!"},
		{Role: RoleAssistant, Content: ""}, // streaming placeholder
		{Role: RoleSystem, Content: "internal"},
	}
	got := renderTranscript(msgs)
	want := "User: hi\nAssistant: hello!\n"
	if got != want {
		t.Errorf("renderTranscript =\n%q\nwant\n%q", got, want)
	}
}

func TestEstimateTokens(t *testing.T) {
	cases := []struct {
		in   string
		want int
	}{
		{"", 0},
		{"a", 1},
		{"abcd", 2},
		{strings.Repeat("x", 400), 101},
	}
	for _, c := range cases {
		if got := estimateTokens(c.in); got != c.want {
			t.Errorf("estimateTokens(%q) = %d, want %d", c.in, got, c.want)
		}
	}
}

func TestEstimatePromptTokensIsConservative(t *testing.T) {
	text := strings.Repeat("x", 300)
	if got := estimatePromptTokens(text); got <= estimateTokens(text) {
		t.Errorf("estimatePromptTokens = %d, want greater than estimateTokens = %d", got, estimateTokens(text))
	}
}

func TestTruncateToBudgetAlwaysKeepsLastUserTurn(t *testing.T) {
	// 3 long completed turns + a short trailing user message (the new
	// question). Budget so small only the last turn can fit.
	long := strings.Repeat("x", 400)
	history := []Message{
		{Role: RoleUser, Content: long + " q1"},
		{Role: RoleAssistant, Content: long + " a1"},
		{Role: RoleUser, Content: long + " q2"},
		{Role: RoleAssistant, Content: long + " a2"},
		{Role: RoleUser, Content: long + " q3"},
		{Role: RoleAssistant, Content: long + " a3"},
		{Role: RoleUser, Content: "current question"},
	}
	dropped, kept := truncateToBudget(history, 10)
	if len(kept) == 0 {
		t.Fatal("kept is empty")
	}
	if kept[len(kept)-1].Content != "current question" {
		t.Errorf("last kept message content = %q, want %q",
			kept[len(kept)-1].Content, "current question")
	}
	if len(dropped)+len(kept) != len(history) {
		t.Errorf("dropped+kept = %d, want %d", len(dropped)+len(kept), len(history))
	}
}