package provider import ( "context" "encoding/json" "errors" "fmt" "net/http" "net/http/httptest" "os" "path/filepath" "strings" "sync/atomic" "testing" "time" "llmsproxy/internal/config" "llmsproxy/internal/lua" "llmsproxy/internal/types" ) func newTestProvider(t *testing.T, src config.Source) *Provider { t.Helper() vm := lua.NewVM(filepath.Join(t.TempDir(), "adapters")) if err := vm.Start(); err != nil { t.Fatalf("vm: %v", err) } t.Cleanup(vm.Stop) return New(src, vm) } func src(name, url, adapter string, models ...string) config.Source { s := config.Source{Name: name, BaseURL: url, Adapter: adapter, MaxConcurrent: 4} for _, m := range models { s.Models = append(s.Models, config.Model{ID: m, Priority: 0}) } return s } func TestProviderChat(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var body map[string]interface{} _ = json.NewDecoder(r.Body).Decode(&body) if body["model"] != "mock-model" { t.Errorf("model = %v", body["model"]) } fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`) })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "mock-model")) resp, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "mock-model", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hello")}}, }) if err != nil { t.Fatalf("chat: %v", err) } if resp.Content != "hi" { t.Fatalf("content = %q", resp.Content) } } func TestProviderChatStream(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n") fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"b\"}}]}\n\n") fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n") fmt.Fprint(w, "data: [DONE]\n\n") })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) ch, err := p.ChatStream(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err != nil { t.Fatalf("stream: %v", err) } var text string var done bool for ck := range ch { text += ck.Content if ck.Done { done = true } } if text != "ab" { t.Fatalf("text = %q", text) } if !done { t.Fatal("expected done") } } // TestStreamSuccessClearsBackoff guards P6: a clean streaming end must reset // a previously cooled (source, model) pair. func TestStreamSuccessClearsBackoff(t *testing.T) { var fail atomic.Bool fail.Store(true) up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if fail.Load() { w.WriteHeader(500) return } w.Header().Set("Content-Type", "text/event-stream") fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"a\"}}]}\n\n") fmt.Fprint(w, "data: [DONE]\n\n") })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) if _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }); err == nil { t.Fatal("expected first chat to fail") } if p.ModelAvailable("m") { t.Fatal("m must be cooling after the failed chat") } fail.Store(false) ch, err := p.ChatStream(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err != nil { t.Fatalf("stream: %v", err) } for range ch { } if !p.ModelAvailable("m") { t.Fatal("clean stream must clear the cooldown") } if st := p.state("m"); st.FailCount() != 0 { t.Fatalf("fail count after clean stream = %d", st.FailCount()) } } func TestProviderImage(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { fmt.Fprint(w, `{"created":123,"data":[{"b64_json":"QUJD"}]}`) })) defer up.Close() p := newTestProvider(t, src("img", up.URL, "openai", "flux-1")) resp, err := p.Image(context.Background(), &types.ImageGenRequest{Model: "flux-1", Prompt: "cat"}) if err != nil { t.Fatalf("image: %v", err) } if len(resp.ImageData) != 1 || resp.ImageData[0].B64JSON != "QUJD" { t.Fatalf("image data = %+v", resp.ImageData) } } func TestProviderBackoffPerModel(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(500) fmt.Fprint(w, "boom") })) defer up.Close() // two models on one source: a failure on m1 must not blacklist m2 p := newTestProvider(t, src("mock", up.URL, "openai", "m1", "m2")) _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m1", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil { t.Fatal("expected error") } if p.ModelAvailable("m1") { t.Fatal("expected m1 to be cooling down") } if !p.ModelAvailable("m2") { t.Fatal("m2 must stay schedulable (per-model isolation)") } if !p.Available() { t.Fatal("source must stay available while any model is schedulable") } if st := p.state("m1"); st.FailCount() != 1 { t.Fatalf("fail count = %d, want 1", st.FailCount()) } } func TestProviderAuthFailureSelfHeals(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(401) fmt.Fprint(w, `{"error":"API_KEY_DISABLED"}`) })) defer up.Close() p := newTestProvider(t, src("mock2", up.URL, "openai", "m2")) _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m2", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil { t.Fatal("expected error") } st := p.state("m2") if st.FailCount() != backoffCapN { t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount()) } if until := st.CooldownUntil(); until <= time.Now().Add(authCooldown-time.Minute).Unix() { t.Fatalf("auth failure must cool for authCooldown (until=%d)", until) } if st.Pref() != -2*int64(prefFailStep) { t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref()) } // not permanent: the reset channel and a later success both restore it st.reset() if !p.ModelAvailable("m2") { t.Fatal("reset must restore schedulability") } } func TestModelStateCooldownAndRecovery(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(500) })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil { t.Fatal("expected error") } st := p.state("m") if st.FailCount() != 1 { t.Fatalf("fail count = %d", st.FailCount()) } // one failure -> 5s cooldown from now until := st.CooldownUntil() want := time.Now().Add(backoffBase).Unix() if until < want-2 || until > want+2 { t.Fatalf("cooldown = %d, want ~%d", until, want) } // success resets everything and bumps the preference st.RecordSuccess() if !p.ModelAvailable("m") { t.Fatal("success must clear cooldown") } if st.FailCount() != 0 { t.Fatalf("fail count after success = %d", st.FailCount()) } // failure penalties prefFailStep, success rewards prefSuccessStep: // -5 then +2 = -3 if st.Pref() != int64(prefSuccessStep)-int64(prefFailStep) { t.Fatalf("pref after one failure then success = %d, want %d", st.Pref(), int64(prefSuccessStep)-int64(prefFailStep)) } } func TestModelStateRateLimitShortCooldown(t *testing.T) { // 429 must NOT use the exponential backoff schedule: a rate-limited but // quota-rich source has to re-enter rotation after the short fixed window. p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m")) st := p.state("m") // repeated 429s must stay at the fixed short cooldown, never escalate for i := 0; i < 15; i++ { p.ReportStatus("m", 429) } until := st.CooldownUntil() want := time.Now().Add(rateLimitCooldown).Unix() if until < want-2 || until > want+2 { t.Fatalf("429 cooldown = %d, want ~%d (fixed %v, not exponential)", until, want, rateLimitCooldown) } if st.FailCount() != 15 { t.Fatalf("fail count = %d, want 15 (counted but not escalating)", st.FailCount()) } if st.Pref() != -15*int64(prefQuotaStep) { t.Fatalf("pref after 15 x 429 = %d, want %d (a 429 is a quota-class event, not a fault)", st.Pref(), -15*int64(prefQuotaStep)) } if p.ModelAvailable("m") { t.Fatal("model must be cooling right after a 429") } } func TestThrottleSpacingAndCancel(t *testing.T) { p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m")) p.cfg.RPM = 120 // gap = 500ms p.rpmGap = time.Minute / time.Duration(p.cfg.RPM) start := time.Now() for i := 0; i < 3; i++ { if err := p.Throttle(context.Background()); err != nil { t.Fatalf("throttle %d: %v", i, err) } } elapsed := time.Since(start) // first call passes immediately, the next two wait one gap each want := 2 * time.Minute / time.Duration(p.cfg.RPM) if elapsed < want-time.Duration(100*time.Millisecond) || elapsed > want+time.Second { t.Fatalf("3 throttled calls took %v, want ~%v", elapsed, want) } // unlimited source: Throttle is a no-op p2 := newTestProvider(t, src("free", "http://127.0.0.1:1", "openai", "m")) if err := p2.Throttle(context.Background()); err != nil { t.Fatalf("unlimited throttle: %v", err) } // cancelled context aborts a pending window reservation p3 := newTestProvider(t, src("slow", "http://127.0.0.1:1", "openai", "m")) p3.cfg.RPM = 6 // 10s gap p3.rpmGap = time.Minute / time.Duration(p3.cfg.RPM) // first call passes immediately (fresh provider) and reserves the next window if err := p3.Throttle(context.Background()); err != nil { t.Fatalf("first throttle: %v", err) } ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) defer cancel() // second call must wait ~10s for the reserved window; the short deadline aborts it start = time.Now() err := p3.Throttle(ctx) if !errors.Is(err, context.DeadlineExceeded) { t.Fatalf("cancelled throttle = %v, want DeadlineExceeded", err) } if time.Since(start) > time.Second { t.Fatalf("cancel took %v, want fast abort", time.Since(start)) } } func TestTryAcquire(t *testing.T) { p := newTestProvider(t, src("mock", "http://127.0.0.1:1", "openai", "m")) p.cfg.MaxConcurrent = 1 p.sem = make(chan struct{}, 1) if err := p.TryAcquire(context.Background()); err != nil { t.Fatalf("first acquire: %v", err) } if err := p.TryAcquire(context.Background()); !errors.Is(err, ErrBusy) { t.Fatalf("second acquire = %v, want ErrBusy", err) } p.Release() if err := p.TryAcquire(context.Background()); err != nil { t.Fatalf("acquire after release: %v", err) } p.Release() } func TestChatBusyFailsFast(t *testing.T) { release := make(chan struct{}) started := make(chan struct{}, 10) up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { started <- struct{}{} <-release fmt.Fprint(w, `{"choices":[{"message":{"content":"ok"}}]}`) })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) p.cfg.MaxConcurrent = 1 p.sem = make(chan struct{}, 1) done := make(chan error, 1) go func() { _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) done <- err }() <-started // first request holds the only slot // second request must fail fast with ErrBusy instead of queueing _, err2 := p.Chat(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if !errors.Is(err2, ErrBusy) { t.Fatalf("second chat err = %v, want ErrBusy", err2) } close(release) if err := <-done; err != nil { t.Fatalf("first chat: %v", err) } } func eventually(t *testing.T, timeout time.Duration, cond func() bool, msg string) { t.Helper() deadline := time.Now().Add(timeout) for time.Now().Before(deadline) { if cond() { return } time.Sleep(10 * time.Millisecond) } t.Fatalf("timed out: %s", msg) } // TestChatClientCancelNotRecorded: a client disconnect before the response is // neither success nor failure — the (source, model) state must stay clean. func TestChatClientCancelNotRecorded(t *testing.T) { started := make(chan struct{}) release := make(chan struct{}) up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { close(started) <-release // hold the upstream open; release at cleanup fmt.Fprint(w, `{"choices":[{"message":{"content":"late"}}]}`) })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) ctx, cancel := context.WithCancel(context.Background()) errCh := make(chan error, 1) go func() { _, err := p.Chat(ctx, &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) errCh <- err }() <-started cancel() // client disconnects mid-request if err := <-errCh; err == nil { t.Fatal("cancelled request must return an error") } close(release) eventually(t, 2*time.Second, func() bool { _, fc, _ := p.ModelHealthInfo("m") return fc == 0 }, "client cancel must not record a scheduling failure") } // TestChatStreamEmptyBodyNotSuccess: a 200 that yields zero chunks and no // [DONE] must fail the candidate (error return, no delivered chunks) so the // scheduler falls through to the next source, and record the failure. func TestChatStreamEmptyBodyNotSuccess(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") // 200 with an empty body: no chunks, no [DONE] })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) ch, err := p.ChatStream(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil { if ch != nil { for range ch { } } t.Fatal("empty stream must fail the candidate for scheduler fallback") } eventually(t, 2*time.Second, func() bool { _, fc, _ := p.ModelHealthInfo("m") return fc >= 1 }, "empty stream must record a failure") } // TestChatStreamErrorFinishFailsCandidate: zen free pool answers HTTP 200 // with a single-chunk stream whose only payload is finish_reason:"network_error" // and an empty delta. That must fail the candidate (not deliver a laundered // empty reply), so the scheduler can fall through to the next source. func TestChatStreamErrorFinishFailsCandidate(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"network_error","delta":{"role":"assistant","content":""}}]}`+"\n\n") fmt.Fprint(w, "data: [DONE]\n\n") })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "opencode", "m")) _, err := p.ChatStream(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil { t.Fatal("network_error-only stream must fail the candidate") } eventually(t, 2*time.Second, func() bool { _, fc, _ := p.ModelHealthInfo("m") return fc >= 1 }, "error-only stream must record a failure") } // TestChatStreamInstantEmptyStopStillDelivered: a legitimate completion that // ends immediately with the standard finish_reason:"stop" and zero content is // NOT an upstream error and must still reach the client. func TestChatStreamInstantEmptyStopStillDelivered(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") fmt.Fprint(w, `data: {"choices":[{"index":0,"finish_reason":"stop","delta":{}}]}`+"\n\n") fmt.Fprint(w, "data: [DONE]\n\n") })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "openai", "m")) ch, err := p.ChatStream(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err != nil { t.Fatalf("standard stop stream must be delivered: %v", err) } n := 0 for range ch { n++ } if n == 0 { t.Fatal("expected at least the terminating chunk") } eventually(t, 2*time.Second, func() bool { _, fc, _ := p.ModelHealthInfo("m") return fc == 0 }, "standard stop must count as success") } // TestImageAutoUsesImageModel: AUTO image generation on a mixed source must // send the image-kind model id, never the best chat model. func TestImageAutoUsesImageModel(t *testing.T) { var gotModel string up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { var body map[string]interface{} _ = json.NewDecoder(r.Body).Decode(&body) gotModel, _ = body["model"].(string) fmt.Fprint(w, `{"created":1,"data":[{"url":"http://x/1.png"}]}`) })) defer up.Close() s := config.Source{Name: "mix", BaseURL: up.URL, Adapter: "openai", MaxConcurrent: 4} s.Models = []config.Model{ {ID: "chat-m", Kind: "chat", Priority: 100}, {ID: "img-m", Kind: "image", Priority: 50}, } p := newTestProvider(t, s) resp, err := p.Image(context.Background(), &types.ImageGenRequest{Prompt: "cat", Model: "AUTO"}) if err != nil { t.Fatalf("image: %v", err) } if len(resp.ImageData) == 0 { t.Fatal("no image data returned") } if gotModel != "img-m" { t.Fatalf("AUTO image must use the image-kind model, got %q", gotModel) } } func TestStandardSSEChunkFinishReason(t *testing.T) { out := standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"tool_calls","delta":{}}]}`) if !strings.Contains(out, `"done":true`) || !strings.Contains(out, `"finish_reason":"tool_calls"`) { t.Fatalf("real finish reason must pass through: %s", out) } out = standardSSEChunk(`{"choices":[{"index":0,"finish_reason":"","delta":{"content":"hi"}}]}`) if strings.Contains(out, `"Done":true`) { t.Fatalf("empty-string finish reason is not a finish signal: %s", out) } } // TestAdapterHookCondensesError: the adapter transform_error hook owns the // per-source error format; its reason must reach the client verbatim. func TestAdapterHookCondensesError(t *testing.T) { up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") w.WriteHeader(503) fmt.Fprint(w, `{"error":{"type":"FreeUsageLimitError","message":"Rate limit exceeded. Please try again later."}}`) })) defer up.Close() p := newTestProvider(t, src("mock", up.URL, "opencode", "m")) _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil || !strings.Contains(err.Error(), "zen free pool quota exhausted") { t.Fatalf("adapter hook reason must surface, got %v", err) } if strings.Contains(err.Error(), "Rate limit exceeded") { t.Fatalf("raw upstream body must not leak past the hook: %v", err) } } // TestUnknownErrorFallbackWithoutHook: an adapter without transform_error // gets the uniform core fallback; the raw body must not leak to clients. func TestUnknownErrorFallbackWithoutHook(t *testing.T) { dir := filepath.Join(t.TempDir(), "adapters") if err := os.MkdirAll(dir, 0o755); err != nil { t.Fatal(err) } minimal := `return { name="custom", endpoint="/chat/completions", transform_request=function(raw) return raw end, transform_response=function(raw) return raw end } ` if err := os.WriteFile(filepath.Join(dir, "custom.lua"), []byte(minimal), 0o644); err != nil { t.Fatal(err) } vm := lua.NewVM(dir) if err := vm.Start(); err != nil { t.Fatalf("vm: %v", err) } defer vm.Stop() up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(500) fmt.Fprint(w, `{"weird":{"shape":["no","message"]}}`) })) defer up.Close() p := New(src("mock", up.URL, "custom", "m"), vm) _, err := p.Chat(context.Background(), &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, }) if err == nil || !strings.Contains(err.Error(), "unknown error") { t.Fatalf("hook-less adapter must fall back to unknown error, got %v", err) } if strings.Contains(err.Error(), "shape") { t.Fatalf("raw body must not leak: %v", err) } } // TestToolIndexRemapper verifies sparse upstream tool-call indices are // compacted to dense 0-based OpenAI ordinals (Issue 2). Anthropic emits the // content_block ordinal, so a thinking block at index 0 pushes the first // tool_use to index 1 — OpenAI clients accumulating by index would leave a // hole at 0 and mis-assemble multi-tool responses. func TestToolIndexRemapper(t *testing.T) { t.Run("single tool after thinking block", func(t *testing.T) { r := newToolIndexRemapper() // content_block index 1 (thinking was 0) -> dense 0 out := r.remap(json.RawMessage(`[{"index":1,"id":"toolu_a","type":"function","function":{"name":"calc","arguments":""}}]`)) var got []map[string]interface{} if err := json.Unmarshal(out, &got); err != nil { t.Fatalf("unmarshal: %v (%s)", err, out) } if got[0]["index"] != float64(0) { t.Fatalf("index = %v, want 0: %s", got[0]["index"], out) } // id/name/arguments must survive verbatim if got[0]["id"] != "toolu_a" { t.Fatalf("id lost: %s", out) } fn := got[0]["function"].(map[string]interface{}) if fn["name"] != "calc" { t.Fatalf("function.name lost: %s", out) } // subsequent argument fragments on the SAME upstream index reuse ordinal 0 out2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"{\"a\":1}"}}]`)) var got2 []map[string]interface{} if err := json.Unmarshal(out2, &got2); err != nil { t.Fatalf("unmarshal2: %v", err) } if got2[0]["index"] != float64(0) { t.Fatalf("fragment index = %v, want stable 0: %s", got2[0]["index"], out2) } }) t.Run("multiple tools get distinct dense ordinals", func(t *testing.T) { r := newToolIndexRemapper() // thinking=0, tool_use=1, tool_use=2 -> 0, 1 a := r.remap(json.RawMessage(`[{"index":1,"id":"t1","type":"function","function":{"name":"f1","arguments":""}}]`)) b := r.remap(json.RawMessage(`[{"index":2,"id":"t2","type":"function","function":{"name":"f2","arguments":""}}]`)) var ga, gb []map[string]interface{} json.Unmarshal(a, &ga) json.Unmarshal(b, &gb) if ga[0]["index"] != float64(0) { t.Fatalf("first tool index = %v, want 0", ga[0]["index"]) } if gb[0]["index"] != float64(1) { t.Fatalf("second tool index = %v, want 1", gb[0]["index"]) } // interleaved fragments keep their own ordinals a2 := r.remap(json.RawMessage(`[{"index":1,"id":"","type":"function","function":{"name":"","arguments":"x"}}]`)) var ga2 []map[string]interface{} json.Unmarshal(a2, &ga2) if ga2[0]["index"] != float64(0) { t.Fatalf("t1 fragment index = %v, want 0", ga2[0]["index"]) } }) t.Run("already dense indices pass through untouched", func(t *testing.T) { r := newToolIndexRemapper() in := json.RawMessage(`[{"index":0,"id":"t","type":"function","function":{"name":"f","arguments":""}}]`) out := r.remap(in) if string(out) != string(in) { t.Fatalf("dense input was rewritten:\n in=%s\nout=%s", in, out) } }) t.Run("empty and malformed input is safe", func(t *testing.T) { r := newToolIndexRemapper() if got := r.remap(nil); got != nil { t.Fatalf("nil -> %s", got) } bad := json.RawMessage(`not json`) if got := r.remap(bad); string(got) != string(bad) { t.Fatalf("malformed input must pass through: %s", got) } }) } // ---- half-cooldown probe gate (plan 阶段 1) ---- // TestProbeGateHalfCooldown pins the core rule: a cooling model is fully silent // in the first half of its cooldown window and probeable exactly once in the // second half. func TestProbeGateHalfCooldown(t *testing.T) { var st ModelState // open a 100s window centred so that "now" sits in the first half now := time.Now().Unix() st.cooldownFrom.Store(now - 10) st.cooldownUntil.Store(now + 90) if st.TryProbe() { t.Fatal("first half of the cooldown window must stay silent") } // move into the second half st.cooldownFrom.Store(now - 90) st.cooldownUntil.Store(now + 10) if !st.TryProbe() { t.Fatal("second half of the cooldown window must allow one probe") } if st.TryProbe() { t.Fatal("probe permit must be single-token") } if !st.Probing() { t.Fatal("Probing() must report the claimed permit") } st.ProbeDone() if st.Probing() { t.Fatal("ProbeDone must release the permit") } if !st.TryProbe() { t.Fatal("permit must be reusable after ProbeDone") } } // TestProbeGateNoProbeWhenHealthy: a healthy model needs no probe permit — the // caller schedules it normally. func TestProbeGateNoProbeWhenHealthy(t *testing.T) { var idle ModelState if idle.TryProbe() { t.Fatal("a healthy model must not be probed (schedule it normally)") } } // TestProbeGateRescuesPrefFloor documents a fix that comes with the probe gate: // a slot whose preference sank to prefMin used to be refused by ModelAvailable // forever, even long after every cooldown expired — a permanent blacklist. It is // now probeable, so one success brings it back. func TestProbeGateRescuesPrefFloor(t *testing.T) { var dead ModelState dead.pref.Store(prefMin) if !dead.TryProbe() { t.Fatal("a slot stuck at prefMin must be probeable, not blacklisted") } dead.RecordSuccess() dead.ProbeDone() if dead.Pref() <= prefMin { t.Fatalf("one probe success must lift the slot off the floor, pref=%d", dead.Pref()) } } // TestProbeSuccessRestoresImmediately is the user-visible payoff: a recovered // upstream must not sit out the rest of its cooldown. func TestProbeSuccessRestoresImmediately(t *testing.T) { var st ModelState for i := 0; i < backoffCapN; i++ { st.RecordFailure(false) } if cd := st.CooldownUntil() - time.Now().Unix(); cd < int64(backoffCap.Seconds())-2 { t.Fatalf("10 failures must cool for backoffCap, got %ds", cd) } // jump into the probe window: the midpoint of [from, until) must be in the // past, so push the window's start back past its full length now := time.Now().Unix() st.cooldownFrom.Store(now - int64(backoffCap.Seconds()) - 10) if !st.TryProbe() { t.Fatal("probe must be allowed past the midpoint") } if !st.Probing() { t.Fatal("probe permit must be held while the probe runs") } st.RecordSuccess() st.ProbeDone() if !st.Available() { t.Fatal("probe success must clear the cooldown immediately") } if st.FailCount() != 0 { t.Fatalf("probe success must reset failCount, got %d", st.FailCount()) } if st.CooldownUntil() != 0 || st.CooldownFrom() != 0 { t.Fatal("probe success must clear both cooldown bounds") } } // TestProbeFailureDefersNextProbe: a failed probe must not immediately retry — // the new window's midpoint pushes the next probe out. func TestProbeFailureDefersNextProbe(t *testing.T) { var st ModelState st.RecordFailure(false) // 5s window now := time.Now().Unix() st.cooldownFrom.Store(now - 100) st.cooldownUntil.Store(now + 10) if !st.TryProbe() { t.Fatal("probe must be allowed past the midpoint") } st.RecordFailure(false) // probe failed: a fresh, longer window opens st.ProbeDone() if st.TryProbe() { t.Fatal("a failed probe must defer the next probe to the new midpoint") } if from := st.CooldownFrom(); from < now { t.Fatalf("failed probe must reopen the window (from=%d, now=%d)", from, now) } } // TestAuthCooldownIsSeparateFromCap documents that credential failures cool on // their own (longer) schedule instead of reusing the exponential cap. func TestAuthCooldownIsSeparateFromCap(t *testing.T) { var st ModelState st.RecordFailure(true) cd := st.CooldownUntil() - time.Now().Unix() if cd < int64(authCooldown.Seconds())-2 || cd > int64(authCooldown.Seconds())+2 { t.Fatalf("auth failure must cool for authCooldown (%v), got %ds", authCooldown, cd) } if st.FailCount() != backoffCapN { t.Fatalf("auth failure must persist the capped count, got %d", st.FailCount()) } } // TestQuotaExhaustedDoesNotEscalate is the second half of the user's complaint: // running out of allowance must not be treated as a failure streak. func TestQuotaExhaustedDoesNotEscalate(t *testing.T) { var st ModelState st.RecordQuotaExhausted(2 * time.Minute) if st.FailCount() != 0 { t.Fatalf("quota exhaustion must not bump failCount, got %d", st.FailCount()) } cd := st.CooldownUntil() - time.Now().Unix() if cd < 118 || cd > 122 { t.Fatalf("quota cooldown must honour the reset window, got %ds", cd) } // an unbounded window is clamped so a topped-up allowance is still noticed var big ModelState big.RecordQuotaExhausted(72 * time.Hour) if cd := big.CooldownUntil() - time.Now().Unix(); cd > int64(maxQuotaCooldown.Seconds())+2 { t.Fatalf("quota cooldown must be capped at maxQuotaCooldown, got %ds", cd) } // no hint at all falls back to the short window var none ModelState none.RecordQuotaExhausted(0) if cd := none.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 { t.Fatalf("unknown reset window must fall back to rateLimitCooldown, got %ds", cd) } } // TestQuotaReasonRouting checks that a 429 carrying a quota phrase is cooled as // quota exhaustion while a plain rate limit keeps the short fixed window. func TestQuotaReasonRouting(t *testing.T) { p := newTestProvider(t, src("q", "http://127.0.0.1:1", "openai", "m")) p.ReportStatusReason("m", 429, "429 insufficient_quota: you exceeded your current quota") st := p.state("m") if st.FailCount() != 0 { t.Fatalf("quota rejection must not bump failCount, got %d", st.FailCount()) } if cd := st.CooldownUntil() - time.Now().Unix(); cd <= int64(rateLimitCooldown.Seconds()) { t.Fatalf("quota rejection must cool longer than a plain 429, got %ds", cd) } p2 := newTestProvider(t, src("r", "http://127.0.0.1:1", "openai", "m")) p2.ReportStatusReason("m", 429, "rate limit reached, slow down") st2 := p2.state("m") if cd := st2.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 { t.Fatalf("plain 429 must keep the short window, got %ds", cd) } } // TestModelSchedulableProbeHandoff covers the provider-level gate the scheduler // consumes. func TestModelSchedulableProbeHandoff(t *testing.T) { p := newTestProvider(t, src("s", "http://127.0.0.1:1", "openai", "m")) if ok, probe := p.ModelSchedulable("m"); !ok || probe { t.Fatalf("healthy model must be normally schedulable, got ok=%v probe=%v", ok, probe) } st := p.state("m") st.RecordFailure(false) now := time.Now().Unix() st.cooldownFrom.Store(now - 100) st.cooldownUntil.Store(now + 10) ok, probe := p.ModelSchedulable("m") if !ok || !probe { t.Fatalf("cooling model past midpoint must be probeable, got ok=%v probe=%v", ok, probe) } if ok2, _ := p.ModelSchedulable("m"); ok2 { t.Fatal("only one probe permit may be outstanding") } p.ProbeDone("m") if ok3, probe3 := p.ModelSchedulable("m"); !ok3 || !probe3 { t.Fatal("permit must be reclaimable after ProbeDone") } } // TestProbeSlotsFollowsMaxConcurrent pins the user's example: max_concurrent=10 // yields exactly one probe slot. func TestProbeSlotsFollowsMaxConcurrent(t *testing.T) { for _, tc := range []struct{ mc, want int }{ {0, 1}, {1, 1}, {4, 1}, {10, 1}, {20, 2}, {100, 2}, } { s := src("p", "http://127.0.0.1:1", "openai", "m") s.MaxConcurrent = tc.mc p := newTestProvider(t, s) if got := p.ProbeSlots(); got != tc.want { t.Fatalf("max_concurrent=%d: probe slots = %d, want %d", tc.mc, got, tc.want) } } } // TestModelProbeInfo checks the UI-facing view of the probe window. func TestModelProbeInfo(t *testing.T) { p := newTestProvider(t, src("i", "http://127.0.0.1:1", "openai", "m")) if from, after, probing := p.ModelProbeInfo("m"); from != 0 || after != 0 || probing { t.Fatalf("idle model must report no probe window, got %d %d %v", from, after, probing) } st := p.state("m") now := time.Now().Unix() st.cooldownFrom.Store(now) st.cooldownUntil.Store(now + 100) from, after, _ := p.ModelProbeInfo("m") if from != now || after != now+50 { t.Fatalf("probe window = (%d,%d), want (%d,%d)", from, after, now, now+50) } } // TestUpstreamRecoveryWithoutWaitingOutCooldown is the end-to-end shape of the // user's complaint: an upstream that 500s ten times lands in the capped // cooldown, and used to be unusable for a full 30 minutes even after it // recovered. With the probe gate the slot is silent for the first half of the // window and then serves the very next request itself, clearing the cooldown. // // Real time is not waited out: the window bounds are moved so the test observes // the same state the scheduler would see at cap/2. func TestUpstreamRecoveryWithoutWaitingOutCooldown(t *testing.T) { var healthy atomic.Bool var hits atomic.Int64 up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hits.Add(1) if !healthy.Load() { w.WriteHeader(500) return } fmt.Fprint(w, `{"choices":[{"message":{"content":"back"},"finish_reason":"stop"}]}`) })) defer up.Close() p := newTestProvider(t, src("recov", up.URL, "openai", "m")) req := func() *types.ChatRequest { return &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, } } // 1. drive the slot into the capped cooldown for i := 0; i < backoffCapN; i++ { if _, err := p.Chat(context.Background(), req()); err == nil { t.Fatalf("attempt %d: expected upstream failure", i) } } st := p.state("m") if p.ModelAvailable("m") { t.Fatal("10 consecutive 5xx must take the slot out of normal rotation") } cd := st.CooldownUntil() - time.Now().Unix() if cd < int64(backoffCap.Seconds())-2 { t.Fatalf("cooldown = %ds, want ~%v", cd, backoffCap) } // the old 30-minute lockout is gone if cd > int64((6 * time.Minute).Seconds()) { t.Fatalf("cooldown %ds is too long: the cap must be minutes, not half an hour", cd) } // 2. upstream recovers, but we are still in the FIRST half of the window: // the slot must stay completely silent (no probe traffic at all). healthy.Store(true) quiet := hits.Load() if ok, probe := p.ModelSchedulable("m"); ok || probe { t.Fatalf("first half of the window must stay silent, got ok=%v probe=%v", ok, probe) } if hits.Load() != quiet { t.Fatal("no request may reach the upstream during the silent half") } // 3. advance past the midpoint: exactly one probe is admitted. st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1) ok, isProbe := p.ModelSchedulable("m") if !ok || !isProbe { t.Fatalf("past the midpoint one probe must be admitted, got ok=%v probe=%v", ok, isProbe) } if ok2, _ := p.ModelSchedulable("m"); ok2 { t.Fatal("only ONE probe may be in flight while cooling") } // 4. the probe is a real request: its success restores the slot fully. resp, err := p.Chat(context.Background(), req()) p.ProbeDone("m") if err != nil { t.Fatalf("probe request failed: %v", err) } if resp.Content != "back" { t.Fatalf("probe response = %q, want \"back\"", resp.Content) } if !p.ModelAvailable("m") { t.Fatal("a successful probe must return the slot to normal rotation immediately") } if st.CooldownUntil() != 0 || st.FailCount() != 0 { t.Fatalf("probe success must clear cooldown/failCount, got until=%d fails=%d", st.CooldownUntil(), st.FailCount()) } if ok, probe := p.ModelSchedulable("m"); !ok || probe { t.Fatalf("recovered slot must schedule normally, got ok=%v probe=%v", ok, probe) } } // TestRecoveryProbeBudgetIsBounded quantifies the cost of the shorter cap // against a still-dead upstream: one probe per window, not a retry storm. func TestRecoveryProbeBudgetIsBounded(t *testing.T) { var hits atomic.Int64 up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { hits.Add(1) w.WriteHeader(500) })) defer up.Close() p := newTestProvider(t, src("dead", up.URL, "openai", "m")) req := &types.ChatRequest{ Model: "m", Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}}, } for i := 0; i < backoffCapN; i++ { _, _ = p.Chat(context.Background(), req) } baseline := hits.Load() st := p.state("m") // simulate three consecutive probe windows against a dead upstream for i := 0; i < 3; i++ { st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1) ok, isProbe := p.ModelSchedulable("m") if !ok || !isProbe { t.Fatalf("window %d: expected a probe permit", i) } // a second caller in the same window gets nothing if ok2, _ := p.ModelSchedulable("m"); ok2 { t.Fatalf("window %d: more than one probe admitted", i) } if _, err := p.Chat(context.Background(), req); err == nil { t.Fatalf("window %d: expected the probe to fail", i) } p.ProbeDone("m") } if extra := hits.Load() - baseline; extra != 3 { t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra) } } // TestQuotaPenaltyIsLighterThanFailure is the regression for the report // "今天一次失败没看到,分数目前已经 -10 了" on an allowance-metered source // (sensenova). RecordQuotaExhausted deliberately does not touch failCount — the // WebUI therefore showed fails=0 and cooling=false — yet it used to charge the // SAME score penalty as a real failure, so two ordinary quota resets dragged the // slot to -10 with no visible fault. An exhausted allowance means "no budget // right now", not "this model is broken": the cooldown already keeps it out of // rotation, so the score penalty must be mild. func TestQuotaPenaltyIsLighterThanFailure(t *testing.T) { quota := &ModelState{} quota.RecordQuotaExhausted(time.Minute) fail := &ModelState{} fail.RecordFailure(false) qp, fp := quota.Pref(), fail.Pref() if qp <= fp { t.Fatalf("a quota reset must cost less score than a real failure: quota=%d failure=%d", qp, fp) } if qp != -prefQuotaStep { t.Errorf("quota penalty = %d, want -%d", qp, prefQuotaStep) } if fp != -prefFailStep { t.Errorf("failure penalty = %d, want -%d", fp, prefFailStep) } // The reported symptom: several quota resets in a day must NOT approach the // prefMin floor, since each one is a normal event for a metered source. s := &ModelState{} for i := 0; i < 5; i++ { s.RecordQuotaExhausted(time.Minute) } if got := s.Pref(); got <= prefMin/2 { t.Errorf("5 quota resets sank the score to %d (floor %d); a metered source would be written off", got, prefMin) } if s.FailCount() != 0 { t.Errorf("quota exhaustion must not inflate failCount, got %d", s.FailCount()) } } // TestRateLimitPenaltyIsMild covers the sibling case: a 429 means "too fast", // not "broken", so a popular-but-healthy slot must not sink for being throttled. func TestRateLimitPenaltyIsMild(t *testing.T) { s := &ModelState{} s.RecordRateLimit() if got := s.Pref(); got != -prefQuotaStep { t.Errorf("429 penalty = %d, want -%d (same class as a quota event)", got, prefQuotaStep) } } // TestSuccessRecoversFasterThanBefore pins the penalty/reward asymmetry. Reward // used to be +1 against a -5 penalty, so a slot at -10 needed TEN consecutive // successes just to reach neutral — in practice it never got there, because a // low score makes the slot unattractive to the scheduler in the first place. func TestSuccessRecoversFasterThanBefore(t *testing.T) { s := &ModelState{} s.RecordFailure(false) s.RecordFailure(false) // -10, the score from the bug report n := 0 for s.Pref() < 0 { s.RecordSuccess() n++ if n > 50 { t.Fatal("score never recovered") } } // with +2 per success, -10 needs 5; the old +1 needed 10 if n > 5 { t.Errorf("recovery from -10 took %d successes, want <= 5", n) } // a real failure must still outweigh a single success, or failures stop // meaning anything if prefSuccessStep >= prefFailStep { t.Errorf("success reward %d must stay below the failure penalty %d", prefSuccessStep, prefFailStep) } } // TestPrefDecayRehabilitatesIdleSlot closes the starvation loop: a penalised // slot is unattractive, so the scheduler stops picking it, so it never earns the // successes that would rehabilitate it. Without decay the score stays negative // forever even though the upstream may be perfectly healthy. func TestPrefDecayRehabilitatesIdleSlot(t *testing.T) { s := &ModelState{} s.RecordFailure(false) s.RecordFailure(false) if s.Pref() != -2*prefFailStep { t.Fatalf("setup: pref=%d", s.Pref()) } // backdate the last touch by three decay intervals step := int64(prefDecayAfter / time.Second) s.prefTouched.Store(time.Now().Unix() - 3*step) got := s.Pref() if got != -2*prefFailStep+3 { t.Errorf("after 3 idle intervals pref=%d, want %d", got, -2*prefFailStep+3) } // decay must never overshoot into positive territory s.prefTouched.Store(time.Now().Unix() - 1000*step) if got := s.Pref(); got != 0 { t.Errorf("decay overshot to %d, want to stop at 0", got) } // a positive score must not be dragged down for being idle good := &ModelState{} good.RecordSuccess() good.RecordSuccess() before := good.Pref() good.prefTouched.Store(time.Now().Unix() - 1000*step) if after := good.Pref(); after != before { t.Errorf("idle decay must not touch a positive score: %d -> %d", before, after) } } // TestDecayLiftsSlotOffTheFloor verifies decay actually restores schedulability: // a slot at prefMin is refused by ModelAvailable, and before decay existed only // a probe could rescue it. func TestDecayLiftsSlotOffTheFloor(t *testing.T) { s := &ModelState{} for i := 0; i < 10; i++ { s.RecordFailure(false) } if s.Pref() != prefMin { t.Fatalf("setup: want floor %d, got %d", prefMin, s.Pref()) } // clear the cooldown so only the score gates the slot s.cooldownUntil.Store(0) s.cooldownFrom.Store(0) step := int64(prefDecayAfter / time.Second) s.prefTouched.Store(time.Now().Unix() - 2*step) if got := s.Pref(); got <= prefMin { t.Errorf("decay failed to lift the slot off the floor: %d", got) } }