1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "net/http"
8 "net/http/httptest"
9 "os"
10 "path/filepath"
11 "slices"
12 "strings"
13 "sync/atomic"
14 "testing"
15
16 "github.com/bnema/zut/packages/provider"
17 "github.com/bnema/zut/packages/provider/auth"
18)
19
20// loadTestModelsJSON installs a models.json under ZUT_HOME and restores the
21// provider catalog and custom-provider registry when the test ends.
22func loadTestModelsJSON(t *testing.T, body string) {
23 t.Helper()
24 preserveProviderCatalog(t)
25 if err := os.WriteFile(UserModelsPath(), []byte(body), 0o644); err != nil {
26 t.Fatal(err)
27 }
28 provider.SetManagedModelsForProvider("m4", nil)
29 provider.SetManagedModelsForProvider("m5", nil)
30 LoadUserModels()
31 t.Cleanup(func() {
32 _ = os.WriteFile(UserModelsPath(), []byte(`{"providers":{}}`), 0o644)
33 LoadUserModels()
34 })
35}
36
37func newModelsServer(t *testing.T, requests *atomic.Int32, wantAuth string, ids ...string) *httptest.Server {
38 t.Helper()
39 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
40 requests.Add(1)
41 if r.URL.Path != "/v1/models" {
42 t.Errorf("unexpected path %s", r.URL.Path)
43 }
44 if got := r.Header.Get("Authorization"); got != wantAuth {
45 t.Errorf("Authorization = %q, want %q", got, wantAuth)
46 }
47 var data []map[string]string
48 for _, id := range ids {
49 data = append(data, map[string]string{"id": id})
50 }
51 _ = json.NewEncoder(w).Encode(map[string]any{"object": "list", "data": data})
52 }))
53 t.Cleanup(server.Close)
54 return server
55}
56
57func TestCustomDiscoveryKeepsTLSVerificationWithInsecureInference(t *testing.T) {
58 t.Setenv("ZUT_HOME", t.TempDir())
59 if err := SaveConfig(Config{Insecure: true}); err != nil {
60 t.Fatal(err)
61 }
62 client := customDiscoveryHTTPClient()
63 if transport, ok := client.Transport.(*http.Transport); ok && transport.TLSClientConfig != nil && transport.TLSClientConfig.InsecureSkipVerify {
64 t.Fatal("discovery inherited insecure inference TLS")
65 }
66 server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
67 _, _ = w.Write([]byte(`{"data":[{"id":"test-model"}]}`))
68 }))
69 defer server.Close()
70 if _, err := provider.DiscoverCustomProviderWithClient(context.Background(), client, "test-provider", server.URL, ""); err == nil {
71 t.Fatal("discovery trusted a self-signed server")
72 }
73}
74
75func TestRefreshCustomProviderModelsIsOptIn(t *testing.T) {
76 t.Setenv("ZUT_HOME", t.TempDir())
77 t.Setenv("M4_API_KEY", "")
78 var requests atomic.Int32
79 server := newModelsServer(t, &requests, "", "qwen3-coder")
80 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
81 "m4": {"baseUrl": %q, "discover": true},
82 "m5": {"baseUrl": %q, "models": [{"id": "pinned"}]}
83 }}`, server.URL+"/v1", server.URL+"/v1"))
84
85 if !CustomDiscoveryConfigured() {
86 t.Fatal("discovery not reported as configured")
87 }
88 if got := discoverableCustomProviders(); len(got) != 1 || got[0] != "m4" {
89 t.Fatalf("discoverable = %v", got)
90 }
91 if err := RefreshCustomProviderModels(context.Background()); err != nil {
92 t.Fatal(err)
93 }
94 if requests.Load() != 1 {
95 t.Fatalf("requests = %d, want 1 (only m4)", requests.Load())
96 }
97 m, err := provider.FindModel("m4", "qwen3-coder")
98 if err != nil {
99 t.Fatal(err)
100 }
101 if m.Source != "live" || m.BaseURL != server.URL+"/v1" || m.ContextWindow <= 0 || m.MaxOutput <= 0 {
102 t.Fatalf("model = %+v", m)
103 }
104 if _, err := provider.FindModel("m5", "qwen3-coder"); err == nil {
105 t.Fatal("m5 gained discovered models without opting in")
106 }
107 if _, err := provider.FindModel("m5", "pinned"); err != nil {
108 t.Fatal("m5 user model lost")
109 }
110}
111
112func TestDiscoveryDefaultsOffNoRequestsAtStartup(t *testing.T) {
113 t.Setenv("ZUT_HOME", t.TempDir())
114 var requests atomic.Int32
115 server := newModelsServer(t, &requests, "")
116 loadTestModelsJSON(t, fmt.Sprintf(`{"providers":{"m5":{"baseUrl":%q,"models":[{"id":"pinned"}]}}}`, server.URL+"/v1"))
117 refreshCustomProvidersBackground()
118 if requests.Load() != 0 || CustomDiscoveryConfigured() {
119 t.Fatalf("discovery ran without opt-in (requests=%d)", requests.Load())
120 }
121}
122
123func TestRefreshCustomProviderModelsUsesStoredKeyAndKeepsUserOverrides(t *testing.T) {
124 t.Setenv("ZUT_HOME", t.TempDir())
125 t.Setenv("M4_API_KEY", "")
126 var requests atomic.Int32
127 server := newModelsServer(t, &requests, "Bearer dummy", "qwen3-coder", "big-model")
128 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
129 "m4": {"baseUrl": %q, "discover": true, "models": [{"id": "big-model", "contextWindow": 200000, "maxTokens": 16000}]}
130 }}`, server.URL+"/v1"))
131 if err := AuthStoreFor().SetAPIKey("m4", "dummy"); err != nil {
132 t.Fatal(err)
133 }
134 if err := RefreshCustomProviderModels(context.Background()); err != nil {
135 t.Fatal(err)
136 }
137 if requests.Load() != 1 {
138 t.Fatalf("requests = %d", requests.Load())
139 }
140 big, err := provider.FindModel("m4", "big-model")
141 if err != nil {
142 t.Fatal(err)
143 }
144 if big.Source != "user" || big.ContextWindow != 200000 || big.MaxOutput != 16000 {
145 t.Fatalf("models.json override lost: %+v", big)
146 }
147 if _, err := provider.FindModel("m4", "qwen3-coder"); err != nil {
148 t.Fatal("discovered model missing")
149 }
150}
151
152func TestRefreshCustomProviderModelsExplicitKeylessIgnoresStoredKey(t *testing.T) {
153 t.Setenv("ZUT_HOME", t.TempDir())
154 t.Setenv("M4_API_KEY", "env-key")
155 var requests atomic.Int32
156 server := newModelsServer(t, &requests, "", "qwen3-coder")
157 loadTestModelsJSON(t, fmt.Sprintf(`{"providers":{"m4":{"baseUrl":%q,"discover":true,"auth":"none"}}}`, server.URL+"/v1"))
158 if err := AuthStoreFor().SetAPIKey("m4", "stored-key"); err != nil {
159 t.Fatal(err)
160 }
161 if err := RefreshCustomProviderModels(context.Background()); err != nil {
162 t.Fatal(err)
163 }
164 if requests.Load() != 1 {
165 t.Fatalf("requests = %d, want 1", requests.Load())
166 }
167}
168
169func TestRefreshCustomProviderModelsBackgroundSkipsKeyCommand(t *testing.T) {
170 home := t.TempDir()
171 t.Setenv("ZUT_HOME", home)
172 t.Setenv("M4_API_KEY", "")
173 var requests atomic.Int32
174 server := newModelsServer(t, &requests, "", "m")
175 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server.URL))
176 creds := auth.Credentials{AdditionalAPIKeyCreds: map[string]auth.ProviderCreds{
177 "m4": {APIKeyCommand: &auth.APIKeyCommand{Program: filepath.Join(home, "does-not-exist")}},
178 }}
179 encoded, err := json.Marshal(creds)
180 if err != nil {
181 t.Fatal(err)
182 }
183 if err := os.WriteFile(AuthPath(), encoded, 0o600); err != nil {
184 t.Fatal(err)
185 }
186 if err := refreshCustomProviderModels(context.Background(), apiKeyCommandSkip); err != nil {
187 t.Fatal(err)
188 }
189 if requests.Load() != 0 {
190 t.Fatal("background refresh probed a provider whose key needs a command")
191 }
192 err = RefreshCustomProviderModels(context.Background())
193 if err == nil || !strings.Contains(err.Error(), "m4") {
194 t.Fatalf("command failure not reported: %v", err)
195 }
196 if requests.Load() != 0 {
197 t.Fatal("refresh probed after key command failed")
198 }
199}
200
201func TestRefreshCustomProviderModelsKeepsSnapshotOnFailure(t *testing.T) {
202 t.Setenv("ZUT_HOME", t.TempDir())
203 t.Setenv("M4_API_KEY", "")
204 t.Setenv("M5_API_KEY", "")
205 var requests atomic.Int32
206 good := newModelsServer(t, &requests, "", "alive")
207 failing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
208 http.Error(w, "nope", http.StatusServiceUnavailable)
209 }))
210 defer failing.Close()
211 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
212 "m4": {"baseUrl": %q, "discover": true},
213 "m5": {"baseUrl": %q, "discover": true}
214 }}`, good.URL, failing.URL))
215 provider.SetManagedModelsForProvider("m5", []provider.Model{{Provider: "m5", ID: "previous", Source: "live"}})
216
217 err := RefreshCustomProviderModels(context.Background())
218 if err == nil || !strings.Contains(err.Error(), "m5 discovery HTTP 503") {
219 t.Fatalf("err = %v", err)
220 }
221 if _, err := provider.FindModel("m4", "alive"); err != nil {
222 t.Fatal("healthy provider not refreshed when a sibling failed")
223 }
224 if _, err := provider.FindModel("m5", "previous"); err != nil {
225 t.Fatal("failed refresh dropped the previous snapshot")
226 }
227}
228
229func TestValidateAndRepairConfig_PreservesDiscoverableCustomModel(t *testing.T) {
230 home := t.TempDir()
231 t.Setenv("ZUT_HOME", home)
232 loadTestModelsJSON(t, `{"providers": {
233 "m4": {"baseUrl": "http://127.0.0.1:1/v1", "discover": true},
234 "m5": {"baseUrl": "http://127.0.0.1:1/v1", "models": [{"id": "pinned"}]}
235 }}`)
236 write := func(c Config) {
237 t.Helper()
238 b, _ := json.Marshal(c)
239 if err := os.WriteFile(filepath.Join(home, "config.json"), b, 0o644); err != nil {
240 t.Fatal(err)
241 }
242 }
243 write(Config{Provider: "m4", Model: "qwen3-coder"})
244 ValidateAndRepairConfig()
245 out, _ := LoadConfig()
246 if out.Provider != "m4" || out.Model != "qwen3-coder" {
247 t.Fatalf("discoverable custom model repaired away: %+v", out)
248 }
249 write(Config{Provider: "m5", Model: "gone"})
250 ValidateAndRepairConfig()
251 out, _ = LoadConfig()
252 if out.Provider != "m5" || out.Model != "pinned" {
253 t.Fatalf("non-discoverable custom model not repaired: %+v", out)
254 }
255}
256
257func TestResolveRestoresDiscoveredCustomModelAfterRestart(t *testing.T) {
258 t.Setenv("ZUT_HOME", t.TempDir())
259 t.Setenv("M4_API_KEY", "")
260 var requests atomic.Int32
261 server := newModelsServer(t, &requests, "", "qwen3-coder")
262 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server.URL+"/v1"))
263
264 r, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true)
265 if err != nil {
266 t.Fatal(err)
267 }
268 if requests.Load() != 1 {
269 t.Fatalf("requests = %d, want a restore probe", requests.Load())
270 }
271 if r.Provider != "m4" || r.Model != "qwen3-coder" || r.BaseURL != server.URL+"/v1" || r.MaxOutput <= 0 {
272 t.Fatalf("resolved = %+v", r)
273 }
274 if !r.HasCredential() || r.Credential != "" {
275 t.Fatalf("keyless discovered provider should be usable without a placeholder: %q", r.Credential)
276 }
277 if _, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true); err != nil {
278 t.Fatal(err)
279 }
280 if requests.Load() != 1 {
281 t.Fatalf("requests = %d, want no extra probe", requests.Load())
282 }
283}
284
285func streamOnce(t *testing.T, r Resolved, model string) {
286 t.Helper()
287 events, err := r.NewClient().Stream(context.Background(), provider.Request{
288 Model: model, Messages: []provider.Message{{Role: provider.RoleUser, Content: []provider.Content{provider.TextBlock{Text: "hello"}}}},
289 })
290 if err != nil {
291 t.Fatal(err)
292 }
293 for range events {
294 }
295}
296
297func TestCustomProviderKeylessOmitsAuthorization(t *testing.T) {
298 for _, api := range []string{"openai", "openai-responses", "anthropic"} {
299 t.Run(api, func(t *testing.T) {
300 t.Setenv("ZUT_HOME", t.TempDir())
301 t.Setenv("UNSLOTH_API_KEY", "")
302 got := make(chan http.Header, 1)
303 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
304 got <- r.Header.Clone()
305 w.Header().Set("Content-Type", "text/event-stream")
306 }))
307 t.Cleanup(srv.Close)
308 loadTestModelsJSON(t, `{"providers":{"unsloth":{"baseUrl":"`+srv.URL+`/v1","api":"`+api+`","models":[{"id":"qista"}]}}}`)
309 r, err := Resolve(Args{Provider: "unsloth", Model: "qista", NoSkill: true}, true)
310 if err != nil {
311 t.Fatal(err)
312 }
313 if r.Credential != "" || !r.HasCredential() {
314 t.Fatalf("credential = %q usable=%v", r.Credential, r.HasCredential())
315 }
316 streamOnce(t, r, "qista")
317 h := <-got
318 if _, ok := h["Authorization"]; ok {
319 t.Fatal("Authorization sent")
320 }
321 if _, ok := h["X-Api-Key"]; ok {
322 t.Fatal("X-Api-Key sent")
323 }
324 })
325 }
326}
327
328func TestCustomProviderExplicitKeylessIgnoresInheritedKeys(t *testing.T) {
329 for _, api := range []string{"openai", "openai-responses", "anthropic"} {
330 t.Run(api, func(t *testing.T) {
331 home := t.TempDir()
332 t.Setenv("ZUT_HOME", home)
333 t.Setenv("UNSLOTH_API_KEY", "env-key")
334 got := make(chan http.Header, 1)
335 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
336 got <- r.Header.Clone()
337 w.Header().Set("Content-Type", "text/event-stream")
338 }))
339 t.Cleanup(srv.Close)
340 loadTestModelsJSON(t, `{"providers":{"unsloth":{"baseUrl":"`+srv.URL+`/v1","api":"`+api+`","auth":"none","models":[{"id":"qista"}]}}}`)
341 // A stored command-backed key must never run for a keyless provider.
342 marker := filepath.Join(home, "ran")
343 creds := auth.Credentials{AdditionalAPIKeyCreds: map[string]auth.ProviderCreds{
344 "unsloth": {APIKeyCommand: &auth.APIKeyCommand{Program: "/bin/sh", Args: []string{"-c", "touch " + marker + "; echo secret"}}},
345 }}
346 encoded, _ := json.Marshal(creds)
347 if err := os.WriteFile(AuthPath(), encoded, 0o600); err != nil {
348 t.Fatal(err)
349 }
350 r, err := Resolve(Args{Provider: "unsloth", Model: "qista", APIKey: "cli-key", NoSkill: true}, true)
351 if err != nil {
352 t.Fatal(err)
353 }
354 if r.Credential != "" || !r.HasCredential() {
355 t.Fatalf("keyless resolved credential = %q, usable = %v", r.Credential, r.HasCredential())
356 }
357 streamOnce(t, r, "qista")
358 h := <-got
359 if _, ok := h["Authorization"]; ok {
360 t.Fatal("Authorization sent")
361 }
362 if _, ok := h["X-Api-Key"]; ok {
363 t.Fatal("X-Api-Key sent")
364 }
365 if _, err := os.Stat(marker); err == nil {
366 t.Fatal("stored key command executed for explicit keyless provider")
367 }
368 stored, err := AuthStoreFor().Load()
369 if err != nil || stored.AdditionalAPIKeyCreds["unsloth"].APIKeyCommand == nil {
370 t.Fatalf("stored credential changed: %+v, %v", stored.AdditionalAPIKeyCreds["unsloth"], err)
371 }
372 })
373 }
374}
375
376func TestCustomProviderWithKeyStillSendsAuthorization(t *testing.T) {
377 t.Setenv("ZUT_HOME", t.TempDir())
378 t.Setenv("UNSLOTH_API_KEY", "")
379 got := make(chan string, 1)
380 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
381 got <- r.Header.Get("Authorization")
382 w.Header().Set("Content-Type", "text/event-stream")
383 }))
384 t.Cleanup(srv.Close)
385 loadTestModelsJSON(t, `{"providers":{"unsloth":{"baseUrl":"`+srv.URL+`/v1","api":"openai","models":[{"id":"qista"}]}}}`)
386 if err := AuthStoreFor().SetAPIKey("unsloth", "stored-key"); err != nil {
387 t.Fatal(err)
388 }
389 r, err := Resolve(Args{Provider: "unsloth", Model: "qista", NoSkill: true}, true)
390 if err != nil {
391 t.Fatal(err)
392 }
393 streamOnce(t, r, "qista")
394 if a := <-got; a != "Bearer stored-key" {
395 t.Fatalf("Authorization = %q", a)
396 }
397}
398
399func TestBuiltinProviderStillRequiresCredentialAndKeepsPlaceholderRules(t *testing.T) {
400 t.Setenv("ZUT_HOME", t.TempDir())
401 t.Setenv("ANTHROPIC_API_KEY", "")
402 t.Setenv("ANTHROPIC_OAUTH_TOKEN", "")
403 // A custom provider named like a built-in must not become keyless.
404 loadTestModelsJSON(t, `{"providers":{"anthropic":{"baseUrl":"http://127.0.0.1:9/v1","auth":"none","models":[{"id":"claude-sonnet-5"}]}}}`)
405 r, err := Resolve(Args{Provider: "anthropic", Model: "claude-sonnet-5", NoSkill: true}, false)
406 if err != nil {
407 t.Fatal(err)
408 }
409 if r.credentialOptional {
410 t.Fatal("built-in provider treated as keyless custom endpoint")
411 }
412}
413
414func TestModelPickerProvidersIncludeKeylessCustomEndpoint(t *testing.T) {
415 t.Setenv("ZUT_HOME", t.TempDir())
416 t.Setenv("UNSLOTH_API_KEY", "")
417 loadTestModelsJSON(t, `{"providers":{"unsloth":{"baseUrl":"http://127.0.0.1:8888/v1","api":"openai","models":[{"id":"qista"}]}}}`)
418 if CredentialAvailable("unsloth") {
419 t.Fatal("test requires no credential")
420 }
421 if !slices.Contains(modelPickerProviders(), "unsloth") {
422 t.Fatal("keyless custom endpoint missing from model picker")
423 }
424 loadTestModelsJSON(t, `{"providers":{"unsloth":{"models":[{"id":"qista","baseUrl":"http://127.0.0.1:8888/v1"}]}}}`)
425 if !slices.Contains(modelPickerProviders(), "unsloth") {
426 t.Fatal("model-level keyless endpoint missing from model picker")
427 }
428}
429
430func TestResolveUnknownModelRefreshesOnlySelectedCustomProvider(t *testing.T) {
431 home := t.TempDir()
432 t.Setenv("ZUT_HOME", home)
433 t.Setenv("M4_API_KEY", "")
434 t.Setenv("M5_API_KEY", "")
435 var m4Requests, m5Requests atomic.Int32
436 m4 := newModelsServer(t, &m4Requests, "", "qwen3-coder")
437 m5 := newModelsServer(t, &m5Requests, "", "other")
438 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
439 "m4": {"baseUrl": %q, "discover": true},
440 "m5": {"baseUrl": %q, "discover": true}
441 }}`, m4.URL+"/v1", m5.URL+"/v1"))
442 // m5 has a command-backed key; resolving m4 must never run it.
443 marker := filepath.Join(home, "m5-command-ran")
444 creds := auth.Credentials{AdditionalAPIKeyCreds: map[string]auth.ProviderCreds{
445 "m5": {APIKeyCommand: &auth.APIKeyCommand{Program: "/bin/sh", Args: []string{"-c", "touch " + marker + "; echo secret"}}},
446 }}
447 encoded, _ := json.Marshal(creds)
448 if err := os.WriteFile(AuthPath(), encoded, 0o600); err != nil {
449 t.Fatal(err)
450 }
451
452 if _, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true); err != nil {
453 t.Fatal(err)
454 }
455 if m4Requests.Load() != 1 {
456 t.Fatalf("selected provider requests = %d, want 1", m4Requests.Load())
457 }
458 if m5Requests.Load() != 0 {
459 t.Fatal("unrelated provider endpoint was contacted")
460 }
461 if _, err := os.Stat(marker); err == nil {
462 t.Fatal("unrelated provider api-key command executed")
463 }
464}
465
466func TestRefreshSingleCustomProviderKeepsSnapshotOnFailure(t *testing.T) {
467 t.Setenv("ZUT_HOME", t.TempDir())
468 t.Setenv("M4_API_KEY", "")
469 failing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
470 http.Error(w, "nope", http.StatusServiceUnavailable)
471 }))
472 defer failing.Close()
473 loadTestModelsJSON(t, fmt.Sprintf(`{"providers":{"m4":{"baseUrl":%q,"discover":true},"m6":{"baseUrl":%q}}}`, failing.URL, failing.URL))
474 provider.SetManagedModelsForProvider("m4", []provider.Model{{Provider: "m4", ID: "previous", Source: "live"}})
475
476 if err := refreshCustomProviderModel(context.Background(), "m4", apiKeyCommandExecute); err == nil {
477 t.Fatal("expected refresh failure")
478 }
479 if _, err := provider.FindModel("m4", "previous"); err != nil {
480 t.Fatal("failed scoped refresh dropped the previous snapshot")
481 }
482 // A provider that did not opt in is never contacted or changed.
483 if err := refreshCustomProviderModel(context.Background(), "m6", apiKeyCommandExecute); err != nil {
484 t.Fatalf("non-discoverable provider refreshed: %v", err)
485 }
486}