diff --git a/internal/agent/runner.go b/internal/agent/runner.go index 14f18d2..2c2e441 100644 --- a/internal/agent/runner.go +++ b/internal/agent/runner.go @@ -5,12 +5,19 @@ // straight Q&A flow: no tools, no multi-iteration reasoning. Calling the // underlying provider keeps the prompt under our full control so we can // inject RAG context into the system message exactly where we want it. +// +// Auto-compaction (see CompactionConfig / Compact) folds the older portion +// of a long conversation into a single summary system message when the +// previous turn's input tokens approach the model's reported context +// window. It is opt-in via WithCompaction and short-circuits silently when +// the window is unknown or the history is too short to bother. package agent import ( "context" "fmt" "iter" + "log/slog" "strings" "github.com/VictorVargas/rony-llm-agent/pkg/llm" @@ -40,6 +47,9 @@ type Runner struct { store *portfolio.Store topK int usage *Usage + + compaction CompactionConfig + lastCompact CompactionStats } type Usage struct { @@ -47,18 +57,64 @@ type Usage struct { OutputTokens int } +// CompactionConfig mirrors the YAML `compaction:` block. Zero value means +// compaction is disabled. +type CompactionConfig struct { + Enabled bool + ThresholdRatio float64 // 0–1, default 0.75 when Enabled + KeepRecentTurns int // default 4 when Enabled + SummarySystemPrompt string // empty → built-in bilingual default +} + +// CompactionStats reports the outcome of the most recent Compact call, so +// the HTTP handler can surface a "context compacted" event to the client. +type CompactionStats struct { + Happened bool // true when older turns were folded into a summary + OlderTurns int // user turns that got summarized away + KeptTurns int // user turns kept verbatim after the summary + SummaryTokens int // approx. token count of the summary text + WindowTokens int // model's reported context window at compaction time + UsedTokens int // input tokens reported by the previous turn +} + // New returns a Runner. The store may be nil (RAG disabled). // systemPrompt is the full hand-written prompt from configs/...yaml. func New(client llm.LLMClient, p llmpersona.Persona, systemPrompt string, store *portfolio.Store, topK int) *Runner { return &Runner{client: client, persona: p, systemPrompt: systemPrompt, store: store, topK: topK, usage: &Usage{}} } +// WithCompaction enables auto-compaction with the given config. Returns +// the receiver for chaining. Pass a zero-value CompactionConfig to keep +// compaction disabled. +func (r *Runner) WithCompaction(cfg CompactionConfig) *Runner { + r.compaction = cfg + if cfg.Enabled { + if r.compaction.ThresholdRatio <= 0 || r.compaction.ThresholdRatio > 1 { + r.compaction.ThresholdRatio = 0.75 + } + if r.compaction.KeepRecentTurns <= 0 { + r.compaction.KeepRecentTurns = 4 + } + } + return r +} + // LastUsage returns the token usage recorded on the most recent call. func (r *Runner) LastUsage() Usage { return *r.usage } +// LastCompaction returns the outcome of the most recent Compact call. +// Useful for the HTTP handler to emit a "compaction" SSE event after a +// stream finishes. +func (r *Runner) LastCompaction() CompactionStats { return r.lastCompact } + // BuildMessages prepares the system prompt and turns the chat history into // the upstream message list. The system prompt includes RAG context for the // user's last message (if RAG is enabled). +// +// Compaction is NOT applied here — callers invoke Compact explicitly before +// BuildMessages so the same Compact call doesn't run twice per request +// (Stream calls BuildMessages internally, and handlers sometimes call it +// first to extract the RAG context for the sources event). func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Message, string, error) { ragContext := "" if r.store != nil && len(history) > 0 { @@ -69,7 +125,7 @@ func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Messag return nil, "", fmt.Errorf("rag search: %w", err) } if len(hits) > 0 { - ragContext = formatHits(hits) + ragContext = r.limitRAGContext(formatHits(hits), history) } } } @@ -77,12 +133,205 @@ func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Messag // The system prompt comes from the YAML, not from the persona struct. // Keep the persona around only for the UI greeting. system := botpersona.BuildSystemPrompt(r.systemPrompt, ragContext) + history, err := r.fitHistory(system, history) + if err != nil { + return nil, "", err + } msgs := make([]Message, 0, len(history)+1) msgs = append(msgs, Message{Role: RoleSystem, Content: system}) msgs = append(msgs, history...) return msgs, ragContext, nil } +const maxContextSafetyMargin = 256 + +func contextBudget(window int) int { + margin := int(float64(window) * 0.125) + if margin > maxContextSafetyMargin { + margin = maxContextSafetyMargin + } + if margin < 1 { + margin = 1 + } + return window - margin +} + +func (r *Runner) limitRAGContext(ragContext string, history []Message) string { + if strings.TrimSpace(ragContext) == "" { + return "" + } + budget := contextBudget(r.client.Capabilities().MaxContextWindow) + if budget <= 0 { + return "" + } + + historyTokens := totalPromptTokens(history) + fitted := "" + for _, block := range strings.Split(strings.TrimSpace(ragContext), "\n\n") { + block = strings.TrimSpace(block) + if block == "" { + continue + } + candidate := block + if fitted != "" { + candidate = fitted + "\n\n" + block + } + if estimatePromptTokens(botpersona.BuildSystemPrompt(r.systemPrompt, candidate))+historyTokens > budget { + break + } + fitted = candidate + } + return fitted +} + +func (r *Runner) fitHistory(system string, history []Message) ([]Message, error) { + window := r.client.Capabilities().MaxContextWindow + if window <= 0 { + return history, nil + } + budget := contextBudget(window) + if budget <= 0 { + return nil, fmt.Errorf("context window is too small") + } + + history = dropLeadingAssistantMessages(history) + for { + if estimatePromptTokens(system)+totalPromptTokens(history) <= budget { + return history, nil + } + turnStarts := []int{} + for i, message := range history { + if message.Role == RoleUser { + turnStarts = append(turnStarts, i) + } + } + if len(turnStarts) < 2 { + return nil, fmt.Errorf("prompt exceeds context window (%d tokens)", window) + } + first, next := turnStarts[0], turnStarts[1] + kept := make([]Message, 0, len(history)-(next-first)) + kept = append(kept, history[:first]...) + kept = append(kept, history[next:]...) + history = kept + } +} + +func dropLeadingAssistantMessages(history []Message) []Message { + firstUser := -1 + for i, message := range history { + if message.Role == RoleUser { + firstUser = i + break + } + } + if firstUser <= 0 { + return history + } + out := make([]Message, 0, len(history)) + for _, message := range history[:firstUser] { + if message.Role != RoleAssistant { + out = append(out, message) + } + } + return append(out, history[firstUser:]...) +} + +// Compact reduces the older portion of history to a single system-role +// summary message when the previous turn's input tokens exceed the +// configured threshold. Returns the (possibly compacted) message slice — +// unchanged when compaction is disabled, history is too short, the model +// has no reported context window, or summarization fails. +// +// On summarization failure the runner falls back to dropping the oldest +// turns until the remaining slice fits a conservative budget, so a flaky +// summarize call never fails the user's request. +func (r *Runner) Compact(ctx context.Context, history []Message) ([]Message, error) { + r.lastCompact = CompactionStats{} + if !r.compaction.Enabled || len(history) == 0 { + return history, nil + } + + caps := r.client.Capabilities() + if caps.MaxContextWindow <= 0 { + return history, nil + } + + used := r.usage.InputTokens + estimated := estimatePromptTokens(r.systemPrompt) + totalPromptTokens(history) + if estimated > used { + used = estimated + } + window := caps.MaxContextWindow + threshold := int(float64(window) * r.compaction.ThresholdRatio) + if used < threshold { + return history, nil + } + + older, recent := splitByTurns(history, r.compaction.KeepRecentTurns) + if len(older) == 0 { + return history, nil + } + + summary, err := r.summarize(ctx, older) + if err != nil { + slog.Warn("compaction: summarize failed, falling back to truncation", + "err", err, "older_turns", userTurnCount(older)) + // Fall back: drop oldest turns until the slice is small enough to + // fit in (threshold) tokens, using a 4-chars-per-token heuristic. + dropped, kept := truncateToBudget(history, threshold) + r.lastCompact = CompactionStats{ + Happened: true, + OlderTurns: userTurnCount(dropped), + KeptTurns: userTurnCount(kept), + WindowTokens: window, + UsedTokens: used, + } + return kept, nil + } + + out := make([]Message, 0, 1+len(recent)) + out = append(out, Message{ + Role: RoleSystem, + Content: "Earlier conversation summary:\n" + summary, + }) + out = append(out, recent...) + r.lastCompact = CompactionStats{ + Happened: true, + OlderTurns: userTurnCount(older), + KeptTurns: userTurnCount(recent), + SummaryTokens: estimateTokens(summary), + WindowTokens: window, + UsedTokens: used, + } + return out, nil +} + +// summarize calls the LLM (non-streaming) to condense older turns. We use +// a small max_tokens cap so the compaction step itself stays cheap. +func (r *Runner) summarize(ctx context.Context, older []Message) (string, error) { + transcript := renderTranscript(older) + sys := r.compaction.SummarySystemPrompt + if sys == "" { + sys = defaultSummaryPrompt + } + maxTokens := 512 + resp, err := r.client.Generate(ctx, llm.CompletionRequest{ + Messages: []Message{ + {Role: RoleSystem, Content: sys}, + {Role: RoleUser, Content: transcript}, + }, + MaxTokens: &maxTokens, + }) + if err != nil { + return "", err + } + summary := strings.TrimSpace(resp.Content) + if summary == "" { + return "", fmt.Errorf("empty summary") + } + return summary, nil +} + // Stream runs the model and yields each streamed chunk. The caller // forwards chunk.Delta to the SSE stream; chunk.Usage on the last chunk // carries token counts. @@ -122,4 +371,132 @@ func formatHits(hits []portfolio.SearchResult) string { b.WriteString("\n\n") } return b.String() -} \ No newline at end of file +} + +// defaultSummaryPrompt is used when the YAML doesn't override it. The +// instruction explicitly avoids headers / meta-commentary so the summary +// drops in cleanly as a system message and the model treats it as facts +// about the prior conversation rather than instructions from the user. +const defaultSummaryPrompt = "Summarize the following conversation between a user and the assistant concisely but completely. " + + "Preserve decisions made, concrete facts (names, paths, IDs, preferences) and any pending task or context the assistant needs to keep helping. " + + "Respond only with the summary, in the same language as the conversation, with no headers or extra commentary." + +// splitByTurns divides messages into an older portion to summarize and a +// recent tail to keep verbatim. A turn is a user message plus everything +// up to (but not including) the next user message — so a kept turn's tool +// calls, thoughts and response always stay together. Returns older == nil +// when there are not enough turns to bother compacting. +func splitByTurns(history []Message, keepRecentTurns int) (older, recent []Message) { + turnStarts := []int{} + for i, m := range history { + if m.Role == RoleUser { + turnStarts = append(turnStarts, i) + } + } + if len(turnStarts) <= keepRecentTurns { + return nil, history + } + cut := turnStarts[len(turnStarts)-keepRecentTurns] + return history[:cut], history[cut:] +} + +// renderTranscript flattens chat messages into a plain User/Assistant +// transcript for summarization. Tool messages and empty assistant +// placeholders are skipped — neither contributes facts the model needs +// to remember. +func renderTranscript(messages []Message) string { + var b strings.Builder + for _, m := range messages { + switch m.Role { + case RoleUser: + b.WriteString("User: ") + b.WriteString(m.Content) + b.WriteString("\n") + case RoleAssistant: + if m.Content == "" { + continue + } + b.WriteString("Assistant: ") + b.WriteString(m.Content) + b.WriteString("\n") + } + } + return b.String() +} + +// userTurnCount is the number of user-role messages in a slice — a +// convenient proxy for "turns" in stats reporting (each turn starts +// with a user message). +func userTurnCount(messages []Message) int { + n := 0 + for _, m := range messages { + if m.Role == RoleUser { + n++ + } + } + return n +} + +// estimateTokens gives a rough token count for s. ~4 chars per token is +// the usual rule of thumb for English; close enough for Spanish / mixed +// text that this only drives a trigger threshold. +func estimateTokens(s string) int { + if s == "" { + return 0 + } + return len(s)/4 + 1 +} + +func estimatePromptTokens(s string) int { + if s == "" { + return 0 + } + return len(s)/3 + 1 +} + +func totalPromptTokens(messages []Message) int { + total := 0 + for _, m := range messages { + total += estimatePromptTokens(m.Content) + } + return total +} + +// truncateToBudget is the fallback when summarization fails. It drops +// the oldest user-turns one at a time until the remaining slice fits in +// `budget` tokens (conservative heuristic). At least the last user turn is +// always preserved, so we never return an empty slice. +// +// Implementation note: as `cut` increases, history[turnStarts[cut]:] +// shrinks, so `totalPromptTokens(...)` is monotonically non-increasing in +// `cut`. The loop is therefore guaranteed to terminate by either +// fitting the budget or hitting the "keep at least the last turn" floor. +func truncateToBudget(history []Message, budget int) (dropped, kept []Message) { + if len(history) == 0 { + return nil, nil + } + turnStarts := []int{} + for i, m := range history { + if m.Role == RoleUser { + turnStarts = append(turnStarts, i) + } + } + if len(turnStarts) == 0 { + return nil, history + } + // cut = number of leading turns to drop. 0 = keep everything; we + // grow it until the remaining slice fits in budget (or we hit the + // floor of one turn kept). + cut := 0 + for cut < len(turnStarts)-1 { + start := turnStarts[cut+1] + if totalPromptTokens(history[start:]) <= budget { + break + } + cut++ + } + if cut == 0 { + return nil, history + } + return history[:turnStarts[cut]], history[turnStarts[cut]:] +} diff --git a/internal/agent/runner_compaction_test.go b/internal/agent/runner_compaction_test.go new file mode 100644 index 0000000..126513a --- /dev/null +++ b/internal/agent/runner_compaction_test.go @@ -0,0 +1,425 @@ +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)) + } +}