1package agent
2
3import (
4 "context"
5 "encoding/json"
6 "fmt"
7 "net/http"
8 "net/http/httptest"
9 "os"
10 "path/filepath"
11 "strings"
12 "sync/atomic"
13 "testing"
14
15 "github.com/patriceckhart/zot/packages/provider"
16 "github.com/patriceckhart/zot/packages/provider/auth"
17)
18
19// loadTestModelsJSON installs a models.json under ZOT_HOME and clears the
20// custom provider registry and managed catalog when the test ends.
21func loadTestModelsJSON(t *testing.T, body string) {
22 t.Helper()
23 if err := os.WriteFile(UserModelsPath(), []byte(body), 0o644); err != nil {
24 t.Fatal(err)
25 }
26 provider.SetLiveModels(nil)
27 provider.SetManagedModels(nil)
28 LoadUserModels()
29 t.Cleanup(func() {
30 _ = os.WriteFile(UserModelsPath(), []byte(`{"providers":{}}`), 0o644)
31 LoadUserModels()
32 provider.SetLiveModels(nil)
33 provider.SetManagedModels(nil)
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 TestRefreshCustomProviderModelsIsOptIn(t *testing.T) {
58 t.Setenv("ZOT_HOME", t.TempDir())
59 t.Setenv("M4_API_KEY", "")
60 t.Setenv("M5_API_KEY", "")
61 var requests atomic.Int32
62 server := newModelsServer(t, &requests, "", "qwen3-coder")
63 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
64 "m4": {"baseUrl": %q, "discover": true},
65 "m5": {"baseUrl": %q, "models": [{"id": "pinned"}]}
66 }}`, server.URL+"/v1", server.URL+"/v1"))
67
68 if !CustomDiscoveryConfigured() {
69 t.Fatal("discovery not reported as configured")
70 }
71 if got := discoverableCustomProviders(); len(got) != 1 || got[0] != "m4" {
72 t.Fatalf("discoverable = %v", got)
73 }
74 if err := RefreshCustomProviderModels(context.Background()); err != nil {
75 t.Fatal(err)
76 }
77 if requests.Load() != 1 {
78 t.Fatalf("requests = %d, want 1 (only m4)", requests.Load())
79 }
80 m, err := provider.FindModel("m4", "qwen3-coder")
81 if err != nil {
82 t.Fatal(err)
83 }
84 if m.Source != "live" || m.BaseURL != server.URL+"/v1" || m.ContextWindow <= 0 || m.MaxOutput <= 0 {
85 t.Fatalf("model = %+v", m)
86 }
87 if _, err := provider.FindModel("m5", "qwen3-coder"); err == nil {
88 t.Fatal("m5 gained discovered models without opting in")
89 }
90 if _, err := provider.FindModel("m5", "pinned"); err != nil {
91 t.Fatal("m5 user model lost")
92 }
93}
94
95func TestRefreshCustomProviderModelsUsesStoredKeyAndKeepsUserOverrides(t *testing.T) {
96 t.Setenv("ZOT_HOME", t.TempDir())
97 t.Setenv("M4_API_KEY", "")
98 var requests atomic.Int32
99 server := newModelsServer(t, &requests, "Bearer dummy", "qwen3-coder", "big-model")
100 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
101 "m4": {"baseUrl": %q, "discover": true, "models": [{"id": "big-model", "contextWindow": 200000, "maxTokens": 16000}]}
102 }}`, server.URL+"/v1"))
103 if err := AuthStoreFor().SetAPIKey("m4", "dummy"); err != nil {
104 t.Fatal(err)
105 }
106
107 if err := RefreshCustomProviderModels(context.Background()); err != nil {
108 t.Fatal(err)
109 }
110 if requests.Load() != 1 {
111 t.Fatalf("requests = %d", requests.Load())
112 }
113 big, err := provider.FindModel("m4", "big-model")
114 if err != nil {
115 t.Fatal(err)
116 }
117 if big.Source != "user" || big.ContextWindow != 200000 || big.MaxOutput != 16000 {
118 t.Fatalf("models.json override lost: %+v", big)
119 }
120 if _, err := provider.FindModel("m4", "qwen3-coder"); err != nil {
121 t.Fatal("discovered model missing")
122 }
123
124 // Env var beats auth.json, matching credential resolution order.
125 t.Setenv("M4_API_KEY", "from-env")
126 server2 := newModelsServer(t, &requests, "Bearer from-env", "other")
127 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server2.URL))
128 if err := RefreshCustomProviderModels(context.Background()); err != nil {
129 t.Fatal(err)
130 }
131 if _, err := provider.FindModel("m4", "other"); err != nil {
132 t.Fatal("refresh did not replace snapshot")
133 }
134 if _, err := provider.FindModel("m4", "qwen3-coder"); err == nil {
135 t.Fatal("stale discovered model retained after successful refresh")
136 }
137}
138
139func TestRefreshCustomProviderModelsExplicitKeylessIgnoresStoredKey(t *testing.T) {
140 t.Setenv("ZOT_HOME", t.TempDir())
141 t.Setenv("M4_API_KEY", "env-key")
142 var requests atomic.Int32
143 server := newModelsServer(t, &requests, "", "qwen3-coder")
144 loadTestModelsJSON(t, fmt.Sprintf(`{"providers":{"m4":{"baseUrl":%q,"discover":true,"auth":"none"}}}`, server.URL+"/v1"))
145 if err := AuthStoreFor().SetAPIKey("m4", "stored-key"); err != nil {
146 t.Fatal(err)
147 }
148 if err := RefreshCustomProviderModels(context.Background()); err != nil {
149 t.Fatal(err)
150 }
151 if requests.Load() != 1 {
152 t.Fatalf("requests = %d, want 1", requests.Load())
153 }
154}
155
156func TestRefreshCustomProviderModelsBackgroundSkipsKeyCommand(t *testing.T) {
157 home := t.TempDir()
158 t.Setenv("ZOT_HOME", home)
159 t.Setenv("M4_API_KEY", "")
160 var requests atomic.Int32
161 server := newModelsServer(t, &requests, "", "m")
162 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server.URL))
163 creds := auth.Credentials{AdditionalAPIKeyCreds: map[string]auth.ProviderCreds{
164 "m4": {APIKeyCommand: &auth.APIKeyCommand{Program: filepath.Join(home, "does-not-exist")}},
165 }}
166 encoded, err := json.Marshal(creds)
167 if err != nil {
168 t.Fatal(err)
169 }
170 if err := os.WriteFile(AuthPath(), encoded, 0o600); err != nil {
171 t.Fatal(err)
172 }
173
174 if err := refreshCustomProviderModels(context.Background(), apiKeyCommandSkip); err != nil {
175 t.Fatal(err)
176 }
177 if requests.Load() != 0 {
178 t.Fatal("background refresh probed a provider whose key needs a command")
179 }
180 // Interactive refresh runs the command and reports its failure
181 // instead of silently probing without a key.
182 err = RefreshCustomProviderModels(context.Background())
183 if err == nil || !strings.Contains(err.Error(), "m4") {
184 t.Fatalf("command failure not reported: %v", err)
185 }
186 if requests.Load() != 0 {
187 t.Fatal("refresh probed after key command failed")
188 }
189}
190
191func TestRefreshCustomProviderModelsKeepsSnapshotOnFailure(t *testing.T) {
192 t.Setenv("ZOT_HOME", t.TempDir())
193 t.Setenv("M4_API_KEY", "")
194 t.Setenv("M5_API_KEY", "")
195 var requests atomic.Int32
196 good := newModelsServer(t, &requests, "", "alive")
197 failing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
198 http.Error(w, "nope", http.StatusServiceUnavailable)
199 }))
200 defer failing.Close()
201 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
202 "m4": {"baseUrl": %q, "discover": true},
203 "m5": {"baseUrl": %q, "discover": true}
204 }}`, good.URL, failing.URL))
205 provider.SetManagedModelsForProvider("m5", []provider.Model{{Provider: "m5", ID: "previous", Source: "live"}})
206
207 err := RefreshCustomProviderModels(context.Background())
208 if err == nil || !strings.Contains(err.Error(), "m5 discovery HTTP 503") {
209 t.Fatalf("err = %v", err)
210 }
211 if _, err := provider.FindModel("m4", "alive"); err != nil {
212 t.Fatal("healthy provider not refreshed when a sibling failed")
213 }
214 if _, err := provider.FindModel("m5", "previous"); err != nil {
215 t.Fatal("failed refresh dropped the previous snapshot")
216 }
217}
218
219func TestValidateAndRepairConfig_PreservesDiscoverableCustomModel(t *testing.T) {
220 home := t.TempDir()
221 t.Setenv("ZOT_HOME", home)
222 loadTestModelsJSON(t, `{"providers": {
223 "m4": {"baseUrl": "http://127.0.0.1:1/v1", "discover": true},
224 "m5": {"baseUrl": "http://127.0.0.1:1/v1", "models": [{"id": "pinned"}]}
225 }}`)
226
227 write := func(c Config) {
228 t.Helper()
229 b, _ := json.Marshal(c)
230 if err := os.WriteFile(filepath.Join(home, "config.json"), b, 0o644); err != nil {
231 t.Fatal(err)
232 }
233 }
234
235 // Discovery has not run yet at repair time, so the model is unknown.
236 write(Config{Provider: "m4", Model: "qwen3-coder"})
237 ValidateAndRepairConfig()
238 out, _ := LoadConfig()
239 if out.Provider != "m4" || out.Model != "qwen3-coder" {
240 t.Fatalf("discoverable custom model repaired away: %+v", out)
241 }
242
243 // A provider without discovery keeps the existing repair behavior.
244 write(Config{Provider: "m5", Model: "gone"})
245 ValidateAndRepairConfig()
246 out, _ = LoadConfig()
247 if out.Provider != "m5" || out.Model != "pinned" {
248 t.Fatalf("non-discoverable custom model not repaired: %+v", out)
249 }
250}
251
252func TestResolveRestoresDiscoveredCustomModelAfterRestart(t *testing.T) {
253 t.Setenv("ZOT_HOME", t.TempDir())
254 t.Setenv("M4_API_KEY", "")
255 var requests atomic.Int32
256 server := newModelsServer(t, &requests, "", "qwen3-coder")
257 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server.URL+"/v1"))
258
259 r, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true)
260 if err != nil {
261 t.Fatal(err)
262 }
263 if requests.Load() != 1 {
264 t.Fatalf("requests = %d, want a restore probe", requests.Load())
265 }
266 if r.Provider != "m4" || r.Model != "qwen3-coder" || r.BaseURL != server.URL+"/v1" {
267 t.Fatalf("resolved = provider=%s model=%s base=%s", r.Provider, r.Model, r.BaseURL)
268 }
269 if r.MaxOutput <= 0 {
270 t.Fatalf("discovered model lost output budget: %+v", r.MaxOutput)
271 }
272
273 // Once the model is known, launching again does not probe.
274 if _, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true); err != nil {
275 t.Fatal(err)
276 }
277 if requests.Load() != 1 {
278 t.Fatalf("requests = %d, want no extra probe", requests.Load())
279 }
280}