Files
ModelRouter/internal/provider/provider_test.go
JianFeeeee 5b628b4d6f fix(scheduler): rebalance the pref score so metered sources aren't written off silently
Report: an allowance-metered source (sensenova) had its deepseek model sitting at
pref -10 while the WebUI showed zero failures. All three numbers were accurate,
and they exposed three compounding problems.

1. Quota exhaustion cost the SAME score as a real failure. RecordQuotaExhausted
   intentionally avoids failCount (an exhausted allowance is not a fault), so the
   UI showed fails=0 / cooling=false — yet it deducted the full prefFailStep (5).
   For a metered source, running out of budget is an everyday event, so the score
   drifted deep negative with no visible cause. Now quota and 429 events cost
   prefQuotaStep (1): the cooldown already keeps the slot out of rotation until
   the window resets, the score only needs a mild preference for slots with budget.

2. Recovery was 5:1 asymmetric. A failure cost -5 but a success only +1, so a
   slot at -10 needed ten consecutive successes just to reach neutral — which it
   could never get, because a low score makes the scheduler not pick it in the
   first place (starvation). Success now rewards prefSuccessStep (2): recovery
   from -10 needs five successes, while a real failure still outweighs one.

3. No idle decay. A penalised slot kept its negative score forever once it stopped
   being selected. Pref() now applies lazy decay: after prefDecayAfter (2 min) of
   no outcome, the score drifts one step back toward 0 per interval (never past
   0, never touches positive scores). Applied in Pref() and TryProbe(), so a
   naturally-recovered idle slot is schedulable again without needing a probe.

The failure penalty itself is unchanged (prefFailStep=5), so genuinely broken
upstreams are still marked as clearly worse than healthy ones.

Tests: quota penalty lighter than failure; 5 quota resets stay well above the
floor with failCount untouched; 429 is quota-class; recovery from -10 needs <=5;
idle decay rehabilitates a written-off slot, stops at 0, and never drags a
positive score; decay repeatedly lifts a slot off the prefMin floor. Updated the
two pre-existing tests that asserted the old -5/129 +1 values.
2026-08-31 11:50:47 +08:00

1186 lines
42 KiB
Go

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)
}
}