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
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
|
// Tests for the provider layer, against httptest servers: the fake engine of
// the backup docs speaks plain OpenAI and cannot say "I am llama-server", nor
// fail on demand. What these tests guard: the URL precedence the `dmr` package
// had (env, then file, then fallback when the first does not answer), the
// llama.cpp probe, and the one-line error mapping — each line of the table in
// PROVIDERS.md §5.7 that this step implements.
package engine
import (
"context"
"encoding/json"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"testing"
"mm/internal/config"
"github.com/firebase/genkit/go/ai"
)
// openaiServer is a minimal OpenAI-compatible server: /models lists `models`,
// every POST answers `status` with `body`. Enough to drive Resolve, Probe and
// Explain; Generate's happy path is covered by engine_test.go's fake model.
func openaiServer(t *testing.T, status int, body string, models ...string) *httptest.Server {
t.Helper()
mux := http.NewServeMux()
mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/models") {
var data []map[string]string
for _, m := range models {
data = append(data, map[string]string{"id": m, "object": "model"})
}
_ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})
return
}
if r.Method == http.MethodPost {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_, _ = w.Write([]byte(body))
return
}
http.NotFound(w, r)
})
srv := httptest.NewServer(mux)
t.Cleanup(srv.Close)
return srv
}
// closedURL returns a URL nothing listens on: a server started and stopped, so
// the port was free a moment ago.
func closedURL(t *testing.T) string {
t.Helper()
srv := httptest.NewServer(http.NotFoundHandler())
url := srv.URL
srv.Close()
return url
}
func ptr(s string) *string { return &s }
// clearEnv: the precedence tests must not depend on the developer's shell,
// where DMR_BASE_URL is often exported.
func clearEnv(t *testing.T) {
t.Helper()
for _, k := range []string{"AGENT_BASE_URL", "DMR_BASE_URL", "LLAMA_API_KEY", "AGENT_PROVIDER", "AGENT_MODEL"} {
t.Setenv(k, "")
}
}
func TestLookupListsKnownProviders(t *testing.T) {
_, err := Lookup("ollama")
if err == nil || !strings.Contains(err.Error(), "dmr, llamacpp") {
t.Fatalf("Lookup(ollama) = %v, want an error naming the known providers", err)
}
}
// TestResolveURLPrecedence replays the rules of the former config.DMRBaseURL —
// env wins without probing, the fallback is taken only when baseUrl does not
// answer — and checks that llamacpp neither inherits DMR's URL nor its env var.
func TestResolveURLPrecedence(t *testing.T) {
clearEnv(t)
up := openaiServer(t, 200, "{}", "fake-model")
down := closedURL(t)
cases := []struct {
name string
provider string
cfg config.Config
env map[string]string
wantURL string
wantKey string
}{
{"dmr keeps a reachable baseUrl", "dmr", config.Config{BaseURL: up.URL}, nil, up.URL, "not-needed"},
{"dmr falls back when baseUrl is down", "dmr", config.Config{BaseURL: down, Fallback: ptr(up.URL)}, nil, up.URL, "not-needed"},
{"an explicit empty fallback disables it", "dmr", config.Config{BaseURL: down, Fallback: ptr("")}, nil, down, "not-needed"},
{"AGENT_BASE_URL wins without probing", "dmr", config.Config{BaseURL: up.URL, Fallback: ptr("")}, map[string]string{"AGENT_BASE_URL": down}, down, "not-needed"},
{"DMR_BASE_URL still works for dmr", "dmr", config.Config{BaseURL: up.URL, Fallback: ptr("")}, map[string]string{"DMR_BASE_URL": down}, down, "not-needed"},
{"llamacpp default is llama-server's port", "llamacpp", config.Config{}, nil, "http://127.0.0.1:8080/v1", "not-needed"},
{"llamacpp ignores DMR_BASE_URL", "llamacpp", config.Config{}, map[string]string{"DMR_BASE_URL": down}, "http://127.0.0.1:8080/v1", "not-needed"},
{"llamacpp reads LLAMA_API_KEY when set", "llamacpp", config.Config{}, map[string]string{"LLAMA_API_KEY": "secret"}, "http://127.0.0.1:8080/v1", "secret"},
{"apiKeyEnv names another variable", "llamacpp", config.Config{APIKeyEnv: "MY_KEY"}, map[string]string{"MY_KEY": "k2"}, "http://127.0.0.1:8080/v1", "k2"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
clearEnv(t)
for k, v := range tc.env {
t.Setenv(k, v)
}
p, _ := Lookup(tc.provider)
tc.cfg.Model = "fake-model"
b, err := p.Resolve(tc.cfg)
if err != nil {
t.Fatalf("Resolve: %v", err)
}
if b.BaseURL != tc.wantURL {
t.Errorf("BaseURL = %q, want %q", b.BaseURL, tc.wantURL)
}
if b.APIKey != tc.wantKey {
t.Errorf("APIKey = %q, want %q", b.APIKey, tc.wantKey)
}
})
}
}
// TestProbeLlamaCpp: /props sits one level above /v1 and carries the SERVED
// context size; the config's value, when given, wins over it.
func TestProbeLlamaCpp(t *testing.T) {
mux := http.NewServeMux()
mux.HandleFunc("/props", func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"default_generation_settings":{"n_ctx":32768},"total_slots":1,"chat_template":"{{ messages }}"}`))
})
mux.HandleFunc("/v1/models", func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"object":"list","data":[{"id":"qwen2.5-coder","object":"model"}]}`))
})
srv := httptest.NewServer(mux)
defer srv.Close()
p, _ := Lookup("llamacpp")
info := p.Probe(context.Background(), Backend{Provider: "llamacpp", BaseURL: srv.URL + "/v1", Model: "qwen2.5-coder"})
if !info.Reachable || info.ContextWindow != 32768 || info.ContextSource != "/props" {
t.Errorf("Probe = %+v, want reachable, 32768 from /props", info)
}
if len(info.Warnings) != 0 {
t.Errorf("unexpected warnings: %q", info.Warnings)
}
info = p.Probe(context.Background(), Backend{Provider: "llamacpp", BaseURL: srv.URL + "/v1", Model: "qwen2.5-coder", ContextWindow: 4096})
if info.ContextWindow != 4096 || info.ContextSource != "config" {
t.Errorf("with a config hint: %+v, want 4096 from config", info)
}
}
// TestEnsureContextWindowReprobes: the start-up probe runs once, and a
// llama-server started AFTER the agent left `ctx: unknown` for the whole
// session. The lazy re-probe must turn 0 into n_ctx once /props answers — and
// must stop asking as soon as the window is known, or when the config gave
// it: the hit counter on /props is the measure.
func TestEnsureContextWindowReprobes(t *testing.T) {
var propsHits int
up := false // false = llama-server not ready yet: /props answers 503
mux := http.NewServeMux()
mux.HandleFunc("/props", func(w http.ResponseWriter, _ *http.Request) {
propsHits++
if !up {
w.WriteHeader(http.StatusServiceUnavailable)
return
}
_, _ = w.Write([]byte(`{"default_generation_settings":{"n_ctx":32768},"total_slots":1,"chat_template":"{{ messages }}"}`))
})
mux.HandleFunc("/v1/models", func(w http.ResponseWriter, _ *http.Request) {
_, _ = w.Write([]byte(`{"object":"list","data":[{"id":"m","object":"model"}]}`))
})
srv := httptest.NewServer(mux)
defer srv.Close()
p, _ := Lookup("llamacpp")
e := &Engine{Provider: p, Backend: Backend{Provider: "llamacpp", BaseURL: srv.URL + "/v1", Model: "m"}}
// Server not ready: still unknown, nothing learned, but the probe was made.
if w, src := e.EnsureContextWindow(context.Background()); w != 0 || src != "" || e.ContextWindow != 0 {
t.Errorf("before the server is up: got (%d, %q), ContextWindow=%d, want (0, \"\"), 0", w, src, e.ContextWindow)
}
if propsHits != 1 {
t.Fatalf("/props hits = %d after the first re-probe, want 1", propsHits)
}
// Server up: the window is learned, remembered, and its origin is named.
up = true
if w, src := e.EnsureContextWindow(context.Background()); w != 32768 || src != "/props" || e.ContextWindow != 32768 {
t.Errorf("after the server is up: got (%d, %q), ContextWindow=%d, want (32768, \"/props\"), 32768", w, src, e.ContextWindow)
}
if propsHits != 2 {
t.Fatalf("/props hits = %d after the second re-probe, want 2", propsHits)
}
// Known window: answered from memory, no request, no origin to announce.
if w, src := e.EnsureContextWindow(context.Background()); w != 32768 || src != "" {
t.Errorf("known window: got (%d, %q), want (32768, \"\")", w, src)
}
if propsHits != 2 {
t.Errorf("/props hits = %d after asking with a known window, want still 2", propsHits)
}
// A window from the config is known from the start: never probed at all.
propsHits = 0
known := &Engine{Provider: p, Backend: Backend{Provider: "llamacpp", BaseURL: srv.URL + "/v1", Model: "m", ContextWindow: 4096}, ContextWindow: 4096}
if w, src := known.EnsureContextWindow(context.Background()); w != 4096 || src != "" || propsHits != 0 {
t.Errorf("config window: got (%d, %q) with %d /props hit(s), want (4096, \"\") and 0", w, src, propsHits)
}
}
// TestProbeUnreachableWarns: a server that is down is a warning with the start
// command, not a failure — the demo often starts the server second.
func TestProbeUnreachableWarns(t *testing.T) {
for _, name := range []string{"dmr", "llamacpp"} {
p, _ := Lookup(name)
info := p.Probe(context.Background(), Backend{Provider: name, BaseURL: closedURL(t) + "/v1", Model: "m"})
if info.Reachable || len(info.Warnings) != 1 || !strings.Contains(info.Warnings[0], "nothing answers") {
t.Errorf("%s: Probe on a closed port = %+v", name, info)
}
}
}
// TestProbeModelsWarnsOnMissingModel: DMR lists what is pulled; asking for
// something else fails only once the request is sent, after a long spinner.
func TestProbeModelsWarnsOnMissingModel(t *testing.T) {
srv := openaiServer(t, 200, "{}", "ai/other")
p, _ := Lookup("dmr")
info := p.Probe(context.Background(), Backend{Provider: "dmr", BaseURL: srv.URL, Model: "ai/qwen2.5-coder"})
if !info.Reachable || len(info.Warnings) != 1 || !strings.Contains(info.Warnings[0], "docker model pull ai/qwen2.5-coder") {
t.Errorf("Probe = %+v, want the pull hint", info)
}
}
// TestExplainThroughTheStack drives a real Open + Generate against servers that
// fail on purpose, and checks the ONE line the REPL would print. Going through
// Genkit matters: it is what wraps the OpenAI client's error, and the mapping
// has to survive that wrapping.
func TestExplainThroughTheStack(t *testing.T) {
cases := []struct {
name string
provider string
status int
body string
want string
}{
{"llama-server without --jinja", "llamacpp", 500,
`{"error":{"code":500,"message":"tools param requires --jinja flag","type":"server_error"}}`,
"tool calls need llama-server started with --jinja"},
{"401", "llamacpp", 401, `{"error":{"message":"Invalid API Key","type":"authentication_error"}}`,
"authentication failed"},
{"404 on dmr names the pull command", "dmr", 404, `{"error":{"message":"model not found","type":"not_found"}}`,
"docker model pull fake-model"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
srv := openaiServer(t, tc.status, tc.body, "fake-model")
got := explainAgainst(t, tc.provider, srv.URL)
if !strings.Contains(got, tc.want) {
t.Errorf("Explain = %q, want it to contain %q", got, tc.want)
}
})
}
t.Run("connection refused", func(t *testing.T) {
got := explainAgainst(t, "llamacpp", closedURL(t))
if !strings.Contains(got, "nothing answers") || !strings.Contains(got, "--jinja --port 8080") {
t.Errorf("Explain = %q, want the start hint", got)
}
})
}
// TestExplain401CitesTheResolvedVariable guards the phantom variable: the first
// Explain named $AGENT_API_KEY on a 401, a variable nothing read. The message
// must name the variable Resolve used — the yaml's `apiKeyEnv` when set — and
// name none when the provider has none, rather than invent one.
func TestExplain401CitesTheResolvedVariable(t *testing.T) {
body := `{"error":{"message":"Invalid API Key","type":"authentication_error"}}`
t.Run("yaml apiKeyEnv is the name cited", func(t *testing.T) {
clearEnv(t)
t.Setenv("MY_KEY", "wrong")
srv := openaiServer(t, 401, body, "fake-model")
p, _ := Lookup("llamacpp")
b, err := p.Resolve(config.Config{Model: "fake-model", BaseURL: srv.URL, Fallback: ptr(""), APIKeyEnv: "MY_KEY"})
if err != nil {
t.Fatal(err)
}
got := explainWith(t, p, b)
if !strings.Contains(got, "$MY_KEY") || strings.Contains(got, "LLAMA_API_KEY") {
t.Errorf("Explain = %q, want it to cite $MY_KEY and nothing else", got)
}
})
t.Run("dmr has no variable: none is invented", func(t *testing.T) {
clearEnv(t)
srv := openaiServer(t, 401, body, "fake-model")
p, _ := Lookup("dmr")
b, err := p.Resolve(config.Config{Model: "fake-model", BaseURL: srv.URL, Fallback: ptr("")})
if err != nil {
t.Fatal(err)
}
got := explainWith(t, p, b)
if strings.Contains(got, "$") || !strings.Contains(got, "should not need a key") {
t.Errorf("Explain = %q, want no $VARIABLE and the check-baseUrl hint", got)
}
})
}
// TestResolveKeyPolicy exercises the keyRequired branch with a throwaway
// provider built directly — not registered, so the global registry stays as the
// other tests expect it. Neither shipped provider requires a key, so without
// this the branch a paid API will rely on has never run.
func TestResolveKeyPolicy(t *testing.T) {
clearEnv(t)
paid := &openaiCompat{name: "paid", baseURL: "https://paid.example/v1", apiKeyEnv: "PAID_API_KEY", keyRequired: true, dummyKey: "must-not-appear"}
free := &openaiCompat{name: "free", baseURL: "http://127.0.0.1:1/v1", apiKeyEnv: "FREE_API_KEY", dummyKey: "not-needed"}
cases := []struct {
name string
p *openaiCompat
cfg config.Config
env map[string]string
wantErr string // "" = success
wantKey string
wantEnv string
}{
{"required and unset: the error names the variable", paid, config.Config{}, nil,
"paid: PAID_API_KEY is not set", "", "PAID_API_KEY"},
{"required and set: the value, not the dummy", paid, config.Config{}, map[string]string{"PAID_API_KEY": "sk-live"},
"", "sk-live", "PAID_API_KEY"},
{"required, yaml apiKeyEnv redirects to another variable", paid, config.Config{APIKeyEnv: "OTHER_KEY"}, map[string]string{"OTHER_KEY": "sk-other"},
"", "sk-other", "OTHER_KEY"},
{"required, yaml apiKeyEnv unset: names THAT variable", paid, config.Config{APIKeyEnv: "OTHER_KEY"}, nil,
"paid: OTHER_KEY is not set", "", "OTHER_KEY"},
{"not required and unset: the dummy key", free, config.Config{}, nil,
"", "not-needed", "FREE_API_KEY"},
{"not required and set: the value", free, config.Config{}, map[string]string{"FREE_API_KEY": "abc"},
"", "abc", "FREE_API_KEY"},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
for _, k := range []string{"PAID_API_KEY", "FREE_API_KEY", "OTHER_KEY"} {
t.Setenv(k, "")
}
for k, v := range tc.env {
t.Setenv(k, v)
}
tc.cfg.Model = "m"
tc.cfg.Fallback = ptr("") // no probe: the URLs here do not exist
b, err := tc.p.Resolve(tc.cfg)
switch {
case tc.wantErr != "" && (err == nil || !strings.Contains(err.Error(), tc.wantErr)):
t.Fatalf("Resolve error = %v, want it to contain %q", err, tc.wantErr)
case tc.wantErr == "" && err != nil:
t.Fatalf("Resolve: %v", err)
}
if b.APIKey != tc.wantKey {
t.Errorf("APIKey = %q, want %q", b.APIKey, tc.wantKey)
}
if b.APIKeyEnv != tc.wantEnv {
t.Errorf("APIKeyEnv = %q, want %q", b.APIKeyEnv, tc.wantEnv)
}
})
}
}
// explainAgainst opens the provider on baseURL, asks one question and returns
// the explained error. The spinner and the retry print to stdout: harmless.
func explainAgainst(t *testing.T, provider, baseURL string) string {
t.Helper()
p, _ := Lookup(provider)
return explainWith(t, p, Backend{Provider: provider, BaseURL: baseURL, Model: "fake-model", APIKey: "not-needed"})
}
// explainWith is explainAgainst for a Backend built by Resolve — needed when
// the message under test depends on what Resolve recorded (APIKeyEnv).
func explainWith(t *testing.T, p Provider, b Backend) string {
t.Helper()
g, model, err := p.Open(context.Background(), b)
if err != nil {
t.Fatalf("Open: %v", err)
}
e := &Engine{G: g, Model: model, Provider: p, Backend: b}
_, _, err = e.Generate(context.Background(), []*ai.Message{ai.NewUserTextMessage("hi")}, nil)
if err == nil {
t.Fatal("Generate succeeded against a failing server")
}
return e.Explain(err)
}
// TestConfigCompat: an agent.yaml written for part 07 — no `provider`, an
// explicit `fallback: ""` — must load with the same meaning; and a llamacpp
// file with no baseUrl must land on llama-server's port, not DMR's.
func TestConfigCompat(t *testing.T) {
clearEnv(t)
saved := config.Cfg
defer func() { config.Cfg = saved }()
dir := t.TempDir()
write := func(name, body string) string {
path := filepath.Join(dir, name)
if err := os.WriteFile(path, []byte(body), 0o644); err != nil {
t.Fatal(err)
}
return path
}
old := write("probe.yaml", "model: fake-model\nbaseUrl: http://127.0.0.1:18234/engines/v1\nfallback: \"\"\nmaxOutput: 16000\nmaxTurns: 25\nsystem: |\n Probe.\nsampling:\n temperature: 0.0\n")
if _, err := config.Load(old); err != nil {
t.Fatalf("Load(07 config): %v", err)
}
if config.Cfg.Provider != "dmr" {
t.Errorf("Provider = %q, want dmr by default", config.Cfg.Provider)
}
p, _ := Lookup(config.Cfg.Provider)
b, err := p.Resolve(config.Cfg)
if err != nil {
t.Fatal(err)
}
if b.BaseURL != "http://127.0.0.1:18234/engines/v1" {
t.Errorf("BaseURL = %q: an explicit empty fallback must keep the file's URL even when it is down", b.BaseURL)
}
config.Cfg = saved
llama := write("llama.yaml", "provider: llamacpp\nmodel: qwen2.5-coder\ncontextWindow: 16384\n")
if _, err := config.Load(llama); err != nil {
t.Fatalf("Load(llamacpp config): %v", err)
}
p, _ = Lookup(config.Cfg.Provider)
b, err = p.Resolve(config.Cfg)
if err != nil {
t.Fatal(err)
}
if b.BaseURL != "http://127.0.0.1:8080/v1" || b.ContextWindow != 16384 {
t.Errorf("Backend = %+v, want llama-server's default URL and the 16384 hint", b)
}
config.Cfg = saved
bad := write("bad.yaml", "provider: ollama\nmodel: m\n")
if _, err := config.Load(bad); err != nil {
t.Fatalf("Load must accept an unknown provider (engine.New refuses it): %v", err)
}
if _, err := Lookup(config.Cfg.Provider); err == nil {
t.Error("Lookup(ollama) must fail")
}
}
|