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 TestRefreshCustomProviderModelsBackgroundSkipsKeyCommand(t *testing.T) {
140 home := t.TempDir()
141 t.Setenv("ZOT_HOME", home)
142 t.Setenv("M4_API_KEY", "")
143 var requests atomic.Int32
144 server := newModelsServer(t, &requests, "", "m")
145 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server.URL))
146 creds := auth.Credentials{AdditionalAPIKeyCreds: map[string]auth.ProviderCreds{
147 "m4": {APIKeyCommand: &auth.APIKeyCommand{Program: filepath.Join(home, "does-not-exist")}},
148 }}
149 encoded, err := json.Marshal(creds)
150 if err != nil {
151 t.Fatal(err)
152 }
153 if err := os.WriteFile(AuthPath(), encoded, 0o600); err != nil {
154 t.Fatal(err)
155 }
156
157 if err := refreshCustomProviderModels(context.Background(), apiKeyCommandSkip); err != nil {
158 t.Fatal(err)
159 }
160 if requests.Load() != 0 {
161 t.Fatal("background refresh probed a provider whose key needs a command")
162 }
163 // Interactive refresh runs the command and reports its failure
164 // instead of silently probing without a key.
165 err = RefreshCustomProviderModels(context.Background())
166 if err == nil || !strings.Contains(err.Error(), "m4") {
167 t.Fatalf("command failure not reported: %v", err)
168 }
169 if requests.Load() != 0 {
170 t.Fatal("refresh probed after key command failed")
171 }
172}
173
174func TestRefreshCustomProviderModelsKeepsSnapshotOnFailure(t *testing.T) {
175 t.Setenv("ZOT_HOME", t.TempDir())
176 t.Setenv("M4_API_KEY", "")
177 t.Setenv("M5_API_KEY", "")
178 var requests atomic.Int32
179 good := newModelsServer(t, &requests, "", "alive")
180 failing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
181 http.Error(w, "nope", http.StatusServiceUnavailable)
182 }))
183 defer failing.Close()
184 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {
185 "m4": {"baseUrl": %q, "discover": true},
186 "m5": {"baseUrl": %q, "discover": true}
187 }}`, good.URL, failing.URL))
188 provider.SetManagedModelsForProvider("m5", []provider.Model{{Provider: "m5", ID: "previous", Source: "live"}})
189
190 err := RefreshCustomProviderModels(context.Background())
191 if err == nil || !strings.Contains(err.Error(), "m5 discovery HTTP 503") {
192 t.Fatalf("err = %v", err)
193 }
194 if _, err := provider.FindModel("m4", "alive"); err != nil {
195 t.Fatal("healthy provider not refreshed when a sibling failed")
196 }
197 if _, err := provider.FindModel("m5", "previous"); err != nil {
198 t.Fatal("failed refresh dropped the previous snapshot")
199 }
200}
201
202func TestValidateAndRepairConfig_PreservesDiscoverableCustomModel(t *testing.T) {
203 home := t.TempDir()
204 t.Setenv("ZOT_HOME", home)
205 loadTestModelsJSON(t, `{"providers": {
206 "m4": {"baseUrl": "http://127.0.0.1:1/v1", "discover": true},
207 "m5": {"baseUrl": "http://127.0.0.1:1/v1", "models": [{"id": "pinned"}]}
208 }}`)
209
210 write := func(c Config) {
211 t.Helper()
212 b, _ := json.Marshal(c)
213 if err := os.WriteFile(filepath.Join(home, "config.json"), b, 0o644); err != nil {
214 t.Fatal(err)
215 }
216 }
217
218 // Discovery has not run yet at repair time, so the model is unknown.
219 write(Config{Provider: "m4", Model: "qwen3-coder"})
220 ValidateAndRepairConfig()
221 out, _ := LoadConfig()
222 if out.Provider != "m4" || out.Model != "qwen3-coder" {
223 t.Fatalf("discoverable custom model repaired away: %+v", out)
224 }
225
226 // A provider without discovery keeps the existing repair behavior.
227 write(Config{Provider: "m5", Model: "gone"})
228 ValidateAndRepairConfig()
229 out, _ = LoadConfig()
230 if out.Provider != "m5" || out.Model != "pinned" {
231 t.Fatalf("non-discoverable custom model not repaired: %+v", out)
232 }
233}
234
235func TestResolveRestoresDiscoveredCustomModelAfterRestart(t *testing.T) {
236 t.Setenv("ZOT_HOME", t.TempDir())
237 t.Setenv("M4_API_KEY", "")
238 var requests atomic.Int32
239 server := newModelsServer(t, &requests, "", "qwen3-coder")
240 loadTestModelsJSON(t, fmt.Sprintf(`{"providers": {"m4": {"baseUrl": %q, "discover": true}}}`, server.URL+"/v1"))
241
242 r, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true)
243 if err != nil {
244 t.Fatal(err)
245 }
246 if requests.Load() != 1 {
247 t.Fatalf("requests = %d, want a restore probe", requests.Load())
248 }
249 if r.Provider != "m4" || r.Model != "qwen3-coder" || r.BaseURL != server.URL+"/v1" {
250 t.Fatalf("resolved = provider=%s model=%s base=%s", r.Provider, r.Model, r.BaseURL)
251 }
252 if r.MaxOutput <= 0 {
253 t.Fatalf("discovered model lost output budget: %+v", r.MaxOutput)
254 }
255
256 // Once the model is known, launching again does not probe.
257 if _, err := Resolve(Args{Provider: "m4", Model: "qwen3-coder", NoSkill: true}, true); err != nil {
258 t.Fatal(err)
259 }
260 if requests.Load() != 1 {
261 t.Fatalf("requests = %d, want no extra probe", requests.Load())
262 }
263}