goclaw/internal/agent/pruning_test.go

package agent

import (
	"strings"
	"testing"

	"github.com/nextlevelbuilder/goclaw/internal/config"
	"github.com/nextlevelbuilder/goclaw/internal/providers"
)

// ─── resolvePruningSettings ───────────────────────────────────────────────

func TestResolvePruningSettings_Defaults(t *testing.T) {
	s := resolvePruningSettings(nil)
	if s.keepLastAssistants != defaultKeepLastAssistants {
		t.Errorf("keepLastAssistants = %d, want %d", s.keepLastAssistants, defaultKeepLastAssistants)
	}
	if s.softTrimRatio != defaultSoftTrimRatio {
		t.Errorf("softTrimRatio = %v, want %v", s.softTrimRatio, defaultSoftTrimRatio)
	}
	if s.hardClearRatio != defaultHardClearRatio {
		t.Errorf("hardClearRatio = %v, want %v", s.hardClearRatio, defaultHardClearRatio)
	}
	if !s.hardClearEnabled {
		t.Error("hardClearEnabled should be true by default")
	}
	if s.hardClearPlaceholder != defaultHardClearPlaceholder {
		t.Errorf("placeholder = %q, want %q", s.hardClearPlaceholder, defaultHardClearPlaceholder)
	}
}

func TestResolvePruningSettings_CustomValues(t *testing.T) {
	enabled := false
	cfg := &config.ContextPruningConfig{
		KeepLastAssistants:   5,
		SoftTrimRatio:        0.4,
		HardClearRatio:       0.6,
		MinPrunableToolChars: 10000,
		SoftTrim: &config.ContextPruningSoftTrim{
			MaxChars:  2000,
			HeadChars: 800,
			TailChars: 800,
		},
		HardClear: &config.ContextPruningHardClear{
			Enabled:     &enabled,
			Placeholder: "[gone]",
		},
	}
	s := resolvePruningSettings(cfg)
	if s.keepLastAssistants != 5 {
		t.Errorf("keepLastAssistants = %d, want 5", s.keepLastAssistants)
	}
	if s.softTrimRatio != 0.4 {
		t.Errorf("softTrimRatio = %v, want 0.4", s.softTrimRatio)
	}
	if s.hardClearEnabled {
		t.Error("hardClearEnabled should be false")
	}
	if s.hardClearPlaceholder != "[gone]" {
		t.Errorf("placeholder = %q, want [gone]", s.hardClearPlaceholder)
	}
	if s.softTrimMaxChars != 2000 {
		t.Errorf("softTrimMaxChars = %d, want 2000", s.softTrimMaxChars)
	}
}

func TestResolvePruningSettings_InvalidRatiosIgnored(t *testing.T) {
	// Ratios out of (0,1] range should be ignored, defaults kept.
	cfg := &config.ContextPruningConfig{
		SoftTrimRatio:  1.5, // invalid
		HardClearRatio: 0,   // invalid (0 not > 0)
	}
	s := resolvePruningSettings(cfg)
	if s.softTrimRatio != defaultSoftTrimRatio {
		t.Errorf("invalid SoftTrimRatio should keep default, got %v", s.softTrimRatio)
	}
	if s.hardClearRatio != defaultHardClearRatio {
		t.Errorf("invalid HardClearRatio should keep default, got %v", s.hardClearRatio)
	}
}

// ─── findAssistantCutoff ──────────────────────────────────────────────────

func TestFindAssistantCutoff_FewerThanKeepLast(t *testing.T) {
	msgs := []providers.Message{
		{Role: "user", Content: "hi"},
		{Role: "assistant", Content: "hello"},
	}
	// keepLast=3 but only 1 assistant → returns -1 (not enough to protect)
	idx := findAssistantCutoff(msgs, 3)
	if idx != -1 {
		t.Errorf("cutoff = %d, want -1", idx)
	}
}

func TestFindAssistantCutoff_ExactlyKeepLast(t *testing.T) {
	msgs := []providers.Message{
		{Role: "user", Content: "hi"},
		{Role: "assistant", Content: "a1"},
		{Role: "user", Content: "next"},
		{Role: "assistant", Content: "a2"},
		{Role: "user", Content: "end"},
		{Role: "assistant", Content: "a3"},
	}
	// keepLast=3: protect all 3 assistant messages → cutoff at index of first assistant
	idx := findAssistantCutoff(msgs, 3)
	if idx != 1 {
		t.Errorf("cutoff = %d, want 1", idx)
	}
}

func TestFindAssistantCutoff_MoreThanKeepLast(t *testing.T) {
	// 5 assistants, keepLast=2 → cutoff at 3rd-from-last assistant
	msgs := []providers.Message{
		{Role: "assistant", Content: "a0"}, // idx 0
		{Role: "assistant", Content: "a1"}, // idx 1
		{Role: "assistant", Content: "a2"}, // idx 2
		{Role: "assistant", Content: "a3"}, // idx 3
		{Role: "assistant", Content: "a4"}, // idx 4
	}
	idx := findAssistantCutoff(msgs, 2)
	// 2nd from last = idx 3
	if idx != 3 {
		t.Errorf("cutoff = %d, want 3", idx)
	}
}

func TestFindAssistantCutoff_ZeroKeepLast(t *testing.T) {
	msgs := []providers.Message{{Role: "assistant", Content: "x"}}
	// keepLast=0 → return len(msgs) (protect nothing)
	idx := findAssistantCutoff(msgs, 0)
	if idx != len(msgs) {
		t.Errorf("cutoff = %d, want %d", idx, len(msgs))
	}
}

// ─── takeHead / takeTail ──────────────────────────────────────────────────

func TestTakeHead(t *testing.T) {
	cases := []struct {
		input string
		n     int
		want  string
	}{
		{"hello", 3, "hel"},
		{"hello", 10, "hello"},
		{"hello", 0, ""},
		{"αβγδε", 3, "αβγ"}, // unicode
	}
	for _, c := range cases {
		got := takeHead(c.input, c.n)
		if got != c.want {
			t.Errorf("takeHead(%q, %d) = %q, want %q", c.input, c.n, got, c.want)
		}
	}
}

func TestTakeTail(t *testing.T) {
	cases := []struct {
		input string
		n     int
		want  string
	}{
		{"hello", 3, "llo"},
		{"hello", 10, "hello"},
		{"hello", 0, ""},
		{"αβγδε", 2, "δε"}, // unicode
	}
	for _, c := range cases {
		got := takeTail(c.input, c.n)
		if got != c.want {
			t.Errorf("takeTail(%q, %d) = %q, want %q", c.input, c.n, got, c.want)
		}
	}
}

// ─── estimateMessageChars ─────────────────────────────────────────────────

func TestEstimateMessageChars(t *testing.T) {
	m := providers.Message{Content: "αβγ"} // 3 runes, 6 bytes
	if got := estimateMessageChars(m); got != 3 {
		t.Errorf("estimateMessageChars = %d, want 3", got)
	}
}

// ─── pruneContextMessages — mode off ────────────────────────────────────

func TestPruneContextMessages_ModeOff_ReturnsOriginal(t *testing.T) {
	msgs := []providers.Message{
		{Role: "user", Content: "hi"},
		{Role: "assistant", Content: "hello"},
	}
	cfg := &config.ContextPruningConfig{Mode: "off"}
	got := pruneContextMessages(msgs, 100000, cfg)
	if len(got) != len(msgs) {
		t.Error("mode=off should return original slice")
	}
}

func TestPruneContextMessages_ZeroWindow_ReturnsOriginal(t *testing.T) {
	msgs := []providers.Message{{Role: "user", Content: "hi"}}
	got := pruneContextMessages(msgs, 0, nil)
	if len(got) != len(msgs) {
		t.Error("zero context window should return original")
	}
}

func TestPruneContextMessages_SmallContext_NoChange(t *testing.T) {
	// Total content << softTrimRatio → no pruning.
	msgs := []providers.Message{
		{Role: "user", Content: "short"},
		{Role: "assistant", Content: "reply", ToolCalls: nil},
		{Role: "user", Content: "another"},
		{Role: "assistant", Content: "done"},
	}
	// Very large window → ratio tiny → nothing pruned.
	got := pruneContextMessages(msgs, 1_000_000, nil)
	if len(got) != len(msgs) {
		t.Errorf("small context: len=%d, want %d", len(got), len(msgs))
	}
}

func TestPruneContextMessages_SoftTrim_LongToolResult(t *testing.T) {
	// Build scenario: 4 assistant messages (so keepLast=3 protects last 3),
	// and an old tool result that is very long.
	toolContent := strings.Repeat("X", 10000) // long

	msgs := []providers.Message{
		{Role: "user", Content: "step1"},
		{Role: "assistant", Content: "a0", ToolCalls: []providers.ToolCall{{ID: "tc1", Name: "t"}}},
		{Role: "tool", Content: toolContent, ToolCallID: "tc1"},
		{Role: "user", Content: "step2"},
		{Role: "assistant", Content: "a1"},
		{Role: "user", Content: "step3"},
		{Role: "assistant", Content: "a2"},
		{Role: "user", Content: "step4"},
		{Role: "assistant", Content: "a3"},
	}

	cfg := &config.ContextPruningConfig{
		SoftTrimRatio: 0.01, // threshold very low → prune immediately
		HardClearRatio: 0.99, // don't hard clear
		SoftTrim: &config.ContextPruningSoftTrim{
			MaxChars:  100,
			HeadChars: 50,
			TailChars: 50,
		},
	}
	// Use a context window that makes ratio exceed threshold.
	got := pruneContextMessages(msgs, 100, cfg)

	// Tool result should be trimmed — shorter than original.
	for _, m := range got {
		if m.Role == "tool" {
			if len(m.Content) >= len(toolContent) {
				t.Error("expected tool result to be trimmed")
			}
			return
		}
	}
	// If tool message not found, that's also acceptable (hard-cleared).
}

func TestPruneContextMessages_HardClear_WhenRatioVeryHigh(t *testing.T) {
	// Set up context that forces hard clear.
	toolContent := strings.Repeat("Y", 5000)

	msgs := []providers.Message{
		{Role: "user", Content: "q"},
		{Role: "assistant", Content: "a0", ToolCalls: []providers.ToolCall{{ID: "tc1", Name: "t"}}},
		{Role: "tool", Content: toolContent, ToolCallID: "tc1"},
		{Role: "user", Content: "q2"},
		{Role: "assistant", Content: "a1"},
		{Role: "user", Content: "q3"},
		{Role: "assistant", Content: "a2"},
		{Role: "user", Content: "q4"},
		{Role: "assistant", Content: "a3"},
	}

	enabled := true
	cfg := &config.ContextPruningConfig{
		SoftTrimRatio:        0.001,
		HardClearRatio:       0.001,
		MinPrunableToolChars: 1, // very low threshold
		SoftTrim: &config.ContextPruningSoftTrim{
			MaxChars:  100,
			HeadChars: 50,
			TailChars: 50,
		},
		HardClear: &config.ContextPruningHardClear{
			Enabled:     &enabled,
			Placeholder: "[cleared]",
		},
	}

	got := pruneContextMessages(msgs, 10, cfg)

	// Tool result should be replaced by placeholder or heavily trimmed.
	for _, m := range got {
		if m.Role == "tool" {
			if len(m.Content) >= len(toolContent) {
				t.Error("tool content should have been reduced")
			}
		}
	}
}

// ─── hasImportantTail ────────────────────────────────────────────────────

func TestHasImportantTail_MatchesErrorKeyword(t *testing.T) {
	content := strings.Repeat("x", 1000) + " error occurred"
	if !hasImportantTail(content) {
		t.Error("expected important tail (error keyword)")
	}
}

func TestHasImportantTail_ShortContentNoMatch(t *testing.T) {
	// No keywords → false
	if hasImportantTail("no important content here at all") {
		t.Error("expected no important tail")
	}
}