mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-19 16:39:15 +00:00
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.
1186 lines
42 KiB
Go
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)
|
|
}
|
|
}
|