1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366

// The tests build histories by hand, the way TestCommandList does in
// internal/engine, and check the one property a compression must never break:
// the next request is still a valid conversation — system prompt first, every
// tool response under the request that asked for it, the recent turns intact.
//
// No network and no model: the Summarizer is a function here. The last test
// goes through engine.Summarize with a fake Genkit model, so the wiring the
// agent uses is exercised too.
package compact

import (
	"context"
	"errors"
	"strings"
	"testing"

	"mm/internal/config"
	"mm/internal/engine"

	"github.com/firebase/genkit/go/ai"
	"github.com/firebase/genkit/go/genkit"
)

// --- fixtures ---------------------------------------------------------------

func req(ref, cmd string) *ai.Part {
	return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: "bash", Input: map[string]any{"command": cmd}})
}

func resp(ref, out string) *ai.Part {
	return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: "bash", Output: out})
}

// turn builds one question turn with n commands: user, then n × (model call,
// tool response), then the model's final text.
func turn(question string, n int, output string) []*ai.Message {
	msgs := []*ai.Message{ai.NewUserTextMessage(question)}
	for i := 0; i < n; i++ {
		ref := question + "-" + string(rune('a'+i))
		msgs = append(msgs,
			ai.NewModelMessage(req(ref, "echo "+ref)),
			ai.NewMessage(ai.RoleTool, nil, resp(ref, output)),
		)
	}
	return append(msgs, ai.NewModelMessage(ai.NewTextPart("done "+question)))
}

// history is a system prompt followed by the given turns.
func history(turns ...[]*ai.Message) []*ai.Message {
	msgs := []*ai.Message{ai.NewSystemTextMessage("You are Bob.")}
	for _, t := range turns {
		msgs = append(msgs, t...)
	}
	return msgs
}

func cfg(keep int) config.ContextConfig {
	return config.ContextConfig{Enabled: true, KeepLastTurns: keep, Threshold: 75, MaxMessages: 0, SummaryMaxTokens: 100}
}

// fixed is a Summarizer that returns a constant and records the request.
func fixed(text string) (Summarizer, *[]*ai.Message) {
	var seen []*ai.Message
	return func(_ context.Context, request []*ai.Message) (string, error) {
		seen = request
		return text, nil
	}, &seen
}

// --- Split / Valid / Estimate -----------------------------------------------

func TestSplitCutsOnUserMessages(t *testing.T) {
	msgs := history(turn("q1", 2, "x"), turn("q2", 0, "x"), turn("q3", 1, "x"))
	head, turns := Split(msgs)
	if len(head) != 1 || head[0].Role != ai.RoleSystem {
		t.Fatalf("head = %d message(s), want the system prompt alone", len(head))
	}
	if len(turns) != 3 {
		t.Fatalf("turns = %d, want 3", len(turns))
	}
	// q1: user + 2×(call, response) + text = 6 ; q2: 2 ; q3: 4.
	for i, want := range []int{6, 2, 4} {
		if len(turns[i]) != want {
			t.Errorf("turn %d has %d messages, want %d", i, len(turns[i]), want)
		}
		if turns[i][0].Role != ai.RoleUser {
			t.Errorf("turn %d does not start with a user message", i)
		}
	}
}

func TestSplitTreatsSummaryPairAsATurn(t *testing.T) {
	pair := SummaryPair("notes", 10, 2, 3)
	msgs := history(pair, turn("q4", 1, "x"))
	_, turns := Split(msgs)
	if len(turns) != 2 || !IsSummary(turns[0][0]) || len(turns[0]) != 2 {
		t.Fatalf("the summary pair should be its own two-message turn, got %d turn(s)", len(turns))
	}
}

func TestValid(t *testing.T) {
	ok := history(turn("q1", 2, "x"))
	if err := Valid(ok); err != nil {
		t.Errorf("a well-formed history is refused: %v", err)
	}

	// A `tool` message whose request was cut away: the classic broken window.
	orphan := []*ai.Message{
		ai.NewSystemTextMessage("s"),
		ai.NewMessage(ai.RoleTool, nil, resp("z", "out")),
	}
	if err := Valid(orphan); err == nil {
		t.Error("an orphan tool response is accepted")
	}

	noSystem := turn("q1", 0, "x")
	if err := Valid(noSystem); err == nil {
		t.Error("a history without the system prompt is accepted")
	}

	// No Ref on either side: pairing by name must still work.
	byName := []*ai.Message{
		ai.NewSystemTextMessage("s"),
		ai.NewUserTextMessage("q"),
		ai.NewModelMessage(req("", "ls")),
		ai.NewMessage(ai.RoleTool, nil, resp("", "out")),
	}
	if err := Valid(byName); err != nil {
		t.Errorf("pairing by name fails: %v", err)
	}
}

func TestEstimateGrowsWithToolOutput(t *testing.T) {
	small := history(turn("q1", 1, "x"))
	big := history(turn("q1", 1, strings.Repeat("0123456789\n", 350))) // 3 850 chars ≈ 1 100 tokens
	s, b := Estimate(small), Estimate(big)
	if b <= s {
		t.Fatalf("Estimate(big) = %d is not above Estimate(small) = %d", b, s)
	}
	// 3.5 characters per token: 3 850 chars of output add about 1 100 tokens.
	if delta := b - s; delta < 1000 || delta > 1200 {
		t.Errorf("3 850 characters of output add %d tokens, want ~1 100", delta)
	}
}

// --- Decide -----------------------------------------------------------------

func TestDecide(t *testing.T) {
	msgs := history(turn("q1", 1, strings.Repeat("x", 3500))) // ≈ 1 000 tokens estimated

	off := cfg(1)
	off.Enabled = false
	if d := Decide(msgs, 0, 1000, off); d.Compact {
		t.Error("disabled config still triggers")
	}

	// 1 000 tokens ≥ 75 % of a 1 000-token window: triggers on tokens.
	if d := Decide(msgs, 0, 1000, cfg(1)); !d.Compact || d.Reason != "tokens" {
		t.Errorf("tokens trigger: got %+v", d)
	}

	// A window ten times wider: the estimate is far below the threshold…
	if d := Decide(msgs, 0, 10000, cfg(1)); d.Compact {
		t.Errorf("estimate below threshold still triggers: %+v", d)
	}
	// …but the engine's own measure wins when it is larger.
	if d := Decide(msgs, 9000, 10000, cfg(1)); !d.Compact || d.Tokens != 9000 {
		t.Errorf("measured tokens should win: %+v", d)
	}

	// Unknown window (neither the yaml nor the probe gave one): only the
	// message count can trigger.
	byCount := cfg(1)
	byCount.MaxMessages = 4
	if d := Decide(msgs, 0, 0, byCount); !d.Compact || d.Reason != "messages" {
		t.Errorf("messages fallback: got %+v", d)
	}
	byCount.MaxMessages = 100
	if d := Decide(msgs, 0, 0, byCount); d.Compact {
		t.Errorf("unknown window and few messages still triggers: %+v", d)
	}
}

// --- Compact ----------------------------------------------------------------

func TestCompactKeepsLastTurnsAndPairs(t *testing.T) {
	// Realistic outputs: a `cat` or a `go test -v` runs to hundreds of lines,
	// and that is the weight the summary has to beat — with three-character
	// outputs the wrapper alone would be heavier than what it replaces.
	big := strings.Repeat("line of tool output\n", 60)
	q1, q2, q3, q4 := turn("q1", 2, big), turn("q2", 1, big), turn("q3", 1, big), turn("q4", 0, "")
	msgs := history(q1, q2, q3, q4)
	summarize, seen := fixed("## Goal\nsummary of q1 and q2")

	res, err := Compact(context.Background(), msgs, cfg(2), summarize)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	out := res.Messages

	// The request the model saw: note-taker system, the two old turns, prompt.
	if len(*seen) != 1+len(q1)+len(q2)+1 {
		t.Errorf("summary request has %d messages, want %d", len(*seen), 1+len(q1)+len(q2)+1)
	}
	if (*seen)[0].Role != ai.RoleSystem || (*seen)[0].Text() == msgs[0].Text() {
		t.Error("the summary request must use the note-taker's system prompt, not the agent's")
	}
	if last := (*seen)[len(*seen)-1]; last.Role != ai.RoleUser || !strings.Contains(last.Text(), "## Files touched") {
		t.Error("the summary request must end with the prompt")
	}
	if strings.Contains((*seen)[len(*seen)-1].Text(), "earlier summary") {
		t.Error("no earlier summary here: the merge line must not be added")
	}

	// The shape: system (same pointer), summary pair, then q3 and q4 untouched.
	if out[0] != msgs[0] {
		t.Error("the system prompt must be kept as-is, in first position")
	}
	if !IsSummary(out[1]) || !strings.Contains(out[1].Text(), "[Context summary — the 10 earlier messages") {
		t.Errorf("message 1 should be the summary with its wrapper, got %q", out[1].Text())
	}
	if !strings.Contains(out[1].Text(), "(2 question(s), 3 command(s))") {
		t.Errorf("wrapper counts are off: %q", strings.SplitN(out[1].Text(), "\n", 2)[0])
	}
	if out[2].Role != ai.RoleModel {
		t.Error("message 2 should be the model's acknowledgement")
	}
	rest := out[3:]
	want := append(append([]*ai.Message{}, q3...), q4...)
	if len(rest) != len(want) {
		t.Fatalf("kept %d messages, want %d", len(rest), len(want))
	}
	for i := range want {
		if rest[i] != want[i] {
			t.Errorf("kept message %d is not the original pointer", i)
		}
	}

	if err := Valid(out); err != nil {
		t.Errorf("compressed history invalid: %v", err)
	}
	if res.Compressed != len(q1)+len(q2) || res.Kept != len(q3)+len(q4) {
		t.Errorf("counts: compressed %d kept %d, want %d and %d", res.Compressed, res.Kept, len(q1)+len(q2), len(q3)+len(q4))
	}
	if res.After >= res.Before {
		t.Errorf("tokens did not go down: %d → %d", res.Before, res.After)
	}
	// The input must not have been modified: same length, same pointers.
	if len(msgs) != 1+len(q1)+len(q2)+len(q3)+len(q4) || msgs[1] != q1[0] {
		t.Error("Compact modified its input")
	}
}

func TestCompactNothingToDo(t *testing.T) {
	msgs := history(turn("q1", 3, "x"))
	called := false
	_, err := Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) {
		called = true
		return "should not be asked", nil
	})
	if !errors.Is(err, ErrNothingToCompact) {
		t.Fatalf("err = %v, want ErrNothingToCompact", err)
	}
	if called {
		t.Error("the model was asked for a summary with nothing to summarise")
	}
}

func TestCompactMergesAnEarlierSummary(t *testing.T) {
	pair := SummaryPair("old notes", 12, 3, 5)
	msgs := history(pair, turn("q4", 1, "x"), turn("q5", 0, ""))
	summarize, seen := fixed("merged notes")

	res, err := Compact(context.Background(), msgs, cfg(1), summarize)
	if err != nil {
		t.Fatalf("Compact: %v", err)
	}
	prompt := (*seen)[len(*seen)-1].Text()
	if !strings.Contains(prompt, "earlier summary") {
		t.Error("the merge line is missing from the prompt")
	}
	// The old summary is IN the request (to be merged), not summarised away.
	if !IsSummary((*seen)[1]) {
		t.Error("the earlier summary should be the first old message of the request")
	}
	// Exactly one summary pair remains, and the wrapper counts q4 alone.
	if n := countSummaries(res.Messages); n != 1 {
		t.Errorf("%d summary messages after merge, want 1", n)
	}
	if !strings.Contains(res.Messages[1].Text(), "(1 question(s), 1 command(s))") {
		t.Errorf("wrapper should count the real question only: %q", strings.SplitN(res.Messages[1].Text(), "\n", 2)[0])
	}

	// A summary that is the only old turn is not worth a model call.
	onlySummary := history(pair, turn("q4", 1, "x"))
	if _, err := Compact(context.Background(), onlySummary, cfg(1), summarize); !errors.Is(err, ErrNothingToCompact) {
		t.Errorf("re-summarising a lone summary: err = %v, want ErrNothingToCompact", err)
	}
}

func TestCompactFailureLeavesNoResult(t *testing.T) {
	msgs := history(turn("q1", 1, "x"), turn("q2", 0, ""))
	boom := errors.New("engine down")
	res, err := Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) {
		return "", boom
	})
	if !errors.Is(err, boom) || res.Messages != nil {
		t.Errorf("a failing summary must surface the error and no messages, got err=%v res=%+v", err, res)
	}
	_, err = Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) {
		return "   \n", nil
	})
	if err == nil {
		t.Error("an empty summary must be an error — an empty history is worse than a long one")
	}
}

func countSummaries(msgs []*ai.Message) int {
	n := 0
	for _, m := range msgs {
		if IsSummary(m) {
			n++
		}
	}
	return n
}

// --- through engine.Summarize -----------------------------------------------

// TestCompactThroughEngineSummarize wires the real Summarizer the agent uses
// onto a fake Genkit model, the same approach as internal/engine's tests: the
// request the model receives must carry no tools, and the text it returns must
// land in the summary message. The Engine is built the way testEngine does it
// there — the model reference is all Summarize reads from it.
func TestCompactThroughEngineSummarize(t *testing.T) {
	g := genkit.Init(context.Background())
	var sawTools int
	genkit.DefineModel(g, "dmr/"+config.Cfg.Model,
		&ai.ModelOptions{Supports: &ai.ModelSupports{Tools: true, Multiturn: true, SystemRole: true}},
		func(_ context.Context, r *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) {
			sawTools = len(r.Tools)
			return &ai.ModelResponse{
				Message:      ai.NewModelMessage(ai.NewTextPart("## Goal\nfrom the fake model")),
				FinishReason: ai.FinishReasonStop,
			}, nil
		})
	e := &engine.Engine{G: g, Model: "dmr/" + config.Cfg.Model}

	msgs := history(turn("q1", 2, "x"), turn("q2", 1, "y"))
	c := cfg(1)
	res, err := Compact(context.Background(), msgs, c, func(ctx context.Context, request []*ai.Message) (string, error) {
		return e.Summarize(ctx, request, c.SummaryMaxTokens)
	})
	if err != nil {
		t.Fatalf("Compact via engine.Summarize: %v", err)
	}
	if sawTools != 0 {
		t.Errorf("the summary request declared %d tool(s), want none", sawTools)
	}
	if !strings.Contains(res.Messages[1].Text(), "from the fake model") {
		t.Errorf("the model's text is not in the summary: %q", res.Messages[1].Text())
	}
	if err := Valid(res.Messages); err != nil {
		t.Errorf("invalid history: %v", err)
	}
}