| 💾 Saved. d722711 k33g 7h ago | 1 | // The tests build histories by hand, the way TestCommandList does in |
| 2 | // internal/engine, and check the one property a compression must never break: |
| 3 | // the next request is still a valid conversation — system prompt first, every |
| 4 | // tool response under the request that asked for it, the recent turns intact. |
| 5 | // |
| 6 | // No network and no model: the Summarizer is a function here. The last test |
| 7 | // goes through engine.Summarize with a fake Genkit model, so the wiring the |
| 8 | // agent uses is exercised too. |
| 9 | package compact |
| 10 | |
| 11 | import ( |
| 12 | "context" |
| 13 | "errors" |
| 14 | "strings" |
| 15 | "testing" |
| 16 | |
| 17 | "mm/internal/config" |
| 18 | "mm/internal/engine" |
| 19 | |
| 20 | "github.com/firebase/genkit/go/ai" |
| 21 | "github.com/firebase/genkit/go/genkit" |
| 22 | ) |
| 23 | |
| 24 | // --- fixtures --------------------------------------------------------------- |
| 25 | |
| 26 | func req(ref, cmd string) *ai.Part { |
| 27 | return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: "bash", Input: map[string]any{"command": cmd}}) |
| 28 | } |
| 29 | |
| 30 | func resp(ref, out string) *ai.Part { |
| 31 | return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: "bash", Output: out}) |
| 32 | } |
| 33 | |
| 34 | // turn builds one question turn with n commands: user, then n × (model call, |
| 35 | // tool response), then the model's final text. |
| 36 | func turn(question string, n int, output string) []*ai.Message { |
| 37 | msgs := []*ai.Message{ai.NewUserTextMessage(question)} |
| 38 | for i := 0; i < n; i++ { |
| 39 | ref := question + "-" + string(rune('a'+i)) |
| 40 | msgs = append(msgs, |
| 41 | ai.NewModelMessage(req(ref, "echo "+ref)), |
| 42 | ai.NewMessage(ai.RoleTool, nil, resp(ref, output)), |
| 43 | ) |
| 44 | } |
| 45 | return append(msgs, ai.NewModelMessage(ai.NewTextPart("done "+question))) |
| 46 | } |
| 47 | |
| 48 | // history is a system prompt followed by the given turns. |
| 49 | func history(turns ...[]*ai.Message) []*ai.Message { |
| 50 | msgs := []*ai.Message{ai.NewSystemTextMessage("You are Bob.")} |
| 51 | for _, t := range turns { |
| 52 | msgs = append(msgs, t...) |
| 53 | } |
| 54 | return msgs |
| 55 | } |
| 56 | |
| 57 | func cfg(keep int) config.ContextConfig { |
| 58 | return config.ContextConfig{Enabled: true, KeepLastTurns: keep, Threshold: 75, MaxMessages: 0, SummaryMaxTokens: 100} |
| 59 | } |
| 60 | |
| 61 | // fixed is a Summarizer that returns a constant and records the request. |
| 62 | func fixed(text string) (Summarizer, *[]*ai.Message) { |
| 63 | var seen []*ai.Message |
| 64 | return func(_ context.Context, request []*ai.Message) (string, error) { |
| 65 | seen = request |
| 66 | return text, nil |
| 67 | }, &seen |
| 68 | } |
| 69 | |
| 70 | // --- Split / Valid / Estimate ----------------------------------------------- |
| 71 | |
| 72 | func TestSplitCutsOnUserMessages(t *testing.T) { |
| 73 | msgs := history(turn("q1", 2, "x"), turn("q2", 0, "x"), turn("q3", 1, "x")) |
| 74 | head, turns := Split(msgs) |
| 75 | if len(head) != 1 || head[0].Role != ai.RoleSystem { |
| 76 | t.Fatalf("head = %d message(s), want the system prompt alone", len(head)) |
| 77 | } |
| 78 | if len(turns) != 3 { |
| 79 | t.Fatalf("turns = %d, want 3", len(turns)) |
| 80 | } |
| 81 | // q1: user + 2×(call, response) + text = 6 ; q2: 2 ; q3: 4. |
| 82 | for i, want := range []int{6, 2, 4} { |
| 83 | if len(turns[i]) != want { |
| 84 | t.Errorf("turn %d has %d messages, want %d", i, len(turns[i]), want) |
| 85 | } |
| 86 | if turns[i][0].Role != ai.RoleUser { |
| 87 | t.Errorf("turn %d does not start with a user message", i) |
| 88 | } |
| 89 | } |
| 90 | } |
| 91 | |
| 92 | func TestSplitTreatsSummaryPairAsATurn(t *testing.T) { |
| 93 | pair := SummaryPair("notes", 10, 2, 3) |
| 94 | msgs := history(pair, turn("q4", 1, "x")) |
| 95 | _, turns := Split(msgs) |
| 96 | if len(turns) != 2 || !IsSummary(turns[0][0]) || len(turns[0]) != 2 { |
| 97 | t.Fatalf("the summary pair should be its own two-message turn, got %d turn(s)", len(turns)) |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | func TestValid(t *testing.T) { |
| 102 | ok := history(turn("q1", 2, "x")) |
| 103 | if err := Valid(ok); err != nil { |
| 104 | t.Errorf("a well-formed history is refused: %v", err) |
| 105 | } |
| 106 | |
| 107 | // A `tool` message whose request was cut away: the classic broken window. |
| 108 | orphan := []*ai.Message{ |
| 109 | ai.NewSystemTextMessage("s"), |
| 110 | ai.NewMessage(ai.RoleTool, nil, resp("z", "out")), |
| 111 | } |
| 112 | if err := Valid(orphan); err == nil { |
| 113 | t.Error("an orphan tool response is accepted") |
| 114 | } |
| 115 | |
| 116 | noSystem := turn("q1", 0, "x") |
| 117 | if err := Valid(noSystem); err == nil { |
| 118 | t.Error("a history without the system prompt is accepted") |
| 119 | } |
| 120 | |
| 121 | // No Ref on either side: pairing by name must still work. |
| 122 | byName := []*ai.Message{ |
| 123 | ai.NewSystemTextMessage("s"), |
| 124 | ai.NewUserTextMessage("q"), |
| 125 | ai.NewModelMessage(req("", "ls")), |
| 126 | ai.NewMessage(ai.RoleTool, nil, resp("", "out")), |
| 127 | } |
| 128 | if err := Valid(byName); err != nil { |
| 129 | t.Errorf("pairing by name fails: %v", err) |
| 130 | } |
| 131 | } |
| 132 | |
| 133 | func TestEstimateGrowsWithToolOutput(t *testing.T) { |
| 134 | small := history(turn("q1", 1, "x")) |
| 135 | big := history(turn("q1", 1, strings.Repeat("0123456789\n", 350))) // 3 850 chars ≈ 1 100 tokens |
| 136 | s, b := Estimate(small), Estimate(big) |
| 137 | if b <= s { |
| 138 | t.Fatalf("Estimate(big) = %d is not above Estimate(small) = %d", b, s) |
| 139 | } |
| 140 | // 3.5 characters per token: 3 850 chars of output add about 1 100 tokens. |
| 141 | if delta := b - s; delta < 1000 || delta > 1200 { |
| 142 | t.Errorf("3 850 characters of output add %d tokens, want ~1 100", delta) |
| 143 | } |
| 144 | } |
| 145 | |
| 146 | // --- Decide ----------------------------------------------------------------- |
| 147 | |
| 148 | func TestDecide(t *testing.T) { |
| 149 | msgs := history(turn("q1", 1, strings.Repeat("x", 3500))) // ≈ 1 000 tokens estimated |
| 150 | |
| 151 | off := cfg(1) |
| 152 | off.Enabled = false |
| 153 | if d := Decide(msgs, 0, 1000, off); d.Compact { |
| 154 | t.Error("disabled config still triggers") |
| 155 | } |
| 156 | |
| 157 | // 1 000 tokens ≥ 75 % of a 1 000-token window: triggers on tokens. |
| 158 | if d := Decide(msgs, 0, 1000, cfg(1)); !d.Compact || d.Reason != "tokens" { |
| 159 | t.Errorf("tokens trigger: got %+v", d) |
| 160 | } |
| 161 | |
| 162 | // A window ten times wider: the estimate is far below the threshold… |
| 163 | if d := Decide(msgs, 0, 10000, cfg(1)); d.Compact { |
| 164 | t.Errorf("estimate below threshold still triggers: %+v", d) |
| 165 | } |
| 166 | // …but the engine's own measure wins when it is larger. |
| 167 | if d := Decide(msgs, 9000, 10000, cfg(1)); !d.Compact || d.Tokens != 9000 { |
| 168 | t.Errorf("measured tokens should win: %+v", d) |
| 169 | } |
| 170 | |
| 171 | // Unknown window (neither the yaml nor the probe gave one): only the |
| 172 | // message count can trigger. |
| 173 | byCount := cfg(1) |
| 174 | byCount.MaxMessages = 4 |
| 175 | if d := Decide(msgs, 0, 0, byCount); !d.Compact || d.Reason != "messages" { |
| 176 | t.Errorf("messages fallback: got %+v", d) |
| 177 | } |
| 178 | byCount.MaxMessages = 100 |
| 179 | if d := Decide(msgs, 0, 0, byCount); d.Compact { |
| 180 | t.Errorf("unknown window and few messages still triggers: %+v", d) |
| 181 | } |
| 182 | } |
| 183 | |
| 184 | // --- Compact ---------------------------------------------------------------- |
| 185 | |
| 186 | func TestCompactKeepsLastTurnsAndPairs(t *testing.T) { |
| 187 | // Realistic outputs: a `cat` or a `go test -v` runs to hundreds of lines, |
| 188 | // and that is the weight the summary has to beat — with three-character |
| 189 | // outputs the wrapper alone would be heavier than what it replaces. |
| 190 | big := strings.Repeat("line of tool output\n", 60) |
| 191 | q1, q2, q3, q4 := turn("q1", 2, big), turn("q2", 1, big), turn("q3", 1, big), turn("q4", 0, "") |
| 192 | msgs := history(q1, q2, q3, q4) |
| 193 | summarize, seen := fixed("## Goal\nsummary of q1 and q2") |
| 194 | |
| 195 | res, err := Compact(context.Background(), msgs, cfg(2), summarize) |
| 196 | if err != nil { |
| 197 | t.Fatalf("Compact: %v", err) |
| 198 | } |
| 199 | out := res.Messages |
| 200 | |
| 201 | // The request the model saw: note-taker system, the two old turns, prompt. |
| 202 | if len(*seen) != 1+len(q1)+len(q2)+1 { |
| 203 | t.Errorf("summary request has %d messages, want %d", len(*seen), 1+len(q1)+len(q2)+1) |
| 204 | } |
| 205 | if (*seen)[0].Role != ai.RoleSystem || (*seen)[0].Text() == msgs[0].Text() { |
| 206 | t.Error("the summary request must use the note-taker's system prompt, not the agent's") |
| 207 | } |
| 208 | if last := (*seen)[len(*seen)-1]; last.Role != ai.RoleUser || !strings.Contains(last.Text(), "## Files touched") { |
| 209 | t.Error("the summary request must end with the prompt") |
| 210 | } |
| 211 | if strings.Contains((*seen)[len(*seen)-1].Text(), "earlier summary") { |
| 212 | t.Error("no earlier summary here: the merge line must not be added") |
| 213 | } |
| 214 | |
| 215 | // The shape: system (same pointer), summary pair, then q3 and q4 untouched. |
| 216 | if out[0] != msgs[0] { |
| 217 | t.Error("the system prompt must be kept as-is, in first position") |
| 218 | } |
| 219 | if !IsSummary(out[1]) || !strings.Contains(out[1].Text(), "[Context summary — the 10 earlier messages") { |
| 220 | t.Errorf("message 1 should be the summary with its wrapper, got %q", out[1].Text()) |
| 221 | } |
| 222 | if !strings.Contains(out[1].Text(), "(2 question(s), 3 command(s))") { |
| 223 | t.Errorf("wrapper counts are off: %q", strings.SplitN(out[1].Text(), "\n", 2)[0]) |
| 224 | } |
| 225 | if out[2].Role != ai.RoleModel { |
| 226 | t.Error("message 2 should be the model's acknowledgement") |
| 227 | } |
| 228 | rest := out[3:] |
| 229 | want := append(append([]*ai.Message{}, q3...), q4...) |
| 230 | if len(rest) != len(want) { |
| 231 | t.Fatalf("kept %d messages, want %d", len(rest), len(want)) |
| 232 | } |
| 233 | for i := range want { |
| 234 | if rest[i] != want[i] { |
| 235 | t.Errorf("kept message %d is not the original pointer", i) |
| 236 | } |
| 237 | } |
| 238 | |
| 239 | if err := Valid(out); err != nil { |
| 240 | t.Errorf("compressed history invalid: %v", err) |
| 241 | } |
| 242 | if res.Compressed != len(q1)+len(q2) || res.Kept != len(q3)+len(q4) { |
| 243 | t.Errorf("counts: compressed %d kept %d, want %d and %d", res.Compressed, res.Kept, len(q1)+len(q2), len(q3)+len(q4)) |
| 244 | } |
| 245 | if res.After >= res.Before { |
| 246 | t.Errorf("tokens did not go down: %d → %d", res.Before, res.After) |
| 247 | } |
| 248 | // The input must not have been modified: same length, same pointers. |
| 249 | if len(msgs) != 1+len(q1)+len(q2)+len(q3)+len(q4) || msgs[1] != q1[0] { |
| 250 | t.Error("Compact modified its input") |
| 251 | } |
| 252 | } |
| 253 | |
| 254 | func TestCompactNothingToDo(t *testing.T) { |
| 255 | msgs := history(turn("q1", 3, "x")) |
| 256 | called := false |
| 257 | _, err := Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) { |
| 258 | called = true |
| 259 | return "should not be asked", nil |
| 260 | }) |
| 261 | if !errors.Is(err, ErrNothingToCompact) { |
| 262 | t.Fatalf("err = %v, want ErrNothingToCompact", err) |
| 263 | } |
| 264 | if called { |
| 265 | t.Error("the model was asked for a summary with nothing to summarise") |
| 266 | } |
| 267 | } |
| 268 | |
| 269 | func TestCompactMergesAnEarlierSummary(t *testing.T) { |
| 270 | pair := SummaryPair("old notes", 12, 3, 5) |
| 271 | msgs := history(pair, turn("q4", 1, "x"), turn("q5", 0, "")) |
| 272 | summarize, seen := fixed("merged notes") |
| 273 | |
| 274 | res, err := Compact(context.Background(), msgs, cfg(1), summarize) |
| 275 | if err != nil { |
| 276 | t.Fatalf("Compact: %v", err) |
| 277 | } |
| 278 | prompt := (*seen)[len(*seen)-1].Text() |
| 279 | if !strings.Contains(prompt, "earlier summary") { |
| 280 | t.Error("the merge line is missing from the prompt") |
| 281 | } |
| 282 | // The old summary is IN the request (to be merged), not summarised away. |
| 283 | if !IsSummary((*seen)[1]) { |
| 284 | t.Error("the earlier summary should be the first old message of the request") |
| 285 | } |
| 286 | // Exactly one summary pair remains, and the wrapper counts q4 alone. |
| 287 | if n := countSummaries(res.Messages); n != 1 { |
| 288 | t.Errorf("%d summary messages after merge, want 1", n) |
| 289 | } |
| 290 | if !strings.Contains(res.Messages[1].Text(), "(1 question(s), 1 command(s))") { |
| 291 | t.Errorf("wrapper should count the real question only: %q", strings.SplitN(res.Messages[1].Text(), "\n", 2)[0]) |
| 292 | } |
| 293 | |
| 294 | // A summary that is the only old turn is not worth a model call. |
| 295 | onlySummary := history(pair, turn("q4", 1, "x")) |
| 296 | if _, err := Compact(context.Background(), onlySummary, cfg(1), summarize); !errors.Is(err, ErrNothingToCompact) { |
| 297 | t.Errorf("re-summarising a lone summary: err = %v, want ErrNothingToCompact", err) |
| 298 | } |
| 299 | } |
| 300 | |
| 301 | func TestCompactFailureLeavesNoResult(t *testing.T) { |
| 302 | msgs := history(turn("q1", 1, "x"), turn("q2", 0, "")) |
| 303 | boom := errors.New("engine down") |
| 304 | res, err := Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) { |
| 305 | return "", boom |
| 306 | }) |
| 307 | if !errors.Is(err, boom) || res.Messages != nil { |
| 308 | t.Errorf("a failing summary must surface the error and no messages, got err=%v res=%+v", err, res) |
| 309 | } |
| 310 | _, err = Compact(context.Background(), msgs, cfg(1), func(context.Context, []*ai.Message) (string, error) { |
| 311 | return " \n", nil |
| 312 | }) |
| 313 | if err == nil { |
| 314 | t.Error("an empty summary must be an error — an empty history is worse than a long one") |
| 315 | } |
| 316 | } |
| 317 | |
| 318 | func countSummaries(msgs []*ai.Message) int { |
| 319 | n := 0 |
| 320 | for _, m := range msgs { |
| 321 | if IsSummary(m) { |
| 322 | n++ |
| 323 | } |
| 324 | } |
| 325 | return n |
| 326 | } |
| 327 | |
| 328 | // --- through engine.Summarize ----------------------------------------------- |
| 329 | |
| 330 | // TestCompactThroughEngineSummarize wires the real Summarizer the agent uses |
| 331 | // onto a fake Genkit model, the same approach as internal/engine's tests: the |
| 332 | // request the model receives must carry no tools, and the text it returns must |
| 333 | // land in the summary message. The Engine is built the way testEngine does it |
| 334 | // there — the model reference is all Summarize reads from it. |
| 335 | func TestCompactThroughEngineSummarize(t *testing.T) { |
| 336 | g := genkit.Init(context.Background()) |
| 337 | var sawTools int |
| 338 | genkit.DefineModel(g, "dmr/"+config.Cfg.Model, |
| 339 | &ai.ModelOptions{Supports: &ai.ModelSupports{Tools: true, Multiturn: true, SystemRole: true}}, |
| 340 | func(_ context.Context, r *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) { |
| 341 | sawTools = len(r.Tools) |
| 342 | return &ai.ModelResponse{ |
| 343 | Message: ai.NewModelMessage(ai.NewTextPart("## Goal\nfrom the fake model")), |
| 344 | FinishReason: ai.FinishReasonStop, |
| 345 | }, nil |
| 346 | }) |
| 347 | e := &engine.Engine{G: g, Model: "dmr/" + config.Cfg.Model} |
| 348 | |
| 349 | msgs := history(turn("q1", 2, "x"), turn("q2", 1, "y")) |
| 350 | c := cfg(1) |
| 351 | res, err := Compact(context.Background(), msgs, c, func(ctx context.Context, request []*ai.Message) (string, error) { |
| 352 | return e.Summarize(ctx, request, c.SummaryMaxTokens) |
| 353 | }) |
| 354 | if err != nil { |
| 355 | t.Fatalf("Compact via engine.Summarize: %v", err) |
| 356 | } |
| 357 | if sawTools != 0 { |
| 358 | t.Errorf("the summary request declared %d tool(s), want none", sawTools) |
| 359 | } |
| 360 | if !strings.Contains(res.Messages[1].Text(), "from the fake model") { |
| 361 | t.Errorf("the model's text is not in the summary: %q", res.Messages[1].Text()) |
| 362 | } |
| 363 | if err := Valid(res.Messages); err != nil { |
| 364 | t.Errorf("invalid history: %v", err) |
| 365 | } |
| 366 | } |