Cyrene_v2/internal/runtime/loop/compactor_test.go

package loop

import (
	"context"
	"errors"
	"reflect"
	"strconv"
	"strings"
	"testing"

	"git.yeij.top/AskaEth/Cyrene_v2/internal/runtime/model"
)

// buildTurns builds a multi-turn conversation: n user/assistant turns followed by
// a trailing tool result, each message padded to around pad chars so the token
// estimate can be pushed over a small budget.
func buildTurns(n, pad int, withTool bool) []model.Message {
	msgs := make([]model.Message, 0, n*2+1)
	for i := 0; i < n; i++ {
		user := model.TextMessage(model.RoleUser, strings.Repeat("u", pad)+" turn-"+strconv.Itoa(i))
		asst := model.TextMessage(model.RoleAssistant, strings.Repeat("a", pad)+" reply-"+strconv.Itoa(i))
		msgs = append(msgs, user, asst)
		if withTool && i == n-1 {
			msgs = append(msgs, model.TextMessage(model.RoleTool, strings.Repeat("r", pad*4)+" tool-result"))
		}
	}
	return msgs
}

func TestCompact_UnderBudget_Unchanged(t *testing.T) {
	t.Parallel()
	msgs := buildTurns(2, 4, false)
	got, err := (defaultCompactor{}).Compact(context.Background(), msgs, 1_000_000)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	if len(got) != len(msgs) {
		t.Fatalf("under budget should be unchanged: got %d messages, want %d", len(got), len(msgs))
	}
	for i := range msgs {
		if !reflect.DeepEqual(got[i], msgs[i]) {
			t.Fatalf("under budget mutated message %d: got %+v want %+v", i, got[i], msgs[i])
		}
	}
}

func TestCompact_OverBudget_PreservesTrailingToolResult(t *testing.T) {
	t.Parallel()
	msgs := buildTurns(3, 400, true)
	c := defaultCompactor{}
	if estimateContextTokens(msgs) <= 1000 {
		t.Fatalf("test precondition: estimated %d tokens, want > 1000", estimateContextTokens(msgs))
	}
	got, err := c.Compact(context.Background(), msgs, 1000)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	if len(got) >= len(msgs) {
		t.Fatalf("compaction should shrink the conversation: got %d, want < %d", len(got), len(msgs))
	}
	// A leading summary system message marks the compaction.
	if got[0].Role != model.RoleSystem || !strings.HasPrefix(messageText(got[0]), summaryMarker) {
		t.Fatalf("first message is not a compaction summary: %+v", got[0])
	}
	// The trailing tool result survives at the very end, unmodified.
	last := got[len(got)-1]
	if last.Role != model.RoleTool || !strings.Contains(last.Content[0].Text, "tool-result") {
		t.Fatalf("trailing tool result was not preserved: %+v", last)
	}
	// Whenever a summary is emitted, the first kept message is a whole-turn
	// boundary (a user message), so compaction never splits a turn.
	if got[1].Role != model.RoleUser {
		t.Fatalf("first kept message role = %q, want a turn-start user message", got[1].Role)
	}
}

func TestCompact_NeverSplitsTurn_MidConversation(t *testing.T) {
	t.Parallel()
	// Two turns of very different sizes: the second turn is huge, the first is
	// small. Compaction must keep the whole huge last turn and summarize the
	// small earlier turn(s), not slice the last turn apart.
	big := strings.Repeat("b", 5000)
	msgs := []model.Message{
		model.TextMessage(model.RoleUser, "question one"),
		model.TextMessage(model.RoleAssistant, "answer one"),
		model.TextMessage(model.RoleUser, "question two"),
		model.TextMessage(model.RoleAssistant, big),
	}
	got, err := (defaultCompactor{}).Compact(context.Background(), msgs, 400)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	if len(got) >= len(msgs) {
		t.Fatalf("expected compaction, got %d >= %d messages", len(got), len(msgs))
	}
	if got[0].Role != model.RoleSystem {
		t.Fatalf("expected a leading summary, got %+v", got[0])
	}
	// The kept region begins at the second user turn and its assistant reply are
	// both intact.
	if got[1].Role != model.RoleUser || got[1].Content[0].Text != "question two" {
		t.Fatalf("kept region should start a whole turn: %+v", got[1])
	}
	if !strings.Contains(got[len(got)-1].Content[0].Text, big) {
		t.Fatalf("last assistant reply was not preserved: %+v", got[len(got)-1])
	}
}

func TestCompact_SummaryCap(t *testing.T) {
	t.Parallel()
	// A huge prefix so the naive summary would try to sample long text, but each
	// sampled message is truncated; verify the capped length stays bounded.
	msgs := buildTurns(8, 5000, false)
	got, err := (defaultCompactor{}).Compact(context.Background(), msgs, 5000)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	if got[0].Role != model.RoleSystem {
		t.Fatalf("expected a leading summary")
	}
	if n := len([]rune(got[0].Content[0].Text)); n > maxSummaryChars+len(summaryMarker)+2 {
		t.Fatalf("summary too long: %d runes, want <= %d", n, maxSummaryChars+len(summaryMarker)+2)
	}
}

func TestCompact_ExistingSummaryMerged(t *testing.T) {
	t.Parallel()
	old := summaryMarker + "\n\n早前摘要内容"
	msgs := append([]model.Message{model.TextMessage(model.RoleSystem, old)}, buildTurns(4, 300, false)...)
	got, err := (defaultCompactor{}).Compact(context.Background(), msgs, 1000)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	// Only one leading summary remains and it merges the prior summary.
	if got[0].Role != model.RoleSystem {
		t.Fatalf("expected a leading summary")
	}
	text := messageText(got[0])
	if !strings.Contains(text, "早前摘要内容") {
		t.Fatalf("existing summary was dropped: %q", text)
	}
	if summaryPrefixEnd(got) != 1 {
		t.Fatalf("expected exactly one leading summary message")
	}
}

func TestCompact_CanceledContext_ReturnsError(t *testing.T) {
	t.Parallel()
	msgs := buildTurns(4, 300, false)
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_, err := (defaultCompactor{}).Compact(ctx, msgs, 1000)
	if !errors.Is(err, context.Canceled) {
		t.Fatalf("Compact error = %v, want wrapped context.Canceled", err)
	}
}

func TestEstimateContextTokens_Monotonic(t *testing.T) {
	t.Parallel()
	small := []model.Message{model.TextMessage(model.RoleUser, "hi")}
	big := []model.Message{model.TextMessage(model.RoleUser, strings.Repeat("x", 1000))}
	if estimateContextTokens(big) <= estimateContextTokens(small) {
		t.Fatalf("token estimate not monotonic: big=%d small=%d", estimateContextTokens(big), estimateContextTokens(small))
	}
}