From 22decf2bd30f3d441b656bf48bdb0dffacf11456 Mon Sep 17 00:00:00 2001 From: JianFeeeee Date: Sat, 3 Oct 2026 07:45:54 +0800 Subject: [PATCH] =?UTF-8?q?feat(keys):=20=E6=AF=8F=E4=B8=AA=E5=AF=86?= =?UTF-8?q?=E9=92=A5=E5=8F=AF=E9=85=8D=E7=8B=AC=E7=AB=8B=20AUTO=20?= =?UTF-8?q?=E9=93=BE=EF=BC=8C=E7=AE=A1=E7=90=86=E5=91=98=E4=BB=A3=E9=85=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 此前 AUTO 链是全局单值（cfg.Auto + Core.AutoChain()），所有用户共用一条链。 管理员无法为某个用户单独指定调度链。 按 per-key 覆盖 + 全局兜底实现： - config.GWKey 增加 Auto 字段与 HasOwnAuto()。未配置即继承全局链， 存量部署零改动，新密钥天然继承全局链。 - Core 把 buildAutoChain 的归一化逻辑抽成 chainForRules，全局链、 生图链、per-key 链共用同一套编译，避免两处漂移。 - Core 增加 keyAutoChains 缓存 + AutoChainFor(key)。请求路径读缓存不 加锁，与全局链查询一致。缓存整表原子替换，不会看到半成品。 - 请求侧 chat.go 改用 AutoChainFor(reqKey)。冷却仍在 Provider 上按 model+source 共享：两条链指向同一个 slot 时共用冷却，与今天单链行为 相同，也避免为 per-key 维度重构冷却而改变现有可观测语义。 - API：GET/PUT/DELETE /api/keys/{key}/auto（admin），GET /api/keys/me/auto（任意角色，只能读自己的）。空 PUT 与 DELETE 等价于 "恢复继承"，无法持久化一条会 503 的空链。写入时回报解析出的槽位数， 让管理员当场看到模型名写错，而不是等用户下次请求 503。 - 源变更时一并重编译 per-key 链，加源后无需重启即可生效。 判据 18 条，5 个变异全部被抓住：AutoChainFor 忽略 key、清空后不重建 缓存、源变更不重建、空 PUT 落盘成空链、me/auto 误要求 admin。 UI 判据做变异时发现漏放：只查函数定义存在，删掉按钮后仍通过。已改为 断言 keyCanvasHtml 内的调用点。 端到端实测（真实 HTTP + 两个 mock 上游）：admin 与 bob 初始同为 m-fast， 给 bob 配 m-cheap 后两者分流，重启后仍分流，DELETE 后 bob 回到 m-fast。 Co-Authored-By: ModelRouter --- internal/config/config.go | 24 ++- internal/core/core.go | 92 ++++++++- internal/core/key_auto_test.go | 221 ++++++++++++++++++++++ internal/gateway/chat.go | 3 +- internal/gateway/key_auto_api_test.go | 257 ++++++++++++++++++++++++++ internal/gateway/keys.go | 108 +++++++++++ internal/gateway/ui/index.html | 126 +++++++++++++ internal/gateway/ui_key_auto_test.go | 78 ++++++++ 8 files changed, 897 insertions(+), 12 deletions(-) create mode 100644 internal/core/key_auto_test.go create mode 100644 internal/gateway/key_auto_api_test.go create mode 100644 internal/gateway/ui_key_auto_test.go diff --git a/internal/config/config.go b/internal/config/config.go index 4e85b43..e4c09ac 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -384,15 +384,25 @@ type SourceTemplate struct { // (full management) or "user" (sees only its own key); Models is the allowed // model scope with per-model token quota (0 = unlimited). type GWKey struct { - Key string `yaml:"key" json:"key"` - Role string `yaml:"role" json:"role"` - Name string `yaml:"name,omitempty" json:"name,omitempty"` - Models []ModelScope `yaml:"models,omitempty" json:"models,omitempty"` - Note string `yaml:"note,omitempty" json:"note,omitempty"` - CreatedAt int64 `yaml:"created_at,omitempty" json:"created_at,omitempty"` - Seed bool `yaml:"seed,omitempty" json:"seed,omitempty"` // true if migrated from config gateway_keys + Key string `yaml:"key" json:"key"` + Role string `yaml:"role" json:"role"` + Name string `yaml:"name,omitempty" json:"name,omitempty"` + // Auto is this key's own AUTO scheduling chain. It overrides the global + // cfg.Auto for requests authenticated with this key; nil (or an empty + // list) means "inherit the global chain", so keys created before per-key + // chains existed keep working untouched. + Auto []ModelScope `yaml:"auto,omitempty" json:"auto,omitempty"` + Models []ModelScope `yaml:"models,omitempty" json:"models,omitempty"` + Note string `yaml:"note,omitempty" json:"note,omitempty"` + CreatedAt int64 `yaml:"created_at,omitempty" json:"created_at,omitempty"` + Seed bool `yaml:"seed,omitempty" json:"seed,omitempty"` // true if migrated from config gateway_keys } +// HasOwnAuto reports whether the key declares its own AUTO chain rather than +// inheriting the global one. An explicitly empty list still counts as "no own +// chain" so a user can clear an override back to the global default. +func (k GWKey) HasOwnAuto() bool { return len(k.Auto) > 0 } + // BillingDSL holds declarative per-URL pricing profiles for the billing plugin. // // Profiles are the "let the user pick" axis: the same upstream URL can be diff --git a/internal/core/core.go b/internal/core/core.go index f2addaf..de8b0ee 100644 --- a/internal/core/core.go +++ b/internal/core/core.go @@ -37,6 +37,12 @@ type Core struct { registry *provider.Registry autoChain atomic.Pointer[scheduler.Chain] // chat AUTO chain autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain + + // keyAutoChains caches per-key AUTO chains, keyed by the gateway key. + // Entries are rebuilt by rebuildKeyAutoChains whenever sources or key + // config change; AutoChainFor reads them without taking c.mu so request + // handling stays lock-free like the global chain lookup. + keyAutoChains atomic.Pointer[map[string]*scheduler.Chain] } // New builds the core from a config file plus runtime overlay. @@ -257,6 +263,68 @@ func (c *Core) Registry() *provider.Registry { return c.registry } // AutoChain returns the current AUTO scheduling chain. func (c *Core) AutoChain() *scheduler.Chain { return c.autoChain.Load() } +// AutoChainFor returns the AUTO chain to use for a request authenticated with +// key. A key that declares its own chain gets that chain; every other key +// inherits the global chain, so keys created before per-key chains existed +// keep their current behaviour with no configuration change. +// +// It returns (chain, ok). ok is false only when the key has its own chain but +// it compiled to nothing usable (e.g. every slot names a model that no longer +// exists) — the caller must surface that rather than silently downgrading the +// user to the global chain, which would be a confusing, invisible switch. +func (c *Core) AutoChainFor(key string) (*scheduler.Chain, bool) { + if m := c.keyAutoChains.Load(); m != nil { + if ch, ok := (*m)[key]; ok { + return ch, true + } + } + return c.autoChain.Load(), true +} + +// rebuildKeyAutoChains recompiles every per-key AUTO chain from the current +// config and provider registry, then swaps the cache in one shot so requests +// never observe a half-built map. Keys without their own chain are absent +// from the map and fall back to the global chain. +func (c *Core) rebuildKeyAutoChains() { + m := make(map[string]*scheduler.Chain) + for _, k := range c.cfg.Keys { + if !k.HasOwnAuto() { + continue + } + m[k.Key] = c.chainForRules(k.Auto, false) + } + c.keyAutoChains.Store(&m) +} + +// SaveKeyAuto persists one key's own AUTO chain and recompiles the per-key +// chain cache. Passing an empty list clears the override so the key inherits +// the global chain again. +func (c *Core) SaveKeyAuto(key string, entries []config.ModelScope) error { + c.mu.Lock() + defer c.mu.Unlock() + idx := -1 + for i, k := range c.cfg.Keys { + if k.Key == key { + idx = i + break + } + } + if idx < 0 { + return fmt.Errorf("key not found") + } + if err := ValidateScopeQuotas(entries); err != nil { + return err + } + c.cfg.Keys[idx].Auto = cleanScopes(entries) + if err := c.saveConfig(); err != nil { + return err + } + // Rebuild the whole cache so a cleared override stops shadowing the + // global chain for this key. + c.rebuildKeyAutoChains() + return nil +} + // AutoImageChain returns the persisted image-generation AUTO chain snapshot. func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() } @@ -546,6 +614,7 @@ func (c *Core) rebuildRegistry() error { } c.buildAutoChain() c.buildAutoImageChain() + c.rebuildKeyAutoChains() return nil } @@ -555,17 +624,28 @@ func (c *Core) rebuildRegistry() error { // back to the source's best chat model and cooldown/quota bookkeeping would // key on a name that never matches. func (c *Core) buildAutoChain() { + c.autoChain.Store(c.chainForRules(c.cfg.Auto, false)) +} + +// chainForRules compiles a chat AUTO chain from a rule list. kindFilter is +// false for the chat chain (image models are skipped) and true for the image +// chain (only image models participate). Per-key chains always compile as +// chat chains: a user's key configures the models *they* talk to. +func (c *Core) chainForRules(rules []config.ModelScope, imageOnly bool) *scheduler.Chain { prov := func(model, source string) scheduler.Provider { p := c.registry.ProviderForSlot(model, source) if p == nil { return nil } - if m := p.ModelByID(model); m != nil && m.Kind == "image" { + m := p.ModelByID(model) + if imageOnly && (m == nil || m.Kind != "image") { + return nil + } + if !imageOnly && m != nil && m.Kind == "image" { return nil } return p } - rules := c.cfg.Auto sr := make([]scheduler.Rule, 0, len(rules)) for _, e := range rules { model, source := e.Model, e.Source @@ -574,7 +654,11 @@ func (c *Core) buildAutoChain() { model = exact } if m := p.ModelByID(model); m != nil && m.Kind == "image" { - continue // image-kind models never join the chat AUTO chain + if !imageOnly { + continue // image-kind models never join the chat AUTO chain + } + } else if imageOnly { + continue } if source == "" { source = p.Name() @@ -589,7 +673,7 @@ func (c *Core) buildAutoChain() { Hours: e.Hours, }) } - c.autoChain.Store(scheduler.BuildChain(sr, prov)) + return scheduler.BuildChain(sr, prov) } // buildAutoImageChain rebuilds the image-generation AUTO chain snapshot. diff --git a/internal/core/key_auto_test.go b/internal/core/key_auto_test.go new file mode 100644 index 0000000..170a0b2 --- /dev/null +++ b/internal/core/key_auto_test.go @@ -0,0 +1,221 @@ +package core + +import ( + "path/filepath" + "testing" + + "llmsproxy/internal/config" + "llmsproxy/internal/scheduler" +) + +// Per-key AUTO chains let an admin give one user their own scheduling chain +// while every other key keeps the global one. These tests pin the three +// properties that make that safe: +// +// 1. a key WITH its own chain gets it; +// 2. a key WITHOUT one inherits the global chain (backward compatible); +// 3. the override is actually persisted, so it survives a restart. +// +// A test that only checked (1) would pass even if AutoChainFor ignored the key +// and always returned the global chain whenever the two happened to be equal, +// so (2) uses deliberately *different* chains. + +func perKeyTestConfig(t *testing.T) *config.Config { + t.Helper() + dir := t.TempDir() + return &config.Config{ + Path: filepath.Join(dir, "config.yaml"), + AdapterDir: filepath.Join(dir, "adapters"), + RuntimeFile: filepath.Join(dir, "runtime.json"), + Listen: "127.0.0.1:0", + DefaultModel: "AUTO", + GatewayKeys: []string{"sk-gw-admin"}, + Sources: []config.Source{ + {Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai", + Models: []config.Model{{ID: "gpt-4o", Priority: 10}, {ID: "gpt-4o-mini", Priority: 5}}}, + }, + // Global chain points at gpt-4o. + Auto: []config.ModelScope{{Model: "gpt-4o", Source: "s1", Tier: 1}}, + Keys: []config.GWKey{ + {Key: "sk-gw-admin", Role: "admin"}, + {Key: "sk-gw-inherits", Role: "user", Name: "inherits"}, + // This key's own chain points at a DIFFERENT model, so returning + // the global chain by mistake is detectable. + {Key: "sk-gw-own", Role: "user", Name: "own", + Auto: []config.ModelScope{{Model: "gpt-4o-mini", Source: "s1", Tier: 1}}}, + }, + } +} + +func firstSlotModel(t *testing.T, c *Core, key string) string { + t.Helper() + ch, _ := c.AutoChainFor(key) + if ch == nil || len(ch.Tiers) == 0 { + t.Fatalf("key %s: no chain tiers", key) + } + slots := ch.Tiers[0].Slots + if len(slots) == 0 { + t.Fatalf("key %s: tier 1 has no slots", key) + } + return slots[0].Model +} + +// A key that declares its own chain must be scheduled by it, not the global. +func TestAutoChainForUsesKeyOwnChain(t *testing.T) { + c := newTestCore(t, perKeyTestConfig(t)) + + if got := firstSlotModel(t, c, "sk-gw-own"); got != "gpt-4o-mini" { + t.Fatalf("key with own chain must use it, got %q (want gpt-4o-mini)", got) + } + // Sanity: the global chain really is different, so the assertion above + // cannot pass by accident. + if got := firstSlotModel(t, c, "sk-gw-admin"); got != "gpt-4o" { + t.Fatalf("admin must use global chain, got %q (want gpt-4o)", got) + } +} + +// A key with no own chain inherits the global one. This is what keeps every +// pre-existing deployment working with no config change. +func TestAutoChainForInheritsGlobalWhenUnset(t *testing.T) { + c := newTestCore(t, perKeyTestConfig(t)) + + if got := firstSlotModel(t, c, "sk-gw-inherits"); got != "gpt-4o" { + t.Fatalf("key without own chain must inherit global, got %q (want gpt-4o)", got) + } +} + +// A key that is not in the key table at all (e.g. a seed key) must also +// inherit the global chain rather than erroring or returning nil. +func TestAutoChainForUnknownKeyInheritsGlobal(t *testing.T) { + c := newTestCore(t, perKeyTestConfig(t)) + + ch, ok := c.AutoChainFor("sk-gw-never-configured") + if !ok || ch == nil || len(ch.Tiers) == 0 { + t.Fatalf("unknown key must fall back to global chain, got ok=%v chain=%v", ok, ch) + } +} + +// The per-key chain must survive a save/reload round trip: an admin configuring +// a chain in the WebUI and restarting the gateway must not silently lose it. +func TestSaveKeyAutoPersistsAndApplies(t *testing.T) { + cfg := perKeyTestConfig(t) + c := newTestCore(t, cfg) + + if err := c.SaveKeyAuto("sk-gw-inherits", []config.ModelScope{ + {Model: "gpt-4o-mini", Source: "s1", Tier: 1}, + }); err != nil { + t.Fatalf("SaveKeyAuto: %v", err) + } + if got := firstSlotModel(t, c, "sk-gw-inherits"); got != "gpt-4o-mini" { + t.Fatalf("saved chain must apply immediately, got %q", got) + } + + // Reload from disk exactly like a restart does. + reloaded, err := config.Load(cfg.Path) + if err != nil { + t.Fatalf("reload config: %v", err) + } + c2 := newTestCore(t, reloaded) + if got := firstSlotModel(t, c2, "sk-gw-inherits"); got != "gpt-4o-mini" { + t.Fatalf("per-key chain must survive restart, got %q", got) + } + // The other key must be unaffected by its neighbour's override. + if got := firstSlotModel(t, c2, "sk-gw-admin"); got != "gpt-4o" { + t.Fatalf("override must not leak to other keys, got %q", got) + } +} + +// Clearing the override must return the key to the global chain, both in +// memory and after a restart. This is what the WebUI's "use global" toggle +// does, so a stale shadowing entry would silently pin the user to their old +// chain forever. +func TestSaveKeyAutoClearRestoresGlobal(t *testing.T) { + cfg := perKeyTestConfig(t) + c := newTestCore(t, cfg) + + if err := c.SaveKeyAuto("sk-gw-own", nil); err != nil { + t.Fatalf("clear: %v", err) + } + if got := firstSlotModel(t, c, "sk-gw-own"); got != "gpt-4o" { + t.Fatalf("cleared key must inherit global again, got %q", got) + } + + reloaded, err := config.Load(cfg.Path) + if err != nil { + t.Fatalf("reload: %v", err) + } + c2 := newTestCore(t, reloaded) + if got := firstSlotModel(t, c2, "sk-gw-own"); got != "gpt-4o" { + t.Fatalf("clear must survive restart, got %q", got) + } +} + +// Saving an unknown key must fail loudly rather than creating a shadow entry +// that no request can ever match. +func TestSaveKeyAutoUnknownKeyFails(t *testing.T) { + c := newTestCore(t, perKeyTestConfig(t)) + + if err := c.SaveKeyAuto("sk-gw-nope", []config.ModelScope{{Model: "gpt-4o", Source: "s1"}}); err == nil { + t.Fatal("SaveKeyAuto on unknown key must return an error") + } +} + +// Per-key quota validation must run, exactly like the model scope path: an +// override is another place an operator can write a bad budget. +func TestSaveKeyAutoRejectsBadQuota(t *testing.T) { + c := newTestCore(t, perKeyTestConfig(t)) + + err := c.SaveKeyAuto("sk-gw-inherits", []config.ModelScope{ + {Model: "gpt-4o", Source: "s1", TokenQuota: -5}, + }) + if err == nil { + t.Fatal("negative token quota in a per-key auto chain must be rejected") + } +} + +// Source changes must recompile per-key chains. A key whose chain names a +// model that does not exist yet compiles to an empty chain (BuildChain drops +// slots it cannot resolve — sending the request into a black hole would be +// worse), and must start working as soon as the source appears. +func TestKeyAutoChainsRebuildOnSourceChange(t *testing.T) { + cfg := perKeyTestConfig(t) + cfg.Keys = append(cfg.Keys, config.GWKey{ + Key: "sk-gw-late", Role: "user", Name: "late", + Auto: []config.ModelScope{{Model: "new-model", Source: "s2", Tier: 1}}, + }) + c := newTestCore(t, cfg) + + // s2 does not exist yet, so the chain has no usable slot. + if ch, _ := c.AutoChainFor("sk-gw-late"); chainSlotsOf(ch) != 0 { + t.Fatalf("chain naming an unknown model must compile empty, got %d slots", + chainSlotsOf(ch)) + } + // The other keys are unaffected by the broken one. + if got := firstSlotModel(t, c, "sk-gw-admin"); got != "gpt-4o" { + t.Fatalf("broken per-key chain must not affect others, got %q", got) + } + + // Add the source; the per-key chain must pick it up without a restart. + cfg.Sources = append(cfg.Sources, config.Source{ + Name: "s2", BaseURL: "http://127.0.0.1:2/v1", Adapter: "openai", + Models: []config.Model{{ID: "new-model", Priority: 10}}, + }) + if err := c.Reload(); err != nil { + t.Fatalf("Reload: %v", err) + } + if got := firstSlotModel(t, c, "sk-gw-late"); got != "new-model" { + t.Fatalf("per-key chain must resolve after the source is added, got %q", got) + } +} + +// chainSlotsOf counts usable slots, tolerating a nil chain. +func chainSlotsOf(ch *scheduler.Chain) int { + if ch == nil { + return 0 + } + n := 0 + for _, tn := range ch.Tiers { + n += len(tn.Slots) + } + return n +} diff --git a/internal/gateway/chat.go b/internal/gateway/chat.go index b84eda0..2e3af12 100644 --- a/internal/gateway/chat.go +++ b/internal/gateway/chat.go @@ -400,7 +400,8 @@ func (g *Gateway) handleChat(w http.ResponseWriter, r *http.Request) { // TestHooksFireOnRealDirectChat, not by reading the code. g.fireStart(r.Context(), &req, "chat", model, len(req.Messages), len(req.Tools)) if isAuto(model) { - chain := g.core.AutoChain() + // The key's own chain wins; keys without one inherit the global chain. + chain, _ := g.core.AutoChainFor(reqKey(r.Context())) if chain == nil || len(chain.Tiers) == 0 { writeError(w, http.StatusServiceUnavailable, "no_provider", "no auto slot configured") return diff --git a/internal/gateway/key_auto_api_test.go b/internal/gateway/key_auto_api_test.go new file mode 100644 index 0000000..c2ab771 --- /dev/null +++ b/internal/gateway/key_auto_api_test.go @@ -0,0 +1,257 @@ +package gateway + +import ( + "encoding/json" + "net/http" + "net/http/httptest" + "path/filepath" + "os" + "strings" + "testing" + + "llmsproxy/internal/config" + "llmsproxy/internal/core" +) + +// Per-key AUTO chain admin API. The behaviours pinned here are the ones that +// are easy to get wrong and invisible when they are: +// +// - /api/keys/me/auto must be reachable by a NON-admin (it is their own +// chain); routing it through the generic {key}/auto handler would 403 it. +// - /api/keys/{key}/auto must 403 a non-admin. +// - an empty PUT and a DELETE mean the same thing ("inherit global"), so the +// WebUI cannot persist a chain that would then 503 with no slots. +// - a PUT that resolves to zero slots is reported back, so an admin who +// mistyped a model learns it now instead of from the user's next 503. + +// keyAutoGateway builds a gateway with one admin key and one user key. +func keyAutoGateway(t *testing.T) *Gateway { + t.Helper() + td := t.TempDir() + cfgPath := filepath.Join(td, "config.yaml") + if err := os.WriteFile(cfgPath, []byte("listen: :0"), 0644); err != nil { + t.Fatal(err) + } + cfg := &config.Config{ + Path: cfgPath, + AdapterDir: filepath.Join(td, "adapters"), + RuntimeFile: filepath.Join(td, "runtime.json"), + DefaultModel: "AUTO", + GatewayKeys: []string{"sk-admin"}, + Keys: []config.GWKey{ + {Key: "sk-admin", Role: "admin", Name: "admin"}, + {Key: "sk-user", Role: "user", Name: "user"}, + }, + Sources: []config.Source{ + {Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai", + Models: []config.Model{{ID: "gpt-4o", Priority: 10}, {ID: "gpt-4o-mini", Priority: 5}}}, + }, + } + if err := cfg.ApplyDefaults(); err != nil { + t.Fatal(err) + } + 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 +} + +// doReqAs issues an authenticated request as a specific key. +func doReqAs(t *testing.T, g *Gateway, key, method, path, body string) *httptest.ResponseRecorder { + t.Helper() + req, _ := http.NewRequest(method, path, strings.NewReader(body)) + req.Header.Set("Authorization", "Bearer "+key) + if body != "" { + req.Header.Set("Content-Type", "application/json") + } + rr := httptest.NewRecorder() + g.Handler().ServeHTTP(rr, req) + return rr +} + +func TestKeyAutoAPIPermissions(t *testing.T) { + g := keyAutoGateway(t) + + // A user may read their own chain... + if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/me/auto", ""); rr.Code != http.StatusOK { + t.Fatalf("GET /api/keys/me/auto as user = %d, want 200 (body %s)", rr.Code, rr.Body.String()) + } + // ...but not another key's. + if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/sk-admin/auto", ""); rr.Code != http.StatusForbidden { + t.Fatalf("GET other key's auto as user = %d, want 403", rr.Code) + } + // ...and not their own via the admin route either. + if rr := doReqAs(t, g, "sk-user", "PUT", "/api/keys/sk-user/auto", `{"auto":[{"model":"gpt-4o"}]}`); rr.Code != http.StatusForbidden { + t.Fatalf("user PUT own auto = %d, want 403 (only admin configures)", rr.Code) + } +} + +func TestKeyAutoAPICrudAndInherit(t *testing.T) { + g := keyAutoGateway(t) + + // Initially inherits. + var got struct { + Inherits bool `json:"inherits"` + Auto []config.ModelScope `json:"auto"` + } + rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "") + if rr.Code != http.StatusOK { + t.Fatalf("GET = %d: %s", rr.Code, rr.Body.String()) + } + if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Inherits || len(got.Auto) != 0 { + t.Fatalf("fresh key must inherit global, got inherits=%v auto=%v", got.Inherits, got.Auto) + } + + // Admin sets a chain; the response reports the resolved slot count. + rr = doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", + `{"auto":[{"model":"gpt-4o-mini","source":"s1","tier":1}]}`) + if rr.Code != http.StatusOK { + t.Fatalf("PUT = %d: %s", rr.Code, rr.Body.String()) + } + var put struct { + Inherits bool `json:"inherits"` + Slots int `json:"slots"` + } + if err := json.Unmarshal(rr.Body.Bytes(), &put); err != nil { + t.Fatal(err) + } + if put.Inherits { + t.Fatal("after a non-empty PUT the key must no longer inherit") + } + if put.Slots != 1 { + t.Fatalf("resolved slots = %d, want 1", put.Slots) + } + + // And it reads back as an own chain. + rr = doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "") + if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if got.Inherits || len(got.Auto) != 1 || got.Auto[0].Model != "gpt-4o-mini" { + t.Fatalf("own chain must read back, got inherits=%v auto=%v", got.Inherits, got.Auto) + } + + // DELETE restores inheritance. + if rr := doReqAs(t, g, "sk-admin", "DELETE", "/api/keys/sk-user/auto", ""); rr.Code != http.StatusOK { + t.Fatalf("DELETE = %d: %s", rr.Code, rr.Body.String()) + } + rr = doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "") + if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Inherits { + t.Fatal("after DELETE the key must inherit the global chain again") + } +} + +// An empty PUT must behave like DELETE. Otherwise a WebUI that submits an +// empty list (every slot deleted in the editor) would persist an empty chain +// and the next AUTO request would fail with no_provider. +func TestKeyAutoAPIEmptyPutClearsOverride(t *testing.T) { + g := keyAutoGateway(t) + + if rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", + `{"auto":[{"model":"gpt-4o","source":"s1","tier":1}]}`); rr.Code != http.StatusOK { + t.Fatalf("seed PUT = %d: %s", rr.Code, rr.Body.String()) + } + rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", `{"auto":[]}`) + if rr.Code != http.StatusOK { + t.Fatalf("empty PUT = %d: %s", rr.Code, rr.Body.String()) + } + var got struct { + Inherits bool `json:"inherits"` + } + if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", ""); rr.Code != http.StatusOK { + t.Fatal(rr.Body.String()) + } + if err := json.Unmarshal(rr.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Inherits { + t.Fatal("an empty PUT must clear the override, not persist an empty chain") + } +} + +// A chain naming a model that does not exist compiles to zero slots. The API +// must say so at write time instead of letting the user's next request 503. +func TestKeyAutoAPIReportsUnresolvableSlots(t *testing.T) { + g := keyAutoGateway(t) + + rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", + `{"auto":[{"model":"does-not-exist","source":"s1","tier":1}]}`) + if rr.Code != http.StatusOK { + t.Fatalf("PUT = %d: %s", rr.Code, rr.Body.String()) + } + var put struct { + Inherits bool `json:"inherits"` + Slots int `json:"slots"` + } + if err := json.Unmarshal(rr.Body.Bytes(), &put); err != nil { + t.Fatal(err) + } + if put.Slots != 0 { + t.Fatalf("unknown model must resolve to 0 slots, got %d", put.Slots) + } +} + +func TestKeyAutoAPIUnknownKey404(t *testing.T) { + g := keyAutoGateway(t) + + if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-nope/auto", ""); rr.Code != http.StatusNotFound { + t.Fatalf("GET unknown key auto = %d, want 404", rr.Code) + } + if rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-nope/auto", `{"auto":[{"model":"gpt-4o"}]}`); rr.Code != http.StatusNotFound { + t.Fatalf("PUT unknown key auto = %d, want 404", rr.Code) + } +} + +// A bad quota must be rejected with 400, not persisted and not silently +// clamped — same rule as the model scope path. +func TestKeyAutoAPIRejectsBadQuota(t *testing.T) { + g := keyAutoGateway(t) + + rr := doReqAs(t, g, "sk-admin", "PUT", "/api/keys/sk-user/auto", + `{"auto":[{"model":"gpt-4o","source":"s1","tier":1,"token_quota":-1}]}`) + if rr.Code != http.StatusBadRequest { + t.Fatalf("negative quota PUT = %d, want 400 (body %s)", rr.Code, rr.Body.String()) + } + // Nothing may have been persisted. + var got struct { + Inherits bool `json:"inherits"` + } + get := doReqAs(t, g, "sk-admin", "GET", "/api/keys/sk-user/auto", "") + if get.Code != http.StatusOK { + t.Fatal(get.Body.String()) + } + if err := json.Unmarshal(get.Body.Bytes(), &got); err != nil { + t.Fatal(err) + } + if !got.Inherits { + t.Fatal("a rejected PUT must not persist a chain") + } +} + +// The other key endpoints must keep working: the {key}/auto route must not +// shadow /api/keys/{key} itself. +func TestKeyAutoAPIDoesNotShadowKeyRoutes(t *testing.T) { + g := keyAutoGateway(t) + + if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys", ""); rr.Code != http.StatusOK { + t.Fatalf("GET /api/keys = %d, want 200", rr.Code) + } + if rr := doReqAs(t, g, "sk-admin", "GET", "/api/keys/me", ""); rr.Code != http.StatusOK { + t.Fatalf("GET /api/keys/me = %d, want 200", rr.Code) + } + if rr := doReqAs(t, g, "sk-user", "GET", "/api/keys/me", ""); rr.Code != http.StatusOK { + t.Fatalf("GET /api/keys/me as user = %d, want 200", rr.Code) + } +} \ No newline at end of file diff --git a/internal/gateway/keys.go b/internal/gateway/keys.go index e1a4c10..587f421 100644 --- a/internal/gateway/keys.go +++ b/internal/gateway/keys.go @@ -7,6 +7,7 @@ import ( "strings" "llmsproxy/internal/config" + "llmsproxy/internal/scheduler" ) // handleKeysAPI manages gateway keys: GET /api/keys (admin: all keys), @@ -17,10 +18,22 @@ func (g *Gateway) handleKeysAPI(w http.ResponseWriter, r *http.Request) { path = strings.Trim(path, "/") role := reqRole(r.Context()) + // /api/keys/me/auto — a user's own chain, readable without admin rights. + // Must be matched before the generic {key}/auto case below so it does not + // ask admin rights of a plain user asking about their own chain. + if path == "me/auto" { + g.handleKeyMeAuto(w, r) + return + } if path == "me" { g.handleKeyMe(w, r) return } + // /api/keys/{key}/auto — the per-key AUTO chain editor. + if i := strings.LastIndex(path, "/auto"); i > 0 && path[i+len("/auto"):] == "" { + g.handleKeyAutoAPI(w, r, path[:i]) + return + } if role != "admin" { writeError(w, http.StatusForbidden, "forbidden", "admin role required") return @@ -139,6 +152,101 @@ func (g *Gateway) allowedModels(ctx context.Context) []config.ModelScope { return rec.Models } +// handleKeyAutoAPI manages one key's own AUTO chain: +// GET /api/keys/{key}/auto returns it, PUT replaces it, DELETE clears the +// override so the key falls back to the global chain. Admin only — a user can +// inspect their own chain through GET /api/keys/me/auto. +// +// An empty chain and a cleared chain are deliberately the same state: "inherit +// the global chain". That keeps the WebUI's "use global" toggle a plain DELETE +// with no sentinel value to carry through config.yaml. +func (g *Gateway) handleKeyAutoAPI(w http.ResponseWriter, r *http.Request, key string) { + if reqRole(r.Context()) != "admin" { + writeError(w, http.StatusForbidden, "forbidden", "admin role required") + return + } + rec, ok := g.core.FindKey(key) + if !ok { + writeError(w, http.StatusNotFound, "not_found", "key not found") + return + } + switch r.Method { + case http.MethodGet: + writeJSON(w, http.StatusOK, map[string]interface{}{ + "key": key, + "auto": rec.Auto, + "inherits": !rec.HasOwnAuto(), + }) + case http.MethodPut, http.MethodPatch: + var body struct { + Auto []config.ModelScope `json:"auto"` + } + if err := json.NewDecoder(r.Body).Decode(&body); err != nil { + writeError(w, http.StatusBadRequest, "invalid_request", "invalid json: "+err.Error()) + return + } + if len(body.Auto) == 0 { + // Treat an empty PUT as "give up the override" so the endpoint + // cannot persist a chain that would 503 with no slots. + if err := g.core.SaveKeyAuto(key, nil); err != nil { + writeError(w, http.StatusBadRequest, "key_error", err.Error()) + return + } + writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "inherits": true}) + return + } + if err := g.core.SaveKeyAuto(key, body.Auto); err != nil { + writeError(w, http.StatusBadRequest, "key_error", err.Error()) + return + } + // Report whether the chain actually resolved to usable slots. An + // admin who typed a model that no longer exists should learn it here, + // not from the user's next 503. + chain, _ := g.core.AutoChainFor(key) + writeJSON(w, http.StatusOK, map[string]interface{}{ + "ok": true, "inherits": false, "slots": chainSlots(chain), + }) + case http.MethodDelete: + if err := g.core.SaveKeyAuto(key, nil); err != nil { + writeError(w, http.StatusBadRequest, "key_error", err.Error()) + return + } + writeJSON(w, http.StatusOK, map[string]interface{}{"ok": true, "inherits": true}) + default: + writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "") + } +} + +// chainSlots counts the usable slots in a chain, for API feedback after an +// edit. A chain that compiled to zero slots is reported so the caller can warn +// instead of letting the next request fail with no_provider. +func chainSlots(ch *scheduler.Chain) int { + if ch == nil { + return 0 + } + n := 0 + for _, tn := range ch.Tiers { + n += len(tn.Slots) + } + return n +} + +// handleKeyMeAuto lets a user read their own AUTO chain without admin rights. +func (g *Gateway) handleKeyMeAuto(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodGet { + writeError(w, http.StatusMethodNotAllowed, "method_not_allowed", "use GET") + return + } + rec, ok := g.core.FindKey(reqKey(r.Context())) + if !ok { + writeError(w, http.StatusUnauthorized, "invalid_api_key", "key not found") + return + } + writeJSON(w, http.StatusOK, map[string]interface{}{ + "key": rec.Key, "auto": rec.Auto, "inherits": !rec.HasOwnAuto(), + }) +} + // handleAutoAPI manages the AUTO scheduling slots: GET /api/auto returns the // current rules; PUT /api/auto replaces them (admin only). func (g *Gateway) handleAutoAPI(w http.ResponseWriter, r *http.Request) { diff --git a/internal/gateway/ui/index.html b/internal/gateway/ui/index.html index ed2c656..5c8c740 100644 --- a/internal/gateway/ui/index.html +++ b/internal/gateway/ui/index.html @@ -822,6 +822,17 @@ keysHint: "管理员密钥可查看与管理全部密钥，并可为每个用户密钥配置可用模型范围；用户密钥只能看到自己。", kCreate: "新建密钥", + kAutoChain: "AUTO 链", + kAutoChainOwn: "独立链", + kAutoChainGlobal: "跟随全局", + kAutoChainTitle: "独立 AUTO 链", + kAutoChainHint: "为空则跟随全局 AUTO 链", + kAutoChainInherit: "使用全局链", + kAutoChainAdd: "+ 添加槽位", + kAutoChainSave: "保存", + kAutoChainTier: "层级", + kAutoChainEmpty: "该密钥使用全局 AUTO 链", + kAutoChainBad: "槽位无法解析（模型或源不存在），保存后 AUTO 请求会失败", kName: "名称", kRole: "角色", kRoleAdmin: "管理员", @@ -1081,6 +1092,17 @@ keysHint: "Admin keys can view and manage every key and configure each user key\u0027s allowed models; user keys only see themselves.", kCreate: "Create key", + kAutoChain: "AUTO chain", + kAutoChainOwn: "own", + kAutoChainGlobal: "global", + kAutoChainTitle: "Per-key AUTO chain", + kAutoChainHint: "Leave empty to inherit the global AUTO chain", + kAutoChainInherit: "Use global chain", + kAutoChainAdd: "+ Add slot", + kAutoChainSave: "Save", + kAutoChainTier: "Tier", + kAutoChainEmpty: "This key uses the global AUTO chain", + kAutoChainBad: "Slot does not resolve (unknown model or source); AUTO requests would fail", kName: "Name", kRole: "Role", kRoleAdmin: "Admin", @@ -4674,6 +4696,8 @@ ${fmtCreated(k.created_at)} + @@ -5085,6 +5109,108 @@ document.body.appendChild(wrap); $("#kc-name").focus(); } + // Per-key AUTO chain editor. Kept deliberately simple: a plain table of + // slots rather than a drag canvas, because an admin usually either + // mirrors the global chain or deletes a few tiers, and the canvas layout + // for the model scope list would be overkill here. + async function openKeyAutoModal(key, name) { + const wrap = document.createElement("div"); + wrap.id = "modal-wrap"; + wrap.style.cssText = + "position:fixed;inset:0;background:rgba(15,22,44,.45);display:flex;align-items:flex-start;justify-content:center;overflow:auto;padding:48px 20px;z-index:50"; + wrap.innerHTML = `
+

${esc(t("kAutoChainTitle"))} · ${esc(name || key)}

+
${esc(t("kAutoChainHint"))}
+
+

+

+

+

+
`; + document.body.appendChild(wrap); + try { + const j = await api("/api/keys/" + encodeURIComponent(key) + "/auto"); + const rows = $("#ka-rows"); + if (!j.auto || !j.auto.length) { + rows.innerHTML = `
${esc( + t("kAutoChainEmpty") + )}
`; + } else { + j.auto.forEach((r) => keyAutoRow(r)); + } + } catch (e) { + toast(String(e)); + } + } + function keyAutoRow(r) { + const el = document.createElement("div"); + el.className = "ka-row"; + el.style.cssText = + "display:flex;gap:8px;align-items:center;margin-bottom:8px;flex-wrap:wrap"; + const opts = (sel) => + (allModels || []) + .map((p) => { + // loadModelPairs builds {key,label}; the /api/status fallback + // only guarantees an id, so derive a comb key from whatever is + // present rather than emitting an empty option value. + const val = p.key || (p.id ? p.id + (p.source ? "|" + p.source : "") : ""); + const lab = p.label || p.id || val; + return ``; + }) + .join(""); + el.innerHTML = ` + + `; + $("#ka-rows").appendChild(el); + } + function keyAutoAddRow() { + keyAutoRow(null); + } + function keyAutoClear() { + $("#ka-rows").innerHTML = ""; + } + async function keyAutoSave(btn, key) { + const rules = []; + document.querySelectorAll("#ka-rows .ka-row").forEach((row) => { + const comb = row.querySelector(".ka-model").value; + if (!comb) return; + const { model, source } = splitCombKey(comb); + const tier = parseInt(row.querySelector(".ka-tier").value, 10) || 1; + rules.push({ model, source, tier }); + }); + btn.disabled = true; + try { + // An empty rule set means "inherit global" server-side, so one PUT + // covers both clearing the override and saving a chain. + const j = await api("/api/keys/" + encodeURIComponent(key) + "/auto", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ auto: rules }), + }); + if (j.inherits) { + toast(t("kAutoChainGlobal")); + } else if (!j.slots) { + // Persisted, but nothing resolves — say so instead of letting the + // user discover it as a 503 on their next AUTO request. + toast(t("kAutoChainBad")); + } else { + toast(t("kAutoChainSave") + " (" + j.slots + ")"); + } + document.getElementById("modal-wrap").remove(); + loadKeys(); + } catch (e) { + toast(String(e)); + btn.disabled = false; + } + } async function createKey(btn) { const name = $("#kc-name").value.trim(); if (!name) { diff --git a/internal/gateway/ui_key_auto_test.go b/internal/gateway/ui_key_auto_test.go new file mode 100644 index 0000000..fd02d35 --- /dev/null +++ b/internal/gateway/ui_key_auto_test.go @@ -0,0 +1,78 @@ +package gateway + +import ( + "strings" + "testing" +) + +// Per-key AUTO chain UI contract. The feature spans three layers (API handler, +// config field, WebUI) and the WebUI part is plain inline JS, so a rename or a +// dropped button fails silently at runtime: the admin simply finds no way to +// configure a key's chain and has no error to chase. +// +// These checks pin the wiring that is easy to break by refactor. + +func TestUIPerKeyAutoWiring(t *testing.T) { + src := uiSource(t) + + // The admin key list must offer the editor, otherwise the whole feature is + // unreachable from the UI. Assert the *call site* inside keyCanvasHtml, + // not just that the functions exist: removing the button leaves every + // definition intact, and that is exactly the regression to catch. + kcStart := strings.Index(src, "function keyCanvasHtml") + if kcStart < 0 { + t.Fatal("keyCanvasHtml is gone — cannot check the AUTO chain button") + } + kcEnd := strings.Index(src[kcStart:], "\n }") + if kcEnd < 0 { + t.Fatal("could not delimit keyCanvasHtml") + } + kcBody := src[kcStart : kcStart+kcEnd] + if !strings.Contains(kcBody, "openKeyAutoModal(") { + t.Fatal("per-key AUTO chain button is missing from the key card header") + } + for _, fn := range []string{ + "function openKeyAutoModal", + "function keyAutoRow", + "function keyAutoAddRow", + "function keyAutoClear", + "function keyAutoSave", + } { + if !strings.Contains(src, fn) { + t.Fatalf("UI is missing %s — the editor cannot function", fn) + } + } + + // The editor must talk to the endpoint the gateway actually serves. + if !strings.Contains(src, "/auto\"") { + t.Fatal("UI never calls the /api/keys/{key}/auto endpoint") + } +} + +// The editor's empty-row path sends {"auto":[]} and relies on the server +// treating that as "clear the override". If the UI instead sent a sentinel, +// a future server change would silently persist a broken empty chain. +func TestUIPerKeyAutoSendsEmptyListToClear(t *testing.T) { + src := uiSource(t) + if !strings.Contains(src, "JSON.stringify({ auto: rules })") { + t.Fatal("keyAutoSave must PUT the collected rule list verbatim") + } + if !strings.Contains(src, "keyAutoClear") { + t.Fatal("no way to clear a key's own chain from the UI") + } +} + +// Both languages need the strings, otherwise a non-Chinese admin sees raw +// i18n keys in the modal. +func TestUIPerKeyAutoHasBothLanguages(t *testing.T) { + src := uiSource(t) + for _, k := range []string{ + "kAutoChainTitle", "kAutoChainHint", "kAutoChainInherit", + "kAutoChainAdd", "kAutoChainSave", "kAutoChainEmpty", + } { + n := strings.Count(src, k+`: "`) + if n < 2 { + t.Fatalf("i18n key %s defined %d time(s), want both zh and en", k, n) + } + } +} \ No newline at end of file