mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 09:47:23 +00:00
用户判据:「priority 直接移除」「即使带 key 也没有 /v1/responses」
「channel 插件的问题还没完全解决」「生图和聊天合并到一个导航项」
## 一、删除 priority(它从未参与跨源调度)
`config.Model.Priority` 号称"高者优先参与 AUTO 选源",**那是错的**:
AUTO 链由 `config.ModelScope.Tier` 驱动(scheduler.Chain),priority 只被
`bestChatModel/bestImageModel` 用于"本源内部该发哪个模型",与跨源调度无关。
把它当调度依据会让人以为拖个数字就能改顺序,改完毫无效果。
- 删 `config.Model.Priority` 字段与 `bestChatModel/bestImageModel` 的优先级比较
- 删 `seedAuto()`:**AUTO 链为空时网关应直接 503**,不能凭空播种用户没配过的档位
- 删 UI `buildLanes` 的 discovery 回退(按源 priority 自动铺开)与 sourceRows
- 清掉 config.yaml 里 140 行 priority
## 二、一通道 = 一个适配器 = 一个源(channels 插件)
旧设计「一账号一源」(`channels-<平台>-<12位hex>`) 有三个结构性后果:
1. 三个账号 ⇒ 三个源、三份同名模型,用户看到三行无法区分的东西
2. 谁先谁后**取决于注册顺序,不是余额**(rankAccounts 只用于展示,从未接进同步路径)
3. 换号要改源名 ⇒ 用户书签里的 `源:模型` 全部失效
新设计:
- 源名只跟通道绑定:`channels-<平台>`,换号 = 同名 add_source 覆盖 api_key
(内核 AddSource 是 UpsertSourceInYAML + 内存替换),对客户端完全透明
- `pickAccountForSource()`:凭据可用 → 余额分档(先快作废后永久)→
跳过已标记失效的号。没有 balance 声明的平台不据此排除(不因"不会查"判死)
- `request_end` 钩子:真实请求失败且是**终态**错误(需充值/封禁/凭据永久失效)
时给当前账号打 termErr,下次选号跳过;成功则清标记。
⚠️ 只认终态 —— 429/超时/抖动都**不**标记(那会把好号一个个废掉)
- `refresh_creds` 每 30 分钟巡检后重新选号(兜底心跳,纯本地判定)
- 迁移:同步时自动清掉旧版按账号拆的源(按已知平台前缀匹配,不用贪婪 match)
- 顺带修:`refresh_creds` 漏认 `refresh_fn` ⇒ codearts 的 DPoP 续期永不触发
## 三、discover 支持用已存凭据(真实缺口)
刷新已有源的模型目录,此前必须让运维**再粘一次明文 key** —— 而库里是
`enc:v1:` 密文,WebUI 也不回显明文。结果就是"想刷新 → 得重新申请 key",
很多人干脆不刷,源的目录越放越旧(实例:某源 3 个生图模型从未登记)。
- `{"name":"oin"}` ⇒ 从已存源取 base_url + 解密凭据 + adapter
- 安全边界:已要求 admin;只能探测**本网关已配置的**源(不能拿任意 URL+key
做探测,那就成了 SSRF/密钥外带);响应不回显 key
## 四、image_models 显式声明(补元数据缺失)
`discoveredModelKind` 只看上游 modality 字段,并刻意**不从模型名推断**
(错误的 kind 会毒死 AUTO 链)。但有些网关的 /v1/models 只返回
`{id,object,created,owned_by}` —— 零 modality,于是生图模型全落成 chat,
生图端点永远 `no image source configured`。
不解这个两难:默认仍完全按元数据(不猜),但允许 `image_models` **显式声明**。
写错的 id 明确报 400,不静默造一个不存在的模型。
## 五、UI:对话与生图合并为一个导航项
两者本是同一件事的两种输入(都选模型、都发请求、都读结果),只是 kind 与端点
不同。合并只发生在**导航层**:模型下拉仍按 kind 分组/禁选(把图片模型选进 chat
会请求到错误端点)。两个面板同时存在于 DOM,切换只改 display —— 来回切不会丢
已输入的文字、对话历史或已生成的图。
## 测试
- `TestUI_ImageModeIsReachableFromChatTab`:改为钉"合并后仍可达"(分段控件、
切换函数、面板容器、端点),并断言**旧 tab 三件套必须消失**(否则 goTab 会找
一个不存在的面板)
- 测试夹具显式建 AUTO 链(内核不再 seedAuto):默认一源一 tier 倒序
(TestChatAUTO 验优先级);TestChatAutoChainTierFailover 自己覆盖成同 tier
(验 tier 内轮转)—— 两者验不同机制,不能共用一个默认
- core 的一批断言从"按账号 id 找源"改为"按平台名找源"
gateway/config/provider 全绿;core 的 channels 用例全绿。
TestRealMutants 需工作树干净后才跑(其自带的变异残留防护)。
BREAKING CHANGE: config.Model.Priority 已删除。旧的 priority 字段会被 YAML
解码静默忽略(内核不校验未知字段),因此既有 config.yaml 仍可加载,但该字段
不再有任何作用 —— 调度顺序请通过 AUTO 链的 tier 表达。
(cherry picked from commit 25970a2135)
1189 lines
42 KiB
Go
1189 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})
|
|
}
|
|
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, "openai", "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"},
|
|
{ID: "img-m", Kind: "image"},
|
|
}
|
|
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) {
|
|
// A non-JSON HTML error page: the adapter's transform_error hook must
|
|
// distill it into a short reason rather than letting the raw markup
|
|
// reach the client or the log.
|
|
w.Header().Set("Content-Type", "text/html")
|
|
w.WriteHeader(503)
|
|
fmt.Fprint(w, `<html><head><title>503 Service Unavailable</title></head><body>maintenance</body></html>`)
|
|
}))
|
|
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 || !strings.Contains(err.Error(), "503 Service Unavailable") {
|
|
t.Fatalf("adapter hook reason must surface, got %v", err)
|
|
}
|
|
if strings.Contains(err.Error(), "<html") || strings.Contains(err.Error(), "<title>") {
|
|
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)
|
|
}
|
|
}
|