first commit

pigo/internal/compaction/tokens_test.go

@@ -0,0 +1,189 @@
+package compaction
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+)
+
+func userMsg(text string) agentcore.UserMessage {
+	return agentcore.UserMessage{
+		RoleField: agentcore.RoleUser,
+		Content:   agentcore.ContentList{agentcore.NewTextContent(text)},
+	}
+}
+
+func assistantMsg(text string, usage *agentcore.Usage, stop string) agentcore.AssistantMessage {
+	return agentcore.AssistantMessage{
+		RoleField:  agentcore.RoleAssistant,
+		Content:    agentcore.ContentList{agentcore.NewTextContent(text)},
+		Usage:      usage,
+		StopReason: stop,
+	}
+}
+
+func TestEstimateTokensText(t *testing.T) {
+	// 8 chars / 4 = 2 tokens.
+	if got := EstimateTokens(userMsg("abcdefgh")); got != 2 {
+		t.Fatalf("text estimate: got %d, want 2", got)
+	}
+	// 9 chars -> ceil(9/4) = 3.
+	if got := EstimateTokens(userMsg("abcdefghi")); got != 3 {
+		t.Fatalf("ceil estimate: got %d, want 3", got)
+	}
+	// empty -> 0.
+	if got := EstimateTokens(userMsg("")); got != 0 {
+		t.Fatalf("empty estimate: got %d, want 0", got)
+	}
+}
+
+func TestEstimateTokensImage(t *testing.T) {
+	m := agentcore.UserMessage{
+		RoleField: agentcore.RoleUser,
+		Content:   agentcore.ContentList{agentcore.NewImageContent("data", "image/png")},
+	}
+	// estimatedImageChars (4800) / 4 = 1200.
+	if got := EstimateTokens(m); got != 1200 {
+		t.Fatalf("image estimate: got %d, want 1200", got)
+	}
+}
+
+func TestEstimateTokensAssistantToolCall(t *testing.T) {
+	args := json.RawMessage(`{"path":"a.go"}`) // 15 chars
+	m := agentcore.AssistantMessage{
+		RoleField: agentcore.RoleAssistant,
+		Content: agentcore.ContentList{
+			agentcore.NewToolCallContent("id1", "read", args), // name "read" = 4 chars
+		},
+	}
+	// (4 + 15) / 4 = ceil(19/4) = 5.
+	if got := EstimateTokens(m); got != 5 {
+		t.Fatalf("toolcall estimate: got %d, want 5", got)
+	}
+}
+
+func TestEstimateTokensUnknownRoleZero(t *testing.T) {
+	if got := EstimateTokens(agentcore.ToolResultMessage{}); got != 0 {
+		t.Fatalf("empty tool result: got %d, want 0", got)
+	}
+}
+
+func TestEstimateContextTokensNoUsageFallsBackToEstimate(t *testing.T) {
+	msgs := []agentcore.Message{
+		userMsg("abcdefgh"),                   // 2
+		assistantMsg("abcd", nil, "end_turn"), // 1
+	}
+	est := EstimateContextTokens(msgs)
+	if est.LastUsageIndex != -1 {
+		t.Fatalf("LastUsageIndex: got %d, want -1", est.LastUsageIndex)
+	}
+	if est.Tokens != 3 {
+		t.Fatalf("Tokens: got %d, want 3", est.Tokens)
+	}
+	if est.TrailingTokens != 3 || est.UsageTokens != 0 {
+		t.Fatalf("trailing/usage: got %d/%d, want 3/0", est.TrailingTokens, est.UsageTokens)
+	}
+}
+
+func TestEstimateContextTokensPrefersUsage(t *testing.T) {
+	usage := &agentcore.Usage{InputTokens: 1000, OutputTokens: 200}
+	msgs := []agentcore.Message{
+		userMsg("first"),
+		assistantMsg("reply", usage, "end_turn"),
+		userMsg("abcdefgh"), // trailing: 2 tokens
+	}
+	est := EstimateContextTokens(msgs)
+	if est.LastUsageIndex != 1 {
+		t.Fatalf("LastUsageIndex: got %d, want 1", est.LastUsageIndex)
+	}
+	if est.UsageTokens != 1200 {
+		t.Fatalf("UsageTokens: got %d, want 1200", est.UsageTokens)
+	}
+	if est.TrailingTokens != 2 {
+		t.Fatalf("TrailingTokens: got %d, want 2", est.TrailingTokens)
+	}
+	if est.Tokens != 1202 {
+		t.Fatalf("Tokens: got %d, want 1202", est.Tokens)
+	}
+}
+
+func TestEstimateContextTokensSkipsAbortedAndErrorUsage(t *testing.T) {
+	good := &agentcore.Usage{InputTokens: 500, OutputTokens: 0}
+	bad := &agentcore.Usage{InputTokens: 9999, OutputTokens: 0}
+	msgs := []agentcore.Message{
+		assistantMsg("ok", good, "end_turn"),
+		assistantMsg("aborted", bad, agentcore.StopReasonAborted),
+		assistantMsg("errored", bad, agentcore.StopReasonError),
+	}
+	est := EstimateContextTokens(msgs)
+	if est.LastUsageIndex != 0 {
+		t.Fatalf("LastUsageIndex: got %d, want 0 (should skip aborted/error)", est.LastUsageIndex)
+	}
+	if est.UsageTokens != 500 {
+		t.Fatalf("UsageTokens: got %d, want 500", est.UsageTokens)
+	}
+}
+
+func TestEstimateContextTokensSkipsZeroUsage(t *testing.T) {
+	zero := &agentcore.Usage{InputTokens: 0, OutputTokens: 0}
+	msgs := []agentcore.Message{
+		userMsg("abcdefgh"), // 2
+		assistantMsg("reply", zero, "end_turn"),
+	}
+	est := EstimateContextTokens(msgs)
+	// zero usage is ignored, so falls back to full estimation.
+	if est.LastUsageIndex != -1 {
+		t.Fatalf("LastUsageIndex: got %d, want -1", est.LastUsageIndex)
+	}
+}
+
+func TestShouldCompactThresholdBoundaries(t *testing.T) {
+	s := CompactionSettings{Enabled: true, ReserveTokens: 16384}
+	window := 200000
+	usable := window - s.ReserveTokens // 183616
+
+	tests := []struct {
+		name          string
+		contextTokens int
+		want          bool
+	}{
+		{"far below", 1000, false},
+		{"equal to usable", usable, false}, // strictly greater required
+		{"one over usable", usable + 1, true},
+		{"far over", window * 2, true},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := ShouldCompact(tt.contextTokens, window, s); got != tt.want {
+				t.Fatalf("ShouldCompact(%d): got %v, want %v", tt.contextTokens, got, tt.want)
+			}
+		})
+	}
+}
+
+func TestShouldCompactDisabled(t *testing.T) {
+	s := CompactionSettings{Enabled: false, ReserveTokens: 16384}
+	if ShouldCompact(1_000_000, 200000, s) {
+		t.Fatal("disabled settings must never compact")
+	}
+}
+
+func TestShouldCompactUnknownWindow(t *testing.T) {
+	s := CompactionSettings{Enabled: true, ReserveTokens: 16384}
+	if ShouldCompact(1_000_000, 0, s) {
+		t.Fatal("unknown (0) context window must never compact")
+	}
+}
+
+func TestDefaultCompactionSettings(t *testing.T) {
+	if DefaultCompactionSettings.ReserveTokens != 16384 {
+		t.Fatalf("ReserveTokens: got %d, want 16384", DefaultCompactionSettings.ReserveTokens)
+	}
+	if DefaultCompactionSettings.KeepRecentTokens != 20000 {
+		t.Fatalf("KeepRecentTokens: got %d, want 20000", DefaultCompactionSettings.KeepRecentTokens)
+	}
+	if !DefaultCompactionSettings.Enabled {
+		t.Fatal("default settings should be enabled")
+	}
+}