first commit

2026-08-14 23:41:57 +08:00
commit 086803a8dd
471 changed files with 91938 additions and 0 deletions
@@ -0,0 +1,194 @@
+package runtime
+
+// Tests for context rebuild (#482): with a checkpoint present, RebuildFromCheckpoint
+// inserts the compression boundary at the watermark — collapsing the pre-watermark
+// prefix into the checkpoint summary and preserving the recent tail verbatim; with
+// no checkpoint it falls back to the lossy compaction path (compaction.Compact).
+// Filesystem access uses t.TempDir(), matching checkpoint_test.go; the summary
+// model is the shared summaryStream fake from compaction_test.go.
+
+import (
+	"context"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/smallnest/pigo/internal/agentcore"
+	"github.com/smallnest/pigo/internal/compaction"
+)
+
+// textUser builds a user message carrying body, used to seed a rebuildable history.
+func textUser(body string) agentcore.UserMessage {
+	return agentcore.UserMessage{
+		RoleField: agentcore.RoleUser,
+		Content:   agentcore.ContentList{agentcore.NewTextContent(body)},
+	}
+}
+
+func TestRebuildFromCheckpoint_InsertsBoundary(t *testing.T) {
+	root := t.TempDir()
+	const sessionID = "sess-rebuild"
+
+	// A 6-message history; the checkpoint collapses the first 4 into a summary and
+	// keeps messages [4:] verbatim.
+	msgs := agentcore.MessageList{
+		textUser("m0"), textUser("m1"), textUser("m2"),
+		textUser("m3"), textUser("keep-A"), textUser("keep-B"),
+	}
+	cp := Checkpoint{
+		Watermark:       4,
+		Summary:         "## Goal\ndistilled prefix",
+		CreatedAt:       time.Now().UTC(),
+		CoveredMessages: 4,
+	}
+	if err := WriteCheckpoint(sessionID, root, cp); err != nil {
+		t.Fatalf("WriteCheckpoint: %v", err)
+	}
+
+	// No summarization stream is needed: the checkpoint path is pure and local.
+	cfg := newRunCfg(nil)
+	res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil)
+	if err != nil {
+		t.Fatalf("RebuildFromCheckpoint: %v", err)
+	}
+	if !res.FromCheckpoint {
+		t.Fatalf("expected FromCheckpoint=true")
+	}
+	if res.NoOp {
+		t.Fatalf("expected a real rebuild, got NoOp")
+	}
+	if res.Watermark != 4 || res.SummarizedCount != 4 {
+		t.Fatalf("watermark/summarized: got %d/%d, want 4/4", res.Watermark, res.SummarizedCount)
+	}
+	// Rebuilt list = 1 compaction message + the retained tail (2 messages).
+	if len(res.Messages) != 3 {
+		t.Fatalf("rebuilt length: got %d, want 3: %+v", len(res.Messages), res.Messages)
+	}
+	if res.KeptCount != 2 {
+		t.Fatalf("kept: got %d, want 2", res.KeptCount)
+	}
+	// The prefix must collapse into a single compaction message carrying the summary.
+	head, ok := res.Messages[0].(agentcore.CompactionMessage)
+	if !ok {
+		t.Fatalf("message[0] should be a compaction checkpoint, got %T", res.Messages[0])
+	}
+	if !strings.Contains(head.Summary, "distilled prefix") {
+		t.Fatalf("summary not carried through: %q", head.Summary)
+	}
+	// The recent tail is preserved verbatim, in order.
+	for i, want := range []string{"keep-A", "keep-B"} {
+		um, ok := res.Messages[i+1].(agentcore.UserMessage)
+		if !ok {
+			t.Fatalf("message[%d] should be preserved user message, got %T", i+1, res.Messages[i+1])
+		}
+		if got := agentcore.ContentToText(um.Content); got != want {
+			t.Fatalf("tail[%d]: got %q, want %q", i, got, want)
+		}
+	}
+}
+
+func TestRebuildFromCheckpoint_ClampsStaleWatermark(t *testing.T) {
+	root := t.TempDir()
+	const sessionID = "sess-stale"
+
+	msgs := agentcore.MessageList{textUser("a"), textUser("b")}
+	// Watermark past the end (history was re-compacted since the checkpoint).
+	cp := Checkpoint{Watermark: 99, Summary: "old summary", CreatedAt: time.Now().UTC()}
+	if err := WriteCheckpoint(sessionID, root, cp); err != nil {
+		t.Fatalf("WriteCheckpoint: %v", err)
+	}
+
+	cfg := newRunCfg(nil)
+	res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil)
+	if err != nil {
+		t.Fatalf("RebuildFromCheckpoint: %v", err)
+	}
+	// Clamped to len(msgs)=2: everything collapses, no verbatim tail, only the head.
+	if res.Watermark != 2 || res.KeptCount != 0 {
+		t.Fatalf("clamp: watermark=%d kept=%d, want 2/0", res.Watermark, res.KeptCount)
+	}
+	if len(res.Messages) != 1 {
+		t.Fatalf("rebuilt length: got %d, want 1 (summary only)", len(res.Messages))
+	}
+}
+
+func TestRebuildFromCheckpoint_WaitCallbackInvoked(t *testing.T) {
+	root := t.TempDir()
+	const sessionID = "sess-wait"
+	cp := Checkpoint{Watermark: 0, Summary: "s", CreatedAt: time.Now().UTC()}
+	if err := WriteCheckpoint(sessionID, root, cp); err != nil {
+		t.Fatalf("WriteCheckpoint: %v", err)
+	}
+
+	waited := false
+	cfg := newRunCfg(nil)
+	_, err := RebuildFromCheckpoint(context.Background(), agentcore.MessageList{textUser("x")}, sessionID, root, &cfg, func() { waited = true })
+	if err != nil {
+		t.Fatalf("RebuildFromCheckpoint: %v", err)
+	}
+	if !waited {
+		t.Fatalf("waitForCheckpoint callback was not invoked before reading")
+	}
+}
+
+func TestRebuildFromCheckpoint_FallsBackToCompaction(t *testing.T) {
+	root := t.TempDir() // empty: no checkpoint on disk
+	const sessionID = "sess-nocp"
+
+	// Seed a long history so FindCutPoint leaves a summarizable prefix.
+	msgs := bigUserMessages(12, 800)
+
+	cfg := newRunCfg(scriptedStream(nil))
+	cfg.SummaryStream = summaryStream("## Goal\nfallback compaction summary")
+	cfg.ContextWindow = 2000
+	cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100}
+
+	res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil)
+	if err != nil {
+		t.Fatalf("RebuildFromCheckpoint: %v", err)
+	}
+	if res.FromCheckpoint {
+		t.Fatalf("expected fallback (FromCheckpoint=false) when no checkpoint exists")
+	}
+	if res.NoOp {
+		t.Fatalf("expected a real compaction fallback, got NoOp")
+	}
+	if res.SummarizedCount <= 0 {
+		t.Fatalf("expected some messages summarized, got %d", res.SummarizedCount)
+	}
+	if res.TokensAfter >= res.TokensBefore {
+		t.Fatalf("fallback should reduce tokens: before=%d after=%d", res.TokensBefore, res.TokensAfter)
+	}
+	// The rebuilt context begins with a compaction checkpoint holding the summary.
+	head, ok := res.Messages[0].(agentcore.CompactionMessage)
+	if !ok {
+		t.Fatalf("message[0] should be a compaction checkpoint, got %T", res.Messages[0])
+	}
+	if !strings.Contains(head.Summary, "fallback compaction summary") {
+		t.Fatalf("fallback summary not carried through: %q", head.Summary)
+	}
+}
+
+func TestRebuildFromCheckpoint_FallbackNoOpWhenNothingToCompact(t *testing.T) {
+	root := t.TempDir() // no checkpoint
+	const sessionID = "sess-empty"
+
+	// A single short message: no valid cut point leaves a summarization range, so
+	// Compact returns (nil, nil) and the rebuild is a no-op.
+	msgs := agentcore.MessageList{textUser("only")}
+	cfg := newRunCfg(scriptedStream(nil))
+	cfg.SummaryStream = summaryStream("unused")
+	cfg.ContextWindow = 2000
+	cfg.Compaction = compaction.CompactionSettings{Enabled: true, ReserveTokens: 500, KeepRecentTokens: 100}
+
+	res, err := RebuildFromCheckpoint(context.Background(), msgs, sessionID, root, &cfg, nil)
+	if err != nil {
+		t.Fatalf("RebuildFromCheckpoint: %v", err)
+	}
+	if !res.NoOp {
+		t.Fatalf("expected NoOp when there is nothing to compact")
+	}
+	if len(res.Messages) != 1 || res.TokensAfter != res.TokensBefore {
+		t.Fatalf("no-op should return the original context unchanged: %+v", res)
+	}
+}