| 💾 Saved. d722711 k33g 4h ago | 1 | // The test that matters here is the one about CAPTURE: Genkit only hands back |
| 2 | // the last message, and the whole point of `Generate` is to hand back the |
| 3 | // intermediate turns too. That is the one thing the `WithMiddleware` → `WithUse` |
| 4 | // migration could have broken silently — the code compiles just as well with a |
| 5 | // hook that never sees anything. |
| 6 | // |
| 7 | // No real model is needed: a fake model registered in Genkit plays one tool turn |
| 8 | // then answers, which is exactly the shape the history must keep a trace of. |
| 9 | // |
| 10 | // The watchdog of this version runs during these tests: since the fake model |
| 11 | // answers instantly it never fires — and if it did, the tests would fail, which |
| 12 | // is the intended behaviour. |
| 13 | package engine |
| 14 | |
| 15 | import ( |
| 16 | "context" |
| 17 | "testing" |
| 18 | |
| 19 | "mm/internal/config" |
| 20 | |
| 21 | "github.com/firebase/genkit/go/ai" |
| 22 | "github.com/firebase/genkit/go/genkit" |
| 23 | ) |
| 24 | |
| 25 | // fakeGenkit registers a "dmr/<Cfg.Model>" model whose every call is served by |
| 26 | // the next answer in `turns`, and a `bash` tool returning a fixed output. |
| 27 | // `calls` counts the calls to the model. |
| 28 | func fakeGenkit(t *testing.T, turns ...*ai.ModelResponse) (*genkit.Genkit, *int) { |
| 29 | t.Helper() |
| 30 | g := genkit.Init(context.Background()) |
| 31 | |
| 32 | calls := 0 |
| 33 | genkit.DefineModel(g, "dmr/"+config.Cfg.Model, |
| 34 | &ai.ModelOptions{Supports: &ai.ModelSupports{Tools: true, Multiturn: true, SystemRole: true}}, |
| 35 | func(_ context.Context, _ *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) { |
| 36 | if calls >= len(turns) { |
| 37 | t.Errorf("the model was called %d times for %d planned turns", calls+1, len(turns)) |
| 38 | return turns[len(turns)-1], nil |
| 39 | } |
| 40 | resp := turns[calls] |
| 41 | calls++ |
| 42 | return resp, nil |
| 43 | }) |
| 44 | |
| 45 | genkit.DefineTool(g, "bash", "run a shell command", |
| 46 | func(_ *ai.ToolContext, in struct { |
| 47 | Command string `json:"command"` |
| 48 | }) (string, error) { |
| 49 | return "hello from " + in.Command, nil |
| 50 | }) |
| 51 | |
| 52 | genkit.DefineTool(g, "read_skill", "load a markdown procedure", |
| 53 | func(_ *ai.ToolContext, in struct { |
| 54 | Name string `json:"name"` |
| 55 | }) (string, error) { |
| 56 | return "# procedure " + in.Name, nil |
| 57 | }) |
| 58 | |
| 59 | return g, &calls |
| 60 | } |
| 61 | |
| 62 | // testEngine wraps a fake Genkit the way New would: the model reference is the |
| 63 | // only thing Generate reads from the Engine, so no provider is needed here. |
| 64 | func testEngine(g *genkit.Genkit) *Engine { |
| 65 | return &Engine{G: g, Model: "dmr/" + config.Cfg.Model} |
| 66 | } |
| 67 | |
| 68 | func modelTurn(parts ...*ai.Part) *ai.ModelResponse { |
| 69 | return &ai.ModelResponse{ |
| 70 | Message: ai.NewModelMessage(parts...), |
| 71 | FinishReason: ai.FinishReasonStop, |
| 72 | } |
| 73 | } |
| 74 | |
| 75 | // TestGenerateReturnsFullConversation is the regression test for the capture: on |
| 76 | // a tool turn, the history handed back must contain the question, the tool |
| 77 | // request, its response AND the final text — not just that last one. |
| 78 | func TestGenerateReturnsFullConversation(t *testing.T) { |
| 79 | toolReq := modelTurn(ai.NewToolRequestPart(&ai.ToolRequest{ |
| 80 | Name: "bash", |
| 81 | Input: map[string]any{"command": "echo hi"}, |
| 82 | })) |
| 83 | final := modelTurn(ai.NewTextPart("done")) |
| 84 | |
| 85 | g, calls := fakeGenkit(t, toolReq, final) |
| 86 | |
| 87 | resp, full, err := testEngine(g).Generate(context.Background(), |
| 88 | []*ai.Message{ai.NewUserTextMessage("run echo hi")}, |
| 89 | []ai.ToolRef{ai.ToolName("bash")}) |
| 90 | if err != nil { |
| 91 | t.Fatalf("Generate: %v", err) |
| 92 | } |
| 93 | if *calls != 2 { |
| 94 | t.Errorf("calls to the model = %d, want 2 (tool, then answer)", *calls) |
| 95 | } |
| 96 | if got := resp.Text(); got != "done" { |
| 97 | t.Errorf("resp.Text() = %q, want %q", got, "done") |
| 98 | } |
| 99 | |
| 100 | // The full history: without the capture, Genkit only hands back the last |
| 101 | // message and `full` would have length 1. |
| 102 | if len(full) < 4 { |
| 103 | t.Fatalf("history of %d messages, want at least 4 (question, call, tool response, text): %+v", len(full), full) |
| 104 | } |
| 105 | if full[len(full)-1] != final.Message { |
| 106 | t.Errorf("the last message of the history is not the final answer") |
| 107 | } |
| 108 | |
| 109 | // The two markers that exist ONLY in the intermediate turns. |
| 110 | var sawRequest, sawResponse bool |
| 111 | for _, m := range full { |
| 112 | for _, p := range m.Content { |
| 113 | if p.IsToolRequest() { |
| 114 | sawRequest = true |
| 115 | } |
| 116 | if p.IsToolResponse() { |
| 117 | sawResponse = true |
| 118 | } |
| 119 | } |
| 120 | } |
| 121 | if !sawRequest { |
| 122 | t.Error("the history contains no tool request") |
| 123 | } |
| 124 | if !sawResponse { |
| 125 | t.Error("the history contains no tool response") |
| 126 | } |
| 127 | |
| 128 | // Commands reads the same trace: one command run, exactly one. |
| 129 | if n := Commands(full); n != 1 { |
| 130 | t.Errorf("Commands(full) = %d, want 1", n) |
| 131 | } |
| 132 | // And CommandList hands back its TEXT, which exists only in the call — that |
| 133 | // is what the green recap displays. |
| 134 | if got := CommandList(full); len(got) != 1 || got[0] != "echo hi" { |
| 135 | t.Errorf("CommandList(full) = %q, want [\"echo hi\"]", got) |
| 136 | } |
| 137 | } |
| 138 | |
| 139 | // TestCommandsCountsBashOnly guards the bug introduced by adding `read_skill`: |
| 140 | // `Commands` counted EVERY tool response, so loading two skills and running a |
| 141 | // single command displayed "⚙ 3 command(s)". The number then said the opposite |
| 142 | // of what it is for — that the model acted. |
| 143 | func TestCommandsCountsBashOnly(t *testing.T) { |
| 144 | skillReq := modelTurn(ai.NewToolRequestPart(&ai.ToolRequest{ |
| 145 | Name: "read_skill", |
| 146 | Input: map[string]any{"name": "go-rename"}, |
| 147 | })) |
| 148 | bashReq := modelTurn(ai.NewToolRequestPart(&ai.ToolRequest{ |
| 149 | Name: "bash", |
| 150 | Input: map[string]any{"command": "echo hi"}, |
| 151 | })) |
| 152 | final := modelTurn(ai.NewTextPart("done")) |
| 153 | |
| 154 | g, calls := fakeGenkit(t, skillReq, bashReq, final) |
| 155 | |
| 156 | _, full, err := testEngine(g).Generate(context.Background(), |
| 157 | []*ai.Message{ai.NewUserTextMessage("rename Greet")}, |
| 158 | []ai.ToolRef{ai.ToolName("bash"), ai.ToolName("read_skill")}) |
| 159 | if err != nil { |
| 160 | t.Fatalf("Generate: %v", err) |
| 161 | } |
| 162 | if *calls != 3 { |
| 163 | t.Errorf("calls to the model = %d, want 3", *calls) |
| 164 | } |
| 165 | if n := Commands(full); n != 1 { |
| 166 | t.Errorf("Commands = %d, want 1 — only the bash command counts", n) |
| 167 | } |
| 168 | if k := Skills(full); k != 1 { |
| 169 | t.Errorf("Skills = %d, want 1", k) |
| 170 | } |
| 171 | // The list follows the same rule as the count: the skill read is not in it. |
| 172 | if got := CommandList(full); len(got) != 1 || got[0] != "echo hi" { |
| 173 | t.Errorf("CommandList = %q, want [\"echo hi\"] — a skill is not a command", got) |
| 174 | } |
| 175 | } |
| 176 | |
| 177 | // TestCommandList covers the shapes the call↔response pairing has to take: |
| 178 | // several commands in order, a Ref that does not follow the order they were |
| 179 | // issued in, a call left without a response (nothing ran: nothing to show) |
| 180 | // and a turn mixing `bash` and `read_skill` with no Ref. |
| 181 | func TestCommandList(t *testing.T) { |
| 182 | req := func(ref, name string, in any) *ai.Part { |
| 183 | return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: name, Input: in}) |
| 184 | } |
| 185 | resp := func(ref, name string) *ai.Part { |
| 186 | return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: name, Output: "ok"}) |
| 187 | } |
| 188 | bash := func(cmd string) map[string]any { return map[string]any{"command": cmd} } |
| 189 | |
| 190 | cases := []struct { |
| 191 | name string |
| 192 | history []*ai.Message |
| 193 | want []string |
| 194 | }{ |
| 195 | { |
| 196 | name: "two commands, in the order they ran", |
| 197 | history: []*ai.Message{ |
| 198 | ai.NewModelMessage(req("a", "bash", bash("ls")), req("b", "bash", bash("pwd"))), |
| 199 | ai.NewMessage(ai.RoleTool, nil, resp("a", "bash"), resp("b", "bash")), |
| 200 | }, |
| 201 | want: []string{"ls", "pwd"}, |
| 202 | }, |
| 203 | { |
| 204 | name: "responses out of order: the Ref decides", |
| 205 | history: []*ai.Message{ |
| 206 | ai.NewModelMessage(req("a", "bash", bash("ls")), req("b", "bash", bash("pwd"))), |
| 207 | ai.NewMessage(ai.RoleTool, nil, resp("b", "bash"), resp("a", "bash")), |
| 208 | }, |
| 209 | want: []string{"pwd", "ls"}, |
| 210 | }, |
| 211 | { |
| 212 | name: "no Ref, mixed tools: paired against the same name", |
| 213 | history: []*ai.Message{ |
| 214 | ai.NewModelMessage( |
| 215 | req("", "read_skill", map[string]any{"name": "go-rename"}), |
| 216 | req("", "bash", bash("ls")), |
| 217 | ), |
| 218 | ai.NewMessage(ai.RoleTool, nil, resp("", "read_skill"), resp("", "bash")), |
| 219 | }, |
| 220 | want: []string{"ls"}, |
| 221 | }, |
| 222 | { |
| 223 | name: "call with no response: the command never ran", |
| 224 | history: []*ai.Message{ |
| 225 | ai.NewModelMessage(req("a", "bash", bash("ls"))), |
| 226 | }, |
| 227 | want: nil, |
| 228 | }, |
| 229 | } |
| 230 | |
| 231 | for _, tc := range cases { |
| 232 | t.Run(tc.name, func(t *testing.T) { |
| 233 | got := CommandList(tc.history) |
| 234 | if len(got) != len(tc.want) { |
| 235 | t.Fatalf("CommandList = %q, want %q", got, tc.want) |
| 236 | } |
| 237 | for i := range got { |
| 238 | if got[i] != tc.want[i] { |
| 239 | t.Errorf("CommandList[%d] = %q, want %q", i, got[i], tc.want[i]) |
| 240 | } |
| 241 | } |
| 242 | }) |
| 243 | } |
| 244 | } |
| 245 | |
| 246 | // TestGenerateNoToolCall checks the bare case: with no tool called, the history |
| 247 | // stays the question plus the answer, and Commands counts zero — a model that |
| 248 | // tells stories without acting must not be credited with a command. |
| 249 | func TestGenerateNoToolCall(t *testing.T) { |
| 250 | final := modelTurn(ai.NewTextPart("hello")) |
| 251 | g, calls := fakeGenkit(t, final) |
| 252 | |
| 253 | _, full, err := testEngine(g).Generate(context.Background(), |
| 254 | []*ai.Message{ai.NewUserTextMessage("say hello")}, |
| 255 | nil) |
| 256 | if err != nil { |
| 257 | t.Fatalf("Generate: %v", err) |
| 258 | } |
| 259 | if *calls != 1 { |
| 260 | t.Errorf("calls to the model = %d, want 1", *calls) |
| 261 | } |
| 262 | if len(full) != 2 { |
| 263 | t.Fatalf("history of %d messages, want 2 (question, answer): %+v", len(full), full) |
| 264 | } |
| 265 | if n := Commands(full); n != 0 { |
| 266 | t.Errorf("Commands(full) = %d, want 0", n) |
| 267 | } |
| 268 | } |
| 269 | |
| 270 | // TestGenerateCapturesInputTokens: the compression trigger prefers the server's |
| 271 | // own count of the context to its estimate, and that count only exists inside |
| 272 | // the WrapModel hook. A fake response carrying Usage must surface through |
| 273 | // LastInputTokens; ForgetInputTokens must clear it. |
| 274 | func TestGenerateCapturesInputTokens(t *testing.T) { |
| 275 | withUsage := modelTurn(ai.NewTextPart("ok")) |
| 276 | withUsage.Usage = &ai.GenerationUsage{InputTokens: 1234, OutputTokens: 5} |
| 277 | |
| 278 | g, _ := fakeGenkit(t, withUsage) |
| 279 | e := testEngine(g) |
| 280 | if _, _, err := e.Generate(context.Background(), |
| 281 | []*ai.Message{ai.NewUserTextMessage("hi")}, nil); err != nil { |
| 282 | t.Fatalf("Generate: %v", err) |
| 283 | } |
| 284 | if got := e.LastInputTokens(); got != 1234 { |
| 285 | t.Errorf("LastInputTokens() = %d, want 1234", got) |
| 286 | } |
| 287 | e.ForgetInputTokens() |
| 288 | if got := e.LastInputTokens(); got != 0 { |
| 289 | t.Errorf("after ForgetInputTokens, LastInputTokens() = %d, want 0", got) |
| 290 | } |
| 291 | } |
| 292 | |
| 293 | // TestSummarize: the summary request declares no tools (the model must write, |
| 294 | // not act), goes to the Engine's own model reference, and the model's text |
| 295 | // comes back trimmed. An empty answer is an error — replacing the history with |
| 296 | // nothing is worse than keeping it. |
| 297 | func TestSummarize(t *testing.T) { |
| 298 | g := genkit.Init(context.Background()) |
| 299 | var sawTools int |
| 300 | answers := []string{" ## Goal\nnotes\n", ""} |
| 301 | call := 0 |
| 302 | genkit.DefineModel(g, "dmr/"+config.Cfg.Model, |
| 303 | &ai.ModelOptions{Supports: &ai.ModelSupports{Tools: true, Multiturn: true, SystemRole: true}}, |
| 304 | func(_ context.Context, r *ai.ModelRequest, _ ai.ModelStreamCallback) (*ai.ModelResponse, error) { |
| 305 | sawTools = len(r.Tools) |
| 306 | a := answers[call] |
| 307 | call++ |
| 308 | return modelTurn(ai.NewTextPart(a)), nil |
| 309 | }) |
| 310 | e := testEngine(g) |
| 311 | |
| 312 | req := []*ai.Message{ai.NewSystemTextMessage("note-taker"), ai.NewUserTextMessage("summarise")} |
| 313 | text, err := e.Summarize(context.Background(), req, 100) |
| 314 | if err != nil { |
| 315 | t.Fatalf("Summarize: %v", err) |
| 316 | } |
| 317 | if text != "## Goal\nnotes" { |
| 318 | t.Errorf("text = %q, want the trimmed answer", text) |
| 319 | } |
| 320 | if sawTools != 0 { |
| 321 | t.Errorf("the summary request declared %d tool(s), want none", sawTools) |
| 322 | } |
| 323 | if _, err := e.Summarize(context.Background(), req, 100); err == nil { |
| 324 | t.Error("an empty answer must be an error") |
| 325 | } |
| 326 | } |
| 327 | |
| 328 | // TestFileOpsKeepsTheOrder: the recap lists file operations in the order they |
| 329 | // ran, across the three tools — a read then an edit must not come out as "all |
| 330 | // reads, then all edits", which a per-tool pass would produce. Commands stay |
| 331 | // bash-only: the file tools have their own column. |
| 332 | func TestFileOpsKeepsTheOrder(t *testing.T) { |
| 333 | req := func(ref, name string, in map[string]any) *ai.Part { |
| 334 | return ai.NewToolRequestPart(&ai.ToolRequest{Ref: ref, Name: name, Input: in}) |
| 335 | } |
| 336 | resp := func(ref, name string) *ai.Part { |
| 337 | return ai.NewToolResponsePart(&ai.ToolResponse{Ref: ref, Name: name, Output: "ok"}) |
| 338 | } |
| 339 | history := []*ai.Message{ |
| 340 | ai.NewModelMessage(req("a", "read_file", map[string]any{"path": "f.go", "start": 1.0, "end": 9.0})), |
| 341 | ai.NewMessage(ai.RoleTool, nil, resp("a", "read_file")), |
| 342 | ai.NewModelMessage(req("b", "bash", map[string]any{"command": "go vet ./..."})), |
| 343 | ai.NewMessage(ai.RoleTool, nil, resp("b", "bash")), |
| 344 | ai.NewModelMessage(req("c", "edit_file", map[string]any{"path": "f.go", "edits": []any{map[string]any{"old": "a", "new": "b"}, map[string]any{"old": "c", "new": "d"}}})), |
| 345 | ai.NewMessage(ai.RoleTool, nil, resp("c", "edit_file")), |
| 346 | ai.NewModelMessage(req("d", "write_file", map[string]any{"path": "notes.md", "content": "x"})), |
| 347 | ai.NewMessage(ai.RoleTool, nil, resp("d", "write_file")), |
| 348 | } |
| 349 | want := []string{"read_file f.go 1-9", "edit_file f.go (2 edit(s))", "write_file notes.md"} |
| 350 | got := FileOps(history) |
| 351 | if len(got) != len(want) { |
| 352 | t.Fatalf("FileOps = %q, want %q", got, want) |
| 353 | } |
| 354 | for i := range want { |
| 355 | if got[i] != want[i] { |
| 356 | t.Errorf("FileOps[%d] = %q, want %q", i, got[i], want[i]) |
| 357 | } |
| 358 | } |
| 359 | if n := Commands(history); n != 1 { |
| 360 | t.Errorf("Commands = %d, want 1 — file ops are not commands", n) |
| 361 | } |
| 362 | } |