Cyrene_v2/internal/runtime/loop/assemble_test.go

package loop

import (
	"context"
	"strings"
	"testing"

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

// TestAssembleContext_ReRetrievesEachRound proves memory is re-retrieved every
// round rather than cached from the first turn: an empty store yields no memory
// block, then after seeding the store a later round picks it up.
func TestAssembleContext_ReRetrievesEachRound(t *testing.T) {
	t.Parallel()
	store := newTestMemoryStore(t)
	ret := memory.NewRetriever(store)
	l := New(Options{Retriever: ret, PersonaID: "persona-a"})

	msgs := []model.Message{model.TextMessage(model.RoleUser, "do the thing")}

	// Round 1: the store is empty, so no memory is injected.
	if view := l.assembleContext(context.Background(), msgs, 1); len(view) != 1 {
		t.Fatalf("round 1 view = %d messages, want 1 (not memory)", len(view))
	}

	// Seed a memory as if it was learned between rounds.
	if _, err := store.Add(context.Background(), memory.MemoryItem{
		PersonaID: "persona-a", Content: "她喜欢冰美式",
		Scope: memory.ScopePersona, Kind: memory.KindRelationship,
	}); err != nil {
		t.Fatalf("seed memory: %v", err)
	}

	// Round 2: the retrieval is re-run and the memory now appears.
	view := l.assembleContext(context.Background(), msgs, 2)
	if len(view) != 2 {
		t.Fatalf("round 2 view = %d messages, want 2 (memory injected)", len(view))
	}
	if view[0].Role != model.RoleSystem || !strings.Contains(view[0].Content[0].Text, "冰美式") {
		t.Fatalf("round 2 memory not injected: %+v", view[0])
	}
}

// TestCompactIfNeeded_OverBudget_Compacts drives a multi-turn conversation over a
// small budget and verifies the loop compacts it into a leading summary.
func TestCompactIfNeeded_OverBudget_Compacts(t *testing.T) {
	t.Parallel()
	l := New(Options{MaxContextTokens: 1000})
	big := buildTurns(6, 600, false)
	if estimateContextTokens(big) <= 1000 {
		t.Fatalf("test precondition: want > 1000 tokens, got %d", estimateContextTokens(big))
	}
	got := l.compactIfNeeded(context.Background(), big)
	if len(got) >= len(big) {
		t.Fatalf("compactIfNeeded did not shrink: got %d >= %d", len(got), len(big))
	}
	if got[0].Role != model.RoleSystem || !strings.HasPrefix(messageText(got[0]), summaryMarker) {
		t.Fatalf("expected a leading compaction summary, got %+v", got[0])
	}
}

// TestCompactIfNeeded_Disabled_NoOp keeps the pre-existing behavior when
// MaxContextTokens is left at its zero value.
func TestCompactIfNeeded_Disabled_NoOp(t *testing.T) {
	t.Parallel()
	l := New(Options{}) // MaxContextTokens == 0 → disabled
	big := buildTurns(6, 300, false)
	got := l.compactIfNeeded(context.Background(), big)
	if len(got) != len(big) {
		t.Fatalf("disabled compaction should be a no-op: got %d, want %d", len(got), len(big))
	}
}

// TestRun_TransformContext_InvokedEachRound verifies the TransformContext hook is
// called once per model step (per round), bridging per-round assemble into the
// Run path.
func TestRun_TransformContext_InvokedEachRound(t *testing.T) {
	t.Parallel()
	call := `{"tool_calls":[{"name":"echo","args":{"q":"x"}}]}`
	ad := &mockAdapter{chat: []mgResp{
		{resp: model.Response{Content: call}},
		{resp: model.Response{Content: "done"}},
	}}
	td := dispatcherWith(t, &mockTool{name: "echo", result: "hi"})
	a := newTestAgent(t, ad, td)

	var rounds []int
	var sizes []int
	l := New(Options{
		Agent: a,
		TransformContext: func(_ context.Context, msgs []model.Message, round int) []model.Message {
			rounds = append(rounds, round)
			sizes = append(sizes, len(msgs))
			return msgs
		},
		MaxRounds: 8,
	})

	if _, err := l.Run(context.Background(), "sess", []model.Message{model.TextMessage(model.RoleUser, "go")}); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(rounds) != 2 {
		t.Fatalf("TransformContext called %d times, want 2 (one per round)", len(rounds))
	}
	if rounds[0] != 1 || rounds[1] != 2 {
		t.Fatalf("rounds passed to TransformContext = %v, want [1 2]", rounds)
	}
	// Round 1 has just the user message; round 2 has the fed-back tool result.
	if sizes[0] != 1 || sizes[1] != 2 {
		t.Fatalf("sizes = %v, want [1 2]", sizes)
	}
}

// TestRun_ContextOverflow_CompactsHistory passes a long multi-turn history and a
// small budget into Run; the first model step should see a compacted conversation
// carrying a leading summary.
func TestRun_ContextOverflow_CompactsHistory(t *testing.T) {
	t.Parallel()
	ad := &mockAdapter{chat: []mgResp{{resp: model.Response{Content: "final answer"}}}}
	a := newTestAgent(t, ad, nil)
	l := New(Options{Agent: a, MaxContextTokens: 600, MaxRounds: 8})

	history := buildTurns(8, 200, false)
	if estimateContextTokens(history) <= 600 {
		t.Fatalf("test precondition: want > 600 tokens, got %d", estimateContextTokens(history))
	}
	if _, err := l.Run(context.Background(), "sess", history); err != nil {
		t.Fatalf("Run: %v", err)
	}
	if len(ad.reqs) == 0 {
		t.Fatal("no request captured")
	}
	found := false
	for _, m := range ad.reqs[0].Messages {
		if m.Role == model.RoleSystem && strings.HasPrefix(messageText(m), summaryMarker) {
			found = true
			break
		}
	}
	if !found {
		t.Fatalf("expected a compaction summary in the first request, got %+v", ad.reqs[0].Messages)
	}
}