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 构建一段多轮对话：n 个 user/assistant 轮，后接
// 一个尾部工具结果，每条消息填充到约 pad 个字符，以便把 token
// 估算推过一个小预算。
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))
	}
	// 一条前导摘要系统消息标记该压缩。
	if got[0].Role != model.RoleSystem || !strings.HasPrefix(messageText(got[0]), summaryMarker) {
		t.Fatalf("first message is not a compaction summary: %+v", got[0])
	}
	// 尾部工具结果原样存活在最后。
	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)
	}
	// 只要发出摘要，第一条保留的消息就是一个整 turn
	// 边界（用户消息），因此压缩绝不切开一个 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()
	// 两个大小悬殊的 turn：第二个极大，第一个
	// 很小。压缩必须保留整个巨大的最后一个 turn，并摘要
	// 更早的小 turn，而不是把最后一个 turn 切开。
	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])
	}
	// 保留区域从第二个用户 turn 开始，其 assistant 回复
	// 都完整无损。
	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()
	// 一个巨大的前缀，使朴素摘要会试图抽样长文本，但每条
	// 被抽样的消息都被截断；验证被限额后的长度保持有界。
	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)
	}
	// 只剩一条前导摘要，且它合并了此前的摘要。
	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))
	}
}