mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-11-07 09:47:23 +00:00
25970a2135
用户判据：「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 表达。
1066 lines
39 KiB
Go
1066 lines
39 KiB
Go
package gateway
|
||
|
||
import (
|
||
"encoding/json"
|
||
"fmt"
|
||
"io"
|
||
"net/http"
|
||
"net/http/httptest"
|
||
"net/url"
|
||
"os"
|
||
"path/filepath"
|
||
"strconv"
|
||
"strings"
|
||
"testing"
|
||
"time"
|
||
|
||
"llmsproxy/internal/config"
|
||
"llmsproxy/internal/core"
|
||
)
|
||
|
||
func mockUpstream() *httptest.Server {
|
||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
body, _ := io.ReadAll(r.Body)
|
||
var req map[string]interface{}
|
||
_ = json.Unmarshal(body, &req)
|
||
if stream, _ := req["stream"].(bool); stream {
|
||
w.Header().Set("Content-Type", "text/event-stream")
|
||
w.WriteHeader(200)
|
||
fmt.Fprintln(w, `data: {"choices":[{"delta":{"content":"Hel"}}]}`)
|
||
fmt.Fprintln(w, `data: {"choices":[{"delta":{"content":"lo"}}]}`)
|
||
fmt.Fprintln(w, `data: {"choices":[{"delta":{},"finish_reason":"stop"}]}`)
|
||
fmt.Fprintln(w, "data: [DONE]")
|
||
return
|
||
}
|
||
w.Header().Set("Content-Type", "application/json")
|
||
w.WriteHeader(200)
|
||
fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
|
||
}))
|
||
}
|
||
|
||
func newTestGateway(t *testing.T, srcs ...config.Source) *Gateway {
|
||
t.Helper()
|
||
td := t.TempDir()
|
||
cfgPath := filepath.Join(td, "config.yaml")
|
||
os.WriteFile(cfgPath, []byte("listen: :0"), 0644)
|
||
cfg := &config.Config{
|
||
Path: cfgPath,
|
||
AdapterDir: filepath.Join(td, "adapters"),
|
||
RuntimeFile: filepath.Join(td, "runtime.json"),
|
||
GatewayKeys: []string{"sk-test"},
|
||
Sources: srcs,
|
||
}
|
||
if err := cfg.ApplyDefaults(); err != nil {
|
||
t.Fatal(err)
|
||
}
|
||
// ★ 测试默认建一条 AUTO 链（每个源一个 tier，按传入顺序），因为内核
|
||
// **不再** seedAuto 自动播种规则了（那条废弃逻辑已删除：AUTO 链为空时
|
||
// 网关应当直接 503，而不是凭空生成一堆用户没配过的档位）。
|
||
//
|
||
// 但既有测试大量依赖"AUTO 能跑"，它们此前是隐式吃 seedAuto 的。这里显式
|
||
// 建链，既保持测试原意（验调度而不是验播种），也让"空链=503"成为
|
||
// 一个**可以单独断言**的行为（见 TestChatAutoEmptyChainIsRejected）。
|
||
//
|
||
// 顺序沿用 seedAuto 的老口径：**源按"最后一个最先"排**（tier 1 = 最后
|
||
// 传入的源）。既有用例靠这个语义挑"high"源（TestChatAUTO 期望 high-m），
|
||
// 改成正序会让一批用例因为"挑到了另一个源"而红，而那与本次重构无关。
|
||
//
|
||
// 所有源放在**同一个 tier**（tier 1），而不是一个源一个 tier：链内同一
|
||
// tier 的多个 slot 会互相 failover，这正是 TestChatAutoChainTierFailover
|
||
// 要验的行为（"a 挂 500 → 请求由 b 服务"）。一源一 tier 会让 a 独占 tier1，
|
||
// b 在 tier2 永远轮不到，"两个源都被打过"这条断言就必然失败。
|
||
if len(cfg.Auto) == 0 && len(srcs) > 0 {
|
||
// 一源一 tier，**倒序**：最后传入的源拿 tier 1（最先尝试）。
|
||
// TestChatAUTO（low 先传、high 后传，期望 high）靠的就是这个语义。
|
||
// 同 tier failover 由 TestChatAutoChainTierFailover 自己显式建链覆盖 ——
|
||
// 两条测试验的是不同机制（tier 优先 vs 同 tier 轮转），不能用同一个
|
||
// 默认夹具同时满足。
|
||
tier := 0
|
||
for i := len(srcs) - 1; i >= 0; i-- {
|
||
tier++
|
||
s := srcs[i]
|
||
for _, m := range s.Models {
|
||
if m.Kind == "image" {
|
||
continue // 图片模型不进 chat 链（与 buildAutoChain 同口径）
|
||
}
|
||
cfg.Auto = append(cfg.Auto, config.ModelScope{
|
||
Model: m.ID, Source: s.Name, Tier: tier,
|
||
})
|
||
}
|
||
}
|
||
}
|
||
c, err := core.NewFromConfig(cfg)
|
||
if err != nil {
|
||
t.Fatalf("core: %v", err)
|
||
}
|
||
t.Cleanup(c.Close)
|
||
g, err := New(c)
|
||
if err != nil {
|
||
t.Fatalf("gateway: %v", err)
|
||
}
|
||
return g
|
||
}
|
||
|
||
func doReq(t *testing.T, g *Gateway, method, path, body string) *httptest.ResponseRecorder {
|
||
t.Helper()
|
||
req, _ := http.NewRequest(method, path, strings.NewReader(body))
|
||
req.Header.Set("Authorization", "Bearer sk-test")
|
||
if body != "" {
|
||
req.Header.Set("Content-Type", "application/json")
|
||
}
|
||
rr := httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
return rr
|
||
}
|
||
|
||
// upstreamCtrl toggles a mocked upstream's behavior between requests.
|
||
type upstreamCtrl struct {
|
||
status int // 0 = healthy; else every request fails with that status
|
||
hits int // chat call count
|
||
}
|
||
|
||
// upstream returns a mocked OpenAI upstream driven by ctrl.status.
|
||
func upstream(t *testing.T, ctrl *upstreamCtrl) *httptest.Server {
|
||
t.Helper()
|
||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
ctrl.hits++
|
||
if ctrl.status != 0 {
|
||
w.WriteHeader(ctrl.status)
|
||
fmt.Fprint(w, `{"error":"boom"}`)
|
||
return
|
||
}
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprintf(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":1,"total_tokens":4}}`)
|
||
}))
|
||
}
|
||
|
||
// TestChatAutoChainTierFailover: AUTO chain, first slot hard-fails, the pass
|
||
// moves on within the same tier and the request is served by the next slot.
|
||
func TestChatAutoChainTierFailover(t *testing.T) {
|
||
a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{}
|
||
aUp := upstream(t, a)
|
||
bUp := upstream(t, b)
|
||
defer aUp.Close()
|
||
defer bUp.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m"}}},
|
||
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m"}}},
|
||
)
|
||
// ★ 显式把两个源放进**同一 tier**：本用例验的是 tier 内失败轮转。
|
||
// 默认夹具是一源一 tier（供 TestChatAUTO 验优先级），两者语义不同，
|
||
// 必须在这里覆盖，否则 b 在 tier2 永远轮不到。
|
||
if err := g.core.SaveAutoRules([]config.ModelScope{
|
||
{Model: "a-m", Source: "a", Tier: 1},
|
||
{Model: "b-m", Source: "b", Tier: 1},
|
||
}); err != nil {
|
||
t.Fatalf("save auto rules: %v", err)
|
||
}
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var cc ChatCompletion
|
||
_ = json.Unmarshal(rr.Body.Bytes(), &cc)
|
||
if cc.Model != "b-m" {
|
||
t.Fatalf("AUTO served %q, want b-m", cc.Model)
|
||
}
|
||
if a.hits == 0 || b.hits == 0 {
|
||
t.Fatalf("hit counts a=%d b=%d, want both > 0", a.hits, b.hits)
|
||
}
|
||
}
|
||
|
||
// TestChatAutoChain503Summary: every AUTO slot fails -> 503 whose message
|
||
// names each failed tier/source/model.
|
||
func TestChatAutoChain503Summary(t *testing.T) {
|
||
a, b := &upstreamCtrl{status: 500}, &upstreamCtrl{status: 500}
|
||
aUp := upstream(t, a)
|
||
bUp := upstream(t, b)
|
||
defer aUp.Close()
|
||
defer bUp.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m"}}},
|
||
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "b-m"}}},
|
||
)
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != http.StatusServiceUnavailable {
|
||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "auto providers failed") ||
|
||
!strings.Contains(rr.Body.String(), "a/a-m") ||
|
||
!strings.Contains(rr.Body.String(), "b/b-m") {
|
||
t.Fatalf("503 must summarize every tier, body=%s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
// TestChatAutoQuotaSkip: a slot whose token quota is exhausted is dropped
|
||
// from scheduling; with no other slot the chain answers 503 naming the quota.
|
||
func TestChatAutoQuotaSkip(t *testing.T) {
|
||
ctrl := &upstreamCtrl{}
|
||
up := upstream(t, ctrl)
|
||
defer up.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m"}}},
|
||
)
|
||
// one slot with an hourly quota of 1 token
|
||
rr := doReq(t, g, "PUT", "/api/auto",
|
||
`{"rules":[{"model":"a-m","tier":0,"token_quota":1,"period":"hour"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
// first request consumes 4 tokens -> quota exhausted
|
||
rr = doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("first status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
// second request must skip the exhausted slot and fail 503
|
||
rr = doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != http.StatusServiceUnavailable {
|
||
t.Fatalf("quota status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "quota exhausted") {
|
||
t.Fatalf("503 must name the quota reason, body=%s", rr.Body.String())
|
||
}
|
||
if ctrl.hits != 1 {
|
||
t.Fatalf("upstream hits = %d, want 1 (exhausted slot must not be called)", ctrl.hits)
|
||
}
|
||
}
|
||
|
||
// TestAutoStatesReportChainHealth: GET /api/auto reports per-slot health for
|
||
// the priority-page UI; a chain edit resets the failure state to zero.
|
||
func TestAutoStatesReportChainHealth(t *testing.T) {
|
||
ctrl := &upstreamCtrl{status: 500}
|
||
up := upstream(t, ctrl)
|
||
defer up.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m"}}},
|
||
)
|
||
doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
|
||
fetch := func() []core.AutoSlotState {
|
||
rr := doReq(t, g, "GET", "/api/auto", "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("get auto status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var body struct {
|
||
Rules []config.ModelScope `json:"rules"`
|
||
States []core.AutoSlotState `json:"states"`
|
||
}
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
|
||
t.Fatalf("unmarshal: %v", err)
|
||
}
|
||
return body.States
|
||
}
|
||
|
||
st := fetch()
|
||
if len(st) != 1 || st[0].Model != "a-m" || st[0].Source != "a" {
|
||
t.Fatalf("want 1 slot a/a-m, got %#v", st)
|
||
}
|
||
if st[0].FailCount == 0 || !st[0].Cooling {
|
||
t.Fatalf("slot must report the failure (fail=%d cooling=%v)", st[0].FailCount, st[0].Cooling)
|
||
}
|
||
|
||
doReq(t, g, "PUT", "/api/auto",
|
||
`{"rules":[{"model":"a-m","tier":0}]}`)
|
||
st = fetch()
|
||
if st[0].FailCount != 0 || st[0].Cooling {
|
||
t.Fatalf("edit must reset health, got %#v", st[0])
|
||
}
|
||
}
|
||
|
||
// TestAutoSaveResetsCooldown: editing the AUTO chain clears the cooldown of
|
||
// its slots, so a fixed upstream is schedulable again without waiting (P1).
|
||
func TestAutoSaveResetsCooldown(t *testing.T) {
|
||
ctrl := &upstreamCtrl{status: 500}
|
||
up := upstream(t, ctrl)
|
||
defer up.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "a-m"}}},
|
||
)
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != http.StatusServiceUnavailable {
|
||
t.Fatalf("expect 503 while upstream down, got %d", rr.Code)
|
||
}
|
||
p := g.core.ProviderForSlot("a-m", "a")
|
||
if p == nil || p.ModelAvailable("a-m") {
|
||
t.Fatal("a-m must be cooling after the failure")
|
||
}
|
||
// editing the chain (same rules) must clear the cooldown immediately
|
||
rr = doReq(t, g, "PUT", "/api/auto",
|
||
`{"rules":[{"model":"a-m","tier":0}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("put auto status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
if !p.ModelAvailable("a-m") {
|
||
t.Fatal("SaveAutoRules must reset the slot cooldown")
|
||
}
|
||
// healed upstream -> AUTO serves again on the next request
|
||
ctrl.status = 0
|
||
rr = doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("AUTO after reset status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestChatSingle(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"mock-model","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var cc ChatCompletion
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &cc); err != nil {
|
||
t.Fatalf("unmarshal: %v", err)
|
||
}
|
||
if cc.Choices[0].Message.Content != "pong" {
|
||
t.Fatalf("content = %q", cc.Choices[0].Message.Content)
|
||
}
|
||
if cc.Usage == nil || cc.Usage.Total != 4 {
|
||
t.Fatalf("usage = %+v", cc.Usage)
|
||
}
|
||
if cc.Model != "mock-model" {
|
||
t.Fatalf("model = %q", cc.Model)
|
||
}
|
||
}
|
||
|
||
func TestChatDisableThinkingPassthrough(t *testing.T) {
|
||
got := ""
|
||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
b, _ := io.ReadAll(r.Body)
|
||
got = string(b)
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}]}`)
|
||
}))
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"mock-model","disable_thinking":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var sent map[string]interface{}
|
||
if err := json.Unmarshal([]byte(got), &sent); err != nil {
|
||
t.Fatalf("upstream body: %v", err)
|
||
}
|
||
if _, has := sent["disable_thinking"]; has {
|
||
t.Fatalf("disable_thinking not stripped: %s", got)
|
||
}
|
||
// openai adapter strips disable_thinking; deepseek would map it to extra_body.thinking.
|
||
// with the passthrough fix the flag now reaches the VM at all.
|
||
}
|
||
|
||
func TestChatMultimodalPassthrough(t *testing.T) {
|
||
gotBody := ""
|
||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
b, _ := io.ReadAll(r.Body)
|
||
gotBody = string(b)
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"choices":[{"message":{"content":"pong"},"finish_reason":"stop"}]}`)
|
||
}))
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"mock-model","messages":[{"role":"user","content":[
|
||
{"type":"text","text":"what is this?"},
|
||
{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}
|
||
]}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var sent struct {
|
||
Messages []struct {
|
||
Content []map[string]interface{} `json:"content"`
|
||
} `json:"messages"`
|
||
}
|
||
if err := json.Unmarshal([]byte(gotBody), &sent); err != nil {
|
||
t.Fatalf("upstream body: %v", err)
|
||
}
|
||
if len(sent.Messages) != 1 || len(sent.Messages[0].Content) != 2 {
|
||
t.Fatalf("multimodal content lost: %s", gotBody)
|
||
}
|
||
}
|
||
|
||
func TestChatAUTO(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "low", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "low-m"}}},
|
||
config.Source{Name: "high", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "high-m"}}},
|
||
)
|
||
// no model -> AUTO -> picks the highest priority source
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"AUTO","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var cc ChatCompletion
|
||
_ = json.Unmarshal(rr.Body.Bytes(), &cc)
|
||
if cc.Model != "high-m" {
|
||
t.Fatalf("AUTO picked %q, want high-m", cc.Model)
|
||
}
|
||
}
|
||
|
||
func TestChatAuthRequired(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
req, _ := http.NewRequest("POST", "/v1/chat/completions",
|
||
strings.NewReader(`{"messages":[{"role":"user","content":"hi"}]}`))
|
||
rr := httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusUnauthorized {
|
||
t.Fatalf("expected 401, got %d", rr.Code)
|
||
}
|
||
}
|
||
|
||
func TestChatStream(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"mock-model","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||
body := rr.Body.String()
|
||
if !strings.Contains(body, "data: [DONE]") {
|
||
t.Fatalf("missing DONE, body=%s", body)
|
||
}
|
||
if !strings.Contains(body, "Hel") || !strings.Contains(body, "lo") {
|
||
t.Fatalf("missing content chunks, body=%s", body)
|
||
}
|
||
}
|
||
|
||
func TestImageGeneration(t *testing.T) {
|
||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`)
|
||
}))
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "img", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "flux-1", Kind: "image"}}})
|
||
rr := doReq(t, g, "POST", "/v1/images/generations",
|
||
`{"model":"flux-1","prompt":"a cat"}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var out map[string]interface{}
|
||
_ = json.Unmarshal(rr.Body.Bytes(), &out)
|
||
data, _ := out["data"].([]interface{})
|
||
if len(data) != 1 {
|
||
t.Fatalf("image data len = %d", len(data))
|
||
}
|
||
}
|
||
|
||
func TestImageAutoFallsOnlyToImageProviders(t *testing.T) {
|
||
imageHits := 0
|
||
chatHits := 0
|
||
img := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
imageHits++
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"created":1,"data":[{"b64_json":"QUJD"}]}`)
|
||
}))
|
||
defer img.Close()
|
||
chatUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
chatHits++
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprint(w, `{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]}`)
|
||
}))
|
||
defer chatUp.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "chat", BaseURL: chatUp.URL, Adapter: "openai", Models: []config.Model{{ID: "chat-m"}}},
|
||
config.Source{Name: "img", BaseURL: img.URL, Adapter: "openai", Models: []config.Model{{ID: "flux", Kind: "image"}}},
|
||
)
|
||
rr := doReq(t, g, "POST", "/v1/images/generations",
|
||
`{"model":"AUTO","prompt":"a cat"}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
if chatHits != 0 {
|
||
t.Fatalf("image AUTO hit chat-only provider: %d chat hits", chatHits)
|
||
}
|
||
if imageHits == 0 {
|
||
t.Fatalf("image AUTO did not hit image provider")
|
||
}
|
||
}
|
||
|
||
func TestKimicodeSigning(t *testing.T) {
|
||
var gotAuth, gotSign string
|
||
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
_, _ = io.ReadAll(r.Body)
|
||
gotAuth = r.Header.Get("Authorization")
|
||
gotSign = r.Header.Get("X-App-Sign")
|
||
w.Header().Set("Content-Type", "application/json")
|
||
fmt.Fprintf(w, `{"choices":[{"message":{"content":"ok"},"finish_reason":"stop"}]}`)
|
||
}))
|
||
defer up.Close()
|
||
|
||
g := newTestGateway(t, config.Source{
|
||
Name: "kimi",
|
||
BaseURL: up.URL,
|
||
Adapter: "kimicode",
|
||
APIKey: "sk-kimi",
|
||
Models: []config.Model{{ID: "kimi-k2"}},
|
||
Meta: map[string]interface{}{"app_id": "app-1", "app_secret": "s3cr3t", "app_agent": "code-agent", "api_key": "sk-kimi"},
|
||
})
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"kimi-k2","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
if gotAuth == "" || !strings.Contains(gotAuth, "sk-kimi") {
|
||
t.Fatalf("expected signed auth, got %q", gotAuth)
|
||
}
|
||
if gotSign == "" {
|
||
t.Fatalf("expected app signature header")
|
||
}
|
||
}
|
||
|
||
func TestModelRoutingPrefix(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-a"}}},
|
||
config.Source{Name: "b", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-b"}}},
|
||
)
|
||
rr := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"model-b","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d", rr.Code)
|
||
}
|
||
}
|
||
|
||
func TestModelsEndpoint(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "a", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-a"}}},
|
||
config.Source{Name: "b", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "model-b"}}},
|
||
)
|
||
rr := doReq(t, g, "GET", "/v1/models", "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d", rr.Code)
|
||
}
|
||
var out map[string]interface{}
|
||
_ = json.Unmarshal(rr.Body.Bytes(), &out)
|
||
if !strings.Contains(rr.Body.String(), "model-a") || !strings.Contains(rr.Body.String(), "model-b") {
|
||
t.Fatalf("missing models: %s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestWebUIServesPage(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
rr := doReq(t, g, "GET", "/", "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status=%d", rr.Code)
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "ModelRouter") {
|
||
t.Fatalf("ui not served")
|
||
}
|
||
}
|
||
|
||
func TestAdaptersAPIUpload(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
// upload adapter
|
||
rr := doReq(t, g, "POST", "/api/adapters",
|
||
`{"name":"testadp","code":"return {name='testadp',endpoint='/chat/completions',transform_request=function(raw) return raw end,transform_response=function(raw) return raw end}"}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("upload status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
rr = doReq(t, g, "GET", "/api/status", "")
|
||
if !strings.Contains(rr.Body.String(), "testadp") {
|
||
t.Fatalf("adapter not listed: %s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestSourcesAPIAddAndPersist(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t)
|
||
rr := doReq(t, g, "POST", "/api/sources",
|
||
fmt.Sprintf(`{"name":"added","base_url":"%s","adapter":"openai","models":[{"id":"new-m","priority":5}]}`, up.URL))
|
||
if rr.Code != 200 {
|
||
t.Fatalf("add source status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
rr = doReq(t, g, "GET", "/v1/models", "")
|
||
if !strings.Contains(rr.Body.String(), "new-m") {
|
||
t.Fatalf("new model not live: %s", rr.Body.String())
|
||
}
|
||
// verify source persisted to config.yaml
|
||
if _, err := os.Stat(g.core.Config().Path); err != nil {
|
||
t.Fatalf("config file not written: %v", err)
|
||
}
|
||
}
|
||
|
||
func TestUIRequiresAuth(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
// / without a key -> redirect to /login
|
||
req, _ := http.NewRequest("GET", "/", nil)
|
||
rr := httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusFound {
|
||
t.Fatalf("expected 302 for /, got %d", rr.Code)
|
||
}
|
||
if loc := rr.Header().Get("Location"); !strings.Contains(loc, "/login") {
|
||
t.Fatalf("expected redirect to /login, got %q", loc)
|
||
}
|
||
// /api/status without a key -> 401
|
||
req, _ = http.NewRequest("GET", "/api/status", nil)
|
||
rr = httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusUnauthorized {
|
||
t.Fatalf("expected 401 for /api/status, got %d", rr.Code)
|
||
}
|
||
// /login page is public
|
||
req, _ = http.NewRequest("GET", "/login", nil)
|
||
rr = httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusOK || !strings.Contains(rr.Body.String(), "登录") {
|
||
t.Fatalf("login page: %d %s", rr.Code, rr.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestLoginAPIAndCookie(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
// bad key -> 401
|
||
req, _ := http.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"wrong"}`))
|
||
req.Header.Set("Content-Type", "application/json")
|
||
rr := httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusUnauthorized {
|
||
t.Fatalf("bad login: %d", rr.Code)
|
||
}
|
||
// good key -> cookie
|
||
req, _ = http.NewRequest("POST", "/api/login", strings.NewReader(`{"key":"sk-test"}`))
|
||
req.Header.Set("Content-Type", "application/json")
|
||
rr = httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusOK {
|
||
t.Fatalf("login: %d", rr.Code)
|
||
}
|
||
cookies := rr.Result().Cookies()
|
||
if len(cookies) == 0 || cookies[0].Name != "gw_key" {
|
||
t.Fatalf("no gw_key cookie set")
|
||
}
|
||
// use cookie to access /api/status
|
||
req, _ = http.NewRequest("GET", "/api/status", nil)
|
||
req.AddCookie(cookies[0])
|
||
rr = httptest.NewRecorder()
|
||
g.Handler().ServeHTTP(rr, req)
|
||
if rr.Code != http.StatusOK {
|
||
t.Fatalf("cookie authed status: %d", rr.Code)
|
||
}
|
||
var out map[string]interface{}
|
||
_ = json.Unmarshal(rr.Body.Bytes(), &out)
|
||
if out["base_url"] == "" {
|
||
t.Fatalf("status missing base_url: %s", rr.Body.String())
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "sk-test") {
|
||
t.Fatalf("status missing gateway_keys: %s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestAPIChatInternal(t *testing.T) {
|
||
up := mockUpstream()
|
||
defer up.Close()
|
||
g := newTestGateway(t, config.Source{Name: "mock", BaseURL: up.URL, Adapter: "openai", Models: []config.Model{{ID: "mock-model"}}})
|
||
rr := doReq(t, g, "POST", "/api/chat",
|
||
`{"model":"mock-model","messages":[{"role":"user","content":"hi"}]}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("api chat status=%d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "pong") {
|
||
t.Fatalf("api chat body=%s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
func TestHasScopeModelWithSourcePrefix(t *testing.T) {
|
||
// zen serves deepseek-v4-flash-free, so the pinning prefix strips; deepseek
|
||
// serves deepseek-v4-pro but NOT v4-flash-free, so a model id like
|
||
// "deepseek-v4-flash-free" (source deepseek + model id with dashes) must NOT
|
||
// be corrupted by prefix stripping.
|
||
g := newTestGateway(t,
|
||
config.Source{Name: "deepseek", BaseURL: "http://d", Adapter: "openai",
|
||
Models: []config.Model{{ID: "deepseek-v4-pro"}}},
|
||
config.Source{Name: "zen", BaseURL: "http://z", Adapter: "openai",
|
||
Models: []config.Model{{ID: "deepseek-v4-flash-free"}}},
|
||
)
|
||
scope := []config.ModelScope{{Model: "deepseek-v4-flash-free"}, {Model: "gpt-5.6-sol"}}
|
||
for _, prefixed := range []string{"zen:deepseek-v4-flash-free", "zen/deepseek-v4-flash-free"} {
|
||
if !g.hasScopeModel(scope, prefixed) {
|
||
t.Errorf("hasScopeModel(scope, %q)=false, want true (prefix must be stripped)", prefixed)
|
||
}
|
||
}
|
||
for _, bare := range []string{"deepseek-v4-flash-free", "gpt-5.6-sol"} {
|
||
if !g.hasScopeModel(scope, bare) {
|
||
t.Errorf("hasScopeModel(scope, %q)=false, want true", bare)
|
||
}
|
||
}
|
||
if g.hasScopeModel(scope, "deepseek-v4-pro") {
|
||
t.Error("hasScopeModel returned true for model outside scope")
|
||
}
|
||
if g.hasScopeModel(scope, "deepseek-v4-flash-free-extra") {
|
||
t.Error("hasScopeModel returned true for unrelated model")
|
||
}
|
||
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "zen:deepseek-v4-flash-free") {
|
||
t.Error("AUTO scope must NOT grant specific model ids even with prefix")
|
||
}
|
||
if g.hasScopeModel([]config.ModelScope{{Model: "AUTO"}}, "deepseek-v4-pro") {
|
||
t.Error("AUTO scope must NOT grant specific model ids")
|
||
}
|
||
}
|
||
|
||
// TestStatusRecorderFlusher pins the SSE-critical contract: the access-log
|
||
// wrapper must implement http.Flusher, otherwise the streaming chat handlers'
|
||
// w.(http.Flusher) assertion yields nil and chunks are never flushed until
|
||
// the response ends (regression for the buffered-SSE bug).
|
||
func TestStatusRecorderFlusher(t *testing.T) {
|
||
var inner *flushRecorder
|
||
rr := httptest.NewRecorder()
|
||
inner = &flushRecorder{ResponseWriter: rr}
|
||
sr := &statusRecorder{ResponseWriter: inner}
|
||
|
||
fl, ok := any(sr).(http.Flusher)
|
||
if !ok {
|
||
t.Fatal("statusRecorder does not implement http.Flusher — SSE streaming is broken")
|
||
}
|
||
if inner.flushed {
|
||
t.Fatal("Flush called before Flush()")
|
||
}
|
||
fl.Flush()
|
||
if !inner.flushed {
|
||
t.Fatal("statusRecorder.Flush did not forward to the underlying writer")
|
||
}
|
||
}
|
||
|
||
// flushRecorder records whether Flush was forwarded.
|
||
type flushRecorder struct {
|
||
http.ResponseWriter
|
||
flushed bool
|
||
}
|
||
|
||
func (f *flushRecorder) Flush() { f.flushed = true }
|
||
|
||
// TestDirectStreamFailoverAuditSource: a direct streaming request whose first
|
||
// candidate hard-fails must be served by the fallback, and the audit record's
|
||
// Source must name the source that actually served the stream (not just the
|
||
// first candidate).
|
||
func TestDirectStreamFailoverAuditSource(t *testing.T) {
|
||
aUp := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||
w.WriteHeader(500)
|
||
fmt.Fprint(w, `{"error":"boom"}`)
|
||
}))
|
||
defer aUp.Close()
|
||
bUp := mockUpstream()
|
||
defer bUp.Close()
|
||
g := newTestGateway(t,
|
||
// Both sources expose the same model id so the direct scheduler has a
|
||
// real fallback candidate after a hard-fails; b is the healthy one.
|
||
config.Source{Name: "a", BaseURL: aUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
||
config.Source{Name: "b", BaseURL: bUp.URL, Adapter: "openai", Models: []config.Model{{ID: "m"}}},
|
||
)
|
||
body := doReq(t, g, "POST", "/v1/chat/completions",
|
||
`{"model":"m","stream":true,"messages":[{"role":"user","content":"hi"}]}`)
|
||
if body.Code != 200 {
|
||
t.Fatalf("status=%d body=%s", body.Code, body.Body.String())
|
||
}
|
||
sb := body.Body.String()
|
||
if !strings.Contains(sb, `"finish_reason":"stop"`) || !strings.Contains(sb, "[DONE]") {
|
||
t.Fatalf("stream incomplete: %s", sb)
|
||
}
|
||
recs := g.stats.AuditRecords(0, 0, "")
|
||
if len(recs) == 0 {
|
||
t.Fatal("no audit records")
|
||
}
|
||
last := recs[len(recs)-1]
|
||
if last.Source != "b" || !last.OK {
|
||
t.Fatalf("audit rec source=%q ok=%v, want source=b ok=true (actual serving source)", last.Source, last.OK)
|
||
}
|
||
}
|
||
|
||
// TestStatusExposesAdapterPools checks that the elastic Lua pool sizing is
|
||
// visible to admins (and only to them) via /api/status, so the grow/shrink
|
||
// algorithm can be observed instead of inferred.
|
||
func TestStatusExposesAdapterPools(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
rr := doReq(t, g, http.MethodGet, "/api/status", "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status = %d, body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var body map[string]interface{}
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
|
||
t.Fatalf("decode: %v", err)
|
||
}
|
||
pools, ok := body["adapter_pools"].([]interface{})
|
||
if !ok {
|
||
t.Fatalf("adapter_pools missing or wrong type: %T", body["adapter_pools"])
|
||
}
|
||
if len(pools) == 0 {
|
||
t.Fatal("adapter_pools must list the loaded adapters")
|
||
}
|
||
first, _ := pools[0].(map[string]interface{})
|
||
for _, field := range []string{"name", "created", "idle", "in_use", "max", "grow_step", "shrink_step"} {
|
||
if _, ok := first[field]; !ok {
|
||
t.Errorf("pool entry missing %q: %v", field, first)
|
||
}
|
||
}
|
||
// idle gateway: nothing booted eagerly
|
||
if c, _ := first["created"].(float64); c != 0 {
|
||
t.Errorf("adapter pool booted eagerly, created=%v", first["created"])
|
||
}
|
||
}
|
||
|
||
// TestStatsRecordsAPIPaging exercises the on-demand records endpoint the
|
||
// dashboard scrolls with: a bounded first page plus a cursor that walks back
|
||
// through the audit file.
|
||
func TestStatsRecordsAPIPaging(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
for i := 0; i < 250; i++ {
|
||
g.stats.Record(Req{
|
||
Time: time.Now().UnixMilli() - int64(250-i)*1000,
|
||
Key: keyID("sk-test"), Type: "chat", Model: "m", Source: "s",
|
||
Prompt: 1, Compl: 1, OK: true, Status: 200,
|
||
})
|
||
}
|
||
|
||
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=50", "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("status = %d body=%s", rr.Code, rr.Body.String())
|
||
}
|
||
var page struct {
|
||
Records []Req `json:"records"`
|
||
NextCursor string `json:"next_cursor"`
|
||
HasMore bool `json:"has_more"`
|
||
}
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
|
||
t.Fatalf("decode: %v", err)
|
||
}
|
||
if len(page.Records) != 50 {
|
||
t.Fatalf("first page = %d records, want 50", len(page.Records))
|
||
}
|
||
if !page.HasMore || page.NextCursor == "" {
|
||
t.Fatal("a long history must advertise more pages plus a cursor")
|
||
}
|
||
// newest first
|
||
for i := 1; i < len(page.Records); i++ {
|
||
if page.Records[i-1].Time < page.Records[i].Time {
|
||
t.Fatalf("page not newest-first at %d", i)
|
||
}
|
||
}
|
||
|
||
// follow the cursor: the next page must continue strictly older
|
||
oldest := page.Records[len(page.Records)-1].Time
|
||
rr2 := doReq(t, g, http.MethodGet,
|
||
"/api/stats/records?limit=50&before="+url.QueryEscape(page.NextCursor), "")
|
||
if rr2.Code != 200 {
|
||
t.Fatalf("page 2 status = %d", rr2.Code)
|
||
}
|
||
var page2 struct {
|
||
Records []Req `json:"records"`
|
||
}
|
||
if err := json.Unmarshal(rr2.Body.Bytes(), &page2); err != nil {
|
||
t.Fatalf("decode page 2: %v", err)
|
||
}
|
||
if len(page2.Records) == 0 {
|
||
t.Fatal("cursor page must return records")
|
||
}
|
||
if page2.Records[0].Time > oldest {
|
||
t.Fatalf("page 2 starts at %d, must continue below %d", page2.Records[0].Time, oldest)
|
||
}
|
||
}
|
||
|
||
// TestStatsSnapshotShipsOneScreen: a dashboard poll must not serialize the whole
|
||
// ring, otherwise the "load the first screen only" contract is broken on the
|
||
// wire even if the UI pages.
|
||
func TestStatsSnapshotShipsOneScreen(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
for i := 0; i < 400; i++ {
|
||
g.stats.Record(Req{
|
||
Time: time.Now().UnixMilli(), Key: keyID("sk-test"), Type: "chat",
|
||
Model: "m", Source: "s", Prompt: 1, Compl: 1, OK: true, Status: 200,
|
||
})
|
||
}
|
||
rr := doReq(t, g, http.MethodGet, "/api/stats", "")
|
||
var snap struct {
|
||
Records []Req `json:"records"`
|
||
Total Stat `json:"total"`
|
||
}
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &snap); err != nil {
|
||
t.Fatalf("decode: %v", err)
|
||
}
|
||
if len(snap.Records) > firstScreenRecords {
|
||
t.Fatalf("snapshot shipped %d records, want <= %d", len(snap.Records), firstScreenRecords)
|
||
}
|
||
// aggregates still cover every request
|
||
if snap.Total.Reqs != 400 {
|
||
t.Fatalf("total reqs = %d, want 400 (aggregates must not be paged away)", snap.Total.Reqs)
|
||
}
|
||
}
|
||
|
||
// TestStatsCsvExportStreams checks the export still emits every row in the
|
||
// window now that it is written straight from the audit walk.
|
||
func TestStatsCsvExportStreams(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
base := time.Now().UnixMilli()
|
||
for i := 0; i < 120; i++ {
|
||
g.stats.Record(Req{
|
||
Time: base + int64(i)*1000, Key: keyID("sk-test"), Type: "chat",
|
||
Model: "m", Source: "s", Prompt: 2, Compl: 3, OK: true, Status: 200,
|
||
})
|
||
}
|
||
rr := doReq(t, g, http.MethodGet,
|
||
"/api/stats?export=csv&from="+strconv.FormatInt(base, 10)+
|
||
"&to="+strconv.FormatInt(base+120*1000, 10), "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("export status = %d", rr.Code)
|
||
}
|
||
lines := strings.Split(strings.TrimSpace(rr.Body.String()), "\n")
|
||
if len(lines) != 121 { // header + 120 rows
|
||
t.Fatalf("csv had %d lines, want 121 (header + 120 records)", len(lines))
|
||
}
|
||
if !strings.HasPrefix(lines[0], "time,key,key_name") {
|
||
t.Fatalf("unexpected csv header: %q", lines[0])
|
||
}
|
||
}
|
||
|
||
// TestStatsRecordsAPIScopedToOwnKey: a non-admin key must only page its own
|
||
// records even when it asks for someone else's.
|
||
func TestStatsRecordsAPIScopedToOwnKey(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
mine := keyID("sk-test")
|
||
for i := 0; i < 5; i++ {
|
||
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: mine, Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
|
||
}
|
||
for i := 0; i < 5; i++ {
|
||
g.stats.Record(Req{Time: time.Now().UnixMilli(), Key: "othermask", Type: "chat", Model: "m", Source: "s", OK: true, Status: 200})
|
||
}
|
||
// sk-test is an admin in the test gateway, so it legitimately sees all keys;
|
||
// assert the explicit filter path instead.
|
||
rr := doReq(t, g, http.MethodGet, "/api/stats/records?limit=100&key=othermask", "")
|
||
var page struct {
|
||
Records []Req `json:"records"`
|
||
}
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &page); err != nil {
|
||
t.Fatalf("decode: %v", err)
|
||
}
|
||
for _, r := range page.Records {
|
||
if r.Key != "othermask" {
|
||
t.Fatalf("key filter leaked %q", r.Key)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestKeyDeleteRoundTrip is the end-to-end version of the bug users hit:
|
||
// deleting a gateway key from the admin UI. It also asserts the GET-on-a-key
|
||
// error names the right verb, since that message is what a mis-methoded client
|
||
// actually sees.
|
||
func TestKeyDeleteRoundTrip(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
|
||
rr := doReq(t, g, http.MethodPost, "/api/keys", `{"name":"doomed","role":"user"}`)
|
||
if rr.Code != 200 {
|
||
t.Fatalf("create: %d %s", rr.Code, rr.Body.String())
|
||
}
|
||
var created struct {
|
||
Key struct {
|
||
Key string `json:"key"`
|
||
Name string `json:"name"`
|
||
} `json:"key"`
|
||
}
|
||
if err := json.Unmarshal(rr.Body.Bytes(), &created); err != nil {
|
||
t.Fatalf("decode create: %v", err)
|
||
}
|
||
if created.Key.Key == "" {
|
||
t.Fatalf("no key returned: %s", rr.Body.String())
|
||
}
|
||
|
||
// DELETE removes it
|
||
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
|
||
if rr.Code != 200 {
|
||
t.Fatalf("delete: %d %s", rr.Code, rr.Body.String())
|
||
}
|
||
for _, k := range g.core.ListKeys() {
|
||
if k.Key == created.Key.Key {
|
||
t.Fatal("key still present after DELETE")
|
||
}
|
||
}
|
||
|
||
// deleting again is a clean 404, not a 500
|
||
rr = doReq(t, g, http.MethodDelete, "/api/keys/"+url.PathEscape(created.Key.Key), "")
|
||
if rr.Code != http.StatusNotFound {
|
||
t.Fatalf("second delete: %d, want 404", rr.Code)
|
||
}
|
||
}
|
||
|
||
// TestKeyGetOnSpecificKeyNamesTheVerb: a GET on /api/keys/{key} is what a
|
||
// client that forgot to set method:"DELETE" ends up sending. The error must say
|
||
// which verb to use instead of just "use GET /api/keys", which reads as if the
|
||
// caller did nothing wrong.
|
||
func TestKeyGetOnSpecificKeyNamesTheVerb(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
rr := doReq(t, g, http.MethodGet, "/api/keys/sk-whatever", "")
|
||
if rr.Code != http.StatusNotFound {
|
||
t.Fatalf("status = %d, want 404", rr.Code)
|
||
}
|
||
body := rr.Body.String()
|
||
for _, want := range []string{"DELETE /api/keys/", "PUT /api/keys/"} {
|
||
if !strings.Contains(body, want) {
|
||
t.Errorf("error message should mention %q, got: %s", want, body)
|
||
}
|
||
}
|
||
}
|
||
|
||
// TestAdapterGetOnSpecificAdapterNamesDelete is the same guard for the adapters
|
||
// endpoint, whose delAdapter helper had the identical missing-method bug.
|
||
func TestAdapterGetOnSpecificAdapterNamesDelete(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
rr := doReq(t, g, http.MethodGet, "/api/adapters/openai", "")
|
||
if rr.Code != http.StatusNotFound {
|
||
t.Fatalf("status = %d, want 404", rr.Code)
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "DELETE /api/adapters/openai") {
|
||
t.Errorf("error should name the DELETE verb, got: %s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
// TestCannotDeleteOwnKey keeps the admin from locking themselves out.
|
||
func TestCannotDeleteOwnKey(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
rr := doReq(t, g, http.MethodDelete, "/api/keys/sk-test", "")
|
||
if rr.Code != http.StatusBadRequest {
|
||
t.Fatalf("status = %d, want 400", rr.Code)
|
||
}
|
||
if !strings.Contains(rr.Body.String(), "logged in with") {
|
||
t.Errorf("unexpected message: %s", rr.Body.String())
|
||
}
|
||
}
|
||
|
||
// TestGatewayRingStaysSmall pins the resident record ring at the small default.
|
||
// The gateway used to construct Stats with a 10000-record ring, which after the
|
||
// full-history aggregate scan meant 10000 resident Req structs — putting back the
|
||
// ~25 MB the on-demand paging work removed. Aggregates cover all of history;
|
||
// the ring only needs the first screen plus the 5-minute status windows.
|
||
func TestGatewayRingStaysSmall(t *testing.T) {
|
||
g := newTestGateway(t)
|
||
g.stats.mu.Lock()
|
||
max := g.stats.maxRecs
|
||
g.stats.mu.Unlock()
|
||
if max != defaultRingSize {
|
||
t.Fatalf("gateway record ring = %d, want %d (a big ring reintroduces the resident-memory cost)",
|
||
max, defaultRingSize)
|
||
}
|
||
}
|