mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-10-03 23:54:06 +00:00
问题:密钥控制只能限制模型范围。实测发现三个缺陷,其中前两个让
per-model token_quota 在真实链路上从未生效:
1. 桶键不含 key。scopeTokens 调 WindowTokens(model, source, win),
桶键是 model / source::model,与调用方无关。实测两把 key 各用
1000 token,窗口报 2000 —— A key 的额度被 B key 消耗。
2. 无 source pin 的桶永远是空的。真实记录 Source 总被填上,桶键存成
"deepseek::m1",而无 pin 的查询找 "m1" —— 读到 0,永远 < quota,
配额形同虚设。实测 WindowTokens("m1","",1h)=0 而 pinned=2000。
3. AUTO scope 走 KeyTokens(key),是全时段累计、永不重置。实测 30 天
前的 200 token 仍计入 1 小时配额(报 210 而非 10)。配了
period: hour 也不会每小时归零。
生产 5 把 user key 全是 token_quota: 0,所以前两条一直没暴露。
改动:
- Stats 新增 per-key 小时桶 keyModelHour(key → model → hour)与
keyHour(key 总量)、keyReqHour(请求数),retention 40 天,与既有
modelHour 对齐以覆盖最长的 month 窗口;LoadAudit 走 aggregateLocked,
所以窗口用量跨重启存活。modelHour 保持 key-blind:它服务的是 AUTO
槽位配额(限制整个网关对某槽位的消耗),语义不同,不应被 per-key
改造污染。
- 每个请求写两份模型桶:裸 model 与 source::model。无 pin 的 scope
条目读前者,有 pin 的读后者。
- GWKey 新增 TokenQuota / ReqQuota / Period / Hours:整钥配额,
跨该 key 所有模型共享一份预算;ReqQuota 覆盖持续请求量(源上的
RPM 只管突发)。
- 配额耗尽返回 429 + Retry-After(rate_limit_exceeded),而不是 403:
403 让客户端以为这把 key 永远不能用该模型,直接放弃;429 + 等待
才能在窗口重置后自动恢复。模型越权仍是 403。
- admin key 永不受配额限制 —— 否则操作者会把自己锁在门外。
- 周期词表在写入时校验,拼错的 period 被拒绝而不是静默当成永不过期
(那与操作者输入的意图正好相反)。
- PUT /api/keys 的配额字段是指针:省略=保留原值,显式 0=解除限制。
否则只改模型范围就会悄悄清空预算。
判据 3 个文件 24 例,9 个变异全部被抓:key 隔离、pin 桶缺失、
AUTO 周期、key-blind 退化、429→403、admin 被限、PUT 清空配额、
Validate 失效、pinned 桶缺失。前三个变异最初漏网 —— 判据只测了
Stats 层没测接线,补了走真实 HTTP 的接线层与 API 层判据后抓住。
端到端验证:真实进程 + 加密配置往返,配额字段与 enc:v1 密钥均正常。
801 lines
22 KiB
Go
801 lines
22 KiB
Go
// Package core wires together the Lua VM, provider registry, scheduler and
|
|
// runtime store and exposes management operations (hot reload, adapters,
|
|
// sources) for the web UI and gateway.
|
|
package core
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
|
|
"llmsproxy/internal/config"
|
|
"llmsproxy/internal/lua"
|
|
"llmsproxy/internal/provider"
|
|
"llmsproxy/internal/scheduler"
|
|
)
|
|
|
|
// Core owns the running configuration and adapters. mu guards every read and
|
|
// write of c.cfg (keys / sources / auto rules) from the management API while
|
|
// request paths look keys up concurrently; the AUTO chain itself is swapped
|
|
// atomically and needs no lock.
|
|
type Core struct {
|
|
mu sync.Mutex
|
|
cfg *config.Config
|
|
vm *lua.VM
|
|
store *config.Store
|
|
scheduler *scheduler.Scheduler
|
|
registry *provider.Registry
|
|
autoChain atomic.Pointer[scheduler.Chain] // chat AUTO chain
|
|
autoImageChain atomic.Pointer[scheduler.Chain] // image-generation AUTO chain
|
|
}
|
|
|
|
// New builds the core from a config file plus runtime overlay.
|
|
func New(cfgPath string) (*Core, error) {
|
|
cfg, err := config.Load(cfgPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return NewFromConfig(cfg)
|
|
}
|
|
|
|
// NewFromConfig builds the core from an already-loaded config.
|
|
func NewFromConfig(cfg *config.Config) (*Core, error) {
|
|
c := &Core{cfg: cfg}
|
|
c.vm = lua.NewVM(cfg.AdapterDir)
|
|
if err := c.vm.Start(); err != nil {
|
|
return nil, fmt.Errorf("lua vm: %w", err)
|
|
}
|
|
c.store = config.NewStore(cfg.RuntimeFile)
|
|
// Share one box between the runtime store and config.yaml so a single
|
|
// master.key seals both files. config.Load left the config holding
|
|
// ciphertext (if it was sealed); unseal it now that the box exists.
|
|
cfg.AttachSecretBox(c.store.SecretBox())
|
|
if err := c.store.Load(); err != nil {
|
|
return nil, fmt.Errorf("runtime store: %w", err)
|
|
}
|
|
c.scheduler = scheduler.New(buildRetries(cfg))
|
|
// One-time migration: lift auto rules and keys from legacy runtime.json
|
|
// into config.yaml so all configuration lives in one place.
|
|
c.migrateFromRuntime()
|
|
if err := c.seedKeys(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := c.seedAuto(); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := c.rebuildRegistry(); err != nil {
|
|
return nil, err
|
|
}
|
|
// Seed built-in preset templates on first run so the WebUI shows
|
|
// "从模板创建" for popular providers out of the box.
|
|
if err := c.seedPresetTemplates(); err != nil {
|
|
return nil, fmt.Errorf("seed preset templates: %w", err)
|
|
}
|
|
// Seal any credential still in the clear in config.yaml. Idempotent: an
|
|
// already-sealed config is not rewritten, so a normal restart writes
|
|
// nothing. This is the only place that rewrites the file on startup.
|
|
if err := cfg.NormalizeSecretsForRun(c.store.SecretBox()); err != nil {
|
|
return nil, fmt.Errorf("normalize secrets: %w", err)
|
|
}
|
|
return c, nil
|
|
}
|
|
|
|
// migrateFromRuntime lifts auto rules and keys from legacy runtime.json into
|
|
// c.cfg — but only if they are not already present in the YAML config. This
|
|
// lets people upgrade without losing their data; once migrated, the YAML file
|
|
// is authoritative and runtime.json's auto/keys are ignored.
|
|
func (c *Core) migrateFromRuntime() {
|
|
legacy := c.store.LoadLegacy()
|
|
if legacy == nil {
|
|
return
|
|
}
|
|
changed := false
|
|
if len(legacy.Auto) > 0 && len(c.cfg.Auto) == 0 {
|
|
c.cfg.Auto = legacy.Auto
|
|
changed = true
|
|
}
|
|
if len(legacy.Keys) > 0 && len(c.cfg.Keys) == 0 {
|
|
c.cfg.Keys = legacy.Keys
|
|
changed = true
|
|
}
|
|
if changed {
|
|
if err := c.cfg.Save(); err != nil {
|
|
fmt.Printf("[core] migrate to config.yaml: %v\n", err)
|
|
}
|
|
}
|
|
}
|
|
|
|
// seedAuto uses the existing auto rules from config.yaml, or creates an
|
|
// initial chain from legacy per-source model priority (first run only).
|
|
func (c *Core) seedAuto() error {
|
|
if len(c.cfg.Auto) > 0 {
|
|
return nil
|
|
}
|
|
type item struct {
|
|
model string
|
|
prio int
|
|
}
|
|
var flat []item
|
|
for _, s := range c.mergedSources() {
|
|
for _, m := range s.Models {
|
|
if m.Kind == "image" {
|
|
continue
|
|
}
|
|
flat = append(flat, item{m.ID, m.Priority})
|
|
}
|
|
}
|
|
sort.SliceStable(flat, func(i, j int) bool {
|
|
if flat[i].prio != flat[j].prio {
|
|
return flat[i].prio > flat[j].prio
|
|
}
|
|
return flat[i].model < flat[j].model
|
|
})
|
|
entries := make([]config.ModelScope, 0, len(flat))
|
|
for _, it := range flat {
|
|
entries = append(entries, config.ModelScope{Model: it.model})
|
|
}
|
|
c.cfg.Auto = entries
|
|
return c.cfg.Save()
|
|
}
|
|
|
|
// seedKeys ensures config.yaml has admin keys for every gateway_keys entry
|
|
// from the YAML config (once). After first save they become GWKey records
|
|
// in c.cfg.Keys and gateway_keys is no longer read for auth.
|
|
func (c *Core) seedKeys() error {
|
|
existing := map[string]bool{}
|
|
for _, k := range c.cfg.Keys {
|
|
existing[k.Key] = true
|
|
}
|
|
changed := false
|
|
for i, raw := range c.cfg.GatewayKeys {
|
|
if raw == "" || existing[raw] {
|
|
continue
|
|
}
|
|
name := "admin"
|
|
if i > 0 {
|
|
name = fmt.Sprintf("admin-%d", i+1)
|
|
}
|
|
c.cfg.Keys = append(c.cfg.Keys, config.GWKey{
|
|
Key: raw,
|
|
Role: "admin",
|
|
Name: name,
|
|
CreatedAt: time.Now().Unix(),
|
|
Seed: true,
|
|
})
|
|
changed = true
|
|
}
|
|
if changed {
|
|
return c.cfg.Save()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// seedPresetTemplates adds the built-in provider presets to the template
|
|
// store on first run. It only adds templates whose name is not already
|
|
// present, so a user's edits to a preset survive restarts and a deliberately
|
|
// deleted preset does not silently return on the next start (its name is
|
|
// remembered in the store's preset marker list).
|
|
func (c *Core) seedPresetTemplates() error {
|
|
if len(config.PresetTemplates) == 0 {
|
|
return nil
|
|
}
|
|
return c.store.SeedTemplates(config.PresetTemplates)
|
|
}
|
|
|
|
// saveConfig writes the current config (including auto rules and keys) back
|
|
// to config.yaml.
|
|
func (c *Core) saveConfig() error {
|
|
return c.cfg.Save()
|
|
}
|
|
|
|
func buildRetries(cfg *config.Config) int {
|
|
return len(cfg.Sources)
|
|
}
|
|
|
|
// VM exposes the Lua adapter runtime.
|
|
func (c *Core) VM() *lua.VM { return c.vm }
|
|
|
|
func (c *Core) Scheduler() *scheduler.Scheduler { return c.scheduler }
|
|
|
|
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() }
|
|
|
|
// AutoImageChain returns the persisted image-generation AUTO chain snapshot.
|
|
func (c *Core) AutoImageChain() *scheduler.Chain { return c.autoImageChain.Load() }
|
|
|
|
func (c *Core) DefaultModel() string { return c.cfg.DefaultModel }
|
|
|
|
func (c *Core) GatewayKeys() []string { return c.cfg.GatewayKeys }
|
|
|
|
func (c *Core) Listen() string { return c.cfg.Listen }
|
|
|
|
func (c *Core) TLS() (cert, key string) {
|
|
return c.cfg.TLSCertFile, c.cfg.TLSKeyFile
|
|
}
|
|
|
|
func (c *Core) PublicBaseURL() string { return c.cfg.PublicBaseURL }
|
|
|
|
// ---- gateway key management (web UI) ----
|
|
|
|
// ListKeys returns all gateway keys (admin view).
|
|
func (c *Core) ListKeys() []config.GWKey {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
out := make([]config.GWKey, len(c.cfg.Keys))
|
|
copy(out, c.cfg.Keys)
|
|
return out
|
|
}
|
|
|
|
// FindKey looks up a gateway key record by its secret value.
|
|
func (c *Core) FindKey(key string) (config.GWKey, bool) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
for _, k := range c.cfg.Keys {
|
|
if k.Key == key {
|
|
return k, true
|
|
}
|
|
}
|
|
return config.GWKey{}, false
|
|
}
|
|
|
|
// CreateKey builds a new random gateway key and persists it to config.yaml.
|
|
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
|
return c.CreateKeyWithQuota(name, role, models, note, config.KeyQuota{})
|
|
}
|
|
|
|
// CreateKeyWithQuota is CreateKey plus the key-wide token/request caps.
|
|
func (c *Core) CreateKeyWithQuota(name, role string, models []config.ModelScope, note string, q config.KeyQuota) (config.GWKey, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
models = cleanScopes(models)
|
|
key := make([]byte, 16)
|
|
if _, err := rand.Read(key); err != nil {
|
|
return config.GWKey{}, err
|
|
}
|
|
if err := q.Validate(); err != nil {
|
|
return config.GWKey{}, err
|
|
}
|
|
rec := config.GWKey{
|
|
Key: "sk-gw-" + hex.EncodeToString(key),
|
|
Role: role,
|
|
Name: name,
|
|
Models: models,
|
|
Note: note,
|
|
CreatedAt: time.Now().Unix(),
|
|
}
|
|
rec.Role = config.NormalizeRole(rec.Role)
|
|
rec.ApplyQuota(q)
|
|
c.cfg.Keys = append(c.cfg.Keys, rec)
|
|
if err := c.saveConfig(); err != nil {
|
|
return config.GWKey{}, err
|
|
}
|
|
return rec, nil
|
|
}
|
|
|
|
// UpdateKey mutates a key's name/role/model scope and persists it.
|
|
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
|
|
return c.UpdateKeyWithQuota(key, name, role, models, note, nil)
|
|
}
|
|
|
|
// UpdateKeyWithQuota is UpdateKey plus the key-wide caps. quota == nil leaves
|
|
// the existing caps untouched, so a caller that only edits the model scope
|
|
// does not silently clear a key's budget.
|
|
func (c *Core) UpdateKeyWithQuota(key, name, role string, models []config.ModelScope, note string, quota *config.KeyQuota) (config.GWKey, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if quota != nil {
|
|
if err := quota.Validate(); err != nil {
|
|
return config.GWKey{}, err
|
|
}
|
|
}
|
|
for i, k := range c.cfg.Keys {
|
|
if k.Key == key {
|
|
if name != "" {
|
|
c.cfg.Keys[i].Name = name
|
|
}
|
|
if role == "admin" || role == "user" {
|
|
c.cfg.Keys[i].Role = role
|
|
}
|
|
// models == nil means the caller did not provide a scope (leave
|
|
// the existing one untouched); an explicit [] clears it.
|
|
if models != nil {
|
|
c.cfg.Keys[i].Models = cleanScopes(models)
|
|
}
|
|
c.cfg.Keys[i].Note = note
|
|
if quota != nil {
|
|
c.cfg.Keys[i].ApplyQuota(*quota)
|
|
}
|
|
if err := c.saveConfig(); err != nil {
|
|
return config.GWKey{}, err
|
|
}
|
|
return c.cfg.Keys[i], nil
|
|
}
|
|
}
|
|
return config.GWKey{}, fmt.Errorf("key not found")
|
|
}
|
|
|
|
// DeleteKey removes a key record; returns false if it did not exist.
|
|
func (c *Core) DeleteKey(key string) (bool, error) {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
for i, k := range c.cfg.Keys {
|
|
if k.Key == key {
|
|
c.cfg.Keys = append(c.cfg.Keys[:i], c.cfg.Keys[i+1:]...)
|
|
return true, c.saveConfig()
|
|
}
|
|
}
|
|
return false, nil
|
|
}
|
|
|
|
// ---- AUTO scheduling slots (web UI) ----
|
|
|
|
// AutoRules returns the AUTO scheduling slots in priority order.
|
|
func (c *Core) AutoRules() []config.ModelScope {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
out := make([]config.ModelScope, len(c.cfg.Auto))
|
|
copy(out, c.cfg.Auto)
|
|
return out
|
|
}
|
|
|
|
// AutoImageRules returns the configured image-generation AUTO chain rules.
|
|
func (c *Core) AutoImageRules() []config.ModelScope {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
out := make([]config.ModelScope, len(c.cfg.AutoImage))
|
|
copy(out, c.cfg.AutoImage)
|
|
return out
|
|
}
|
|
|
|
// cleanScopes drops empty model entries and normalizes placeholder source
|
|
// names; it returns nil when no entries survive so an empty scope means
|
|
// "unrestricted" (nil) instead of a restrictive-but-empty list — a non-nil
|
|
// empty slice would 403 every model in-process yet become unrestricted again
|
|
// after a restart (config omits empty models with omitempty).
|
|
func cleanScopes(entries []config.ModelScope) []config.ModelScope {
|
|
var clean []config.ModelScope
|
|
for _, e := range entries {
|
|
if e.Model == "" {
|
|
continue
|
|
}
|
|
if e.Source == "undefined" || e.Source == "null" {
|
|
e.Source = ""
|
|
}
|
|
clean = append(clean, e)
|
|
}
|
|
return clean
|
|
}
|
|
|
|
// SaveAutoRules persists the AUTO scheduling slots to config.yaml, rebuilds
|
|
// the chain and clears the cooldown of every slot — preference scores are
|
|
// kept, so a reliably good model keeps its edge while an edited chain applies
|
|
// immediately.
|
|
func (c *Core) SaveAutoRules(entries []config.ModelScope) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.cfg.Auto = cleanScopes(entries)
|
|
if err := c.saveConfig(); err != nil {
|
|
return err
|
|
}
|
|
c.buildAutoChain()
|
|
c.resetChainCooldowns(c.autoChain.Load())
|
|
return nil
|
|
}
|
|
|
|
// SaveAutoImageRules persists the image-generation AUTO chain. Image models
|
|
// are kept as-is (unlike buildAutoChain, which skips them).
|
|
func (c *Core) SaveAutoImageRules(entries []config.ModelScope) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
c.cfg.AutoImage = cleanScopes(entries)
|
|
if err := c.saveConfig(); err != nil {
|
|
return err
|
|
}
|
|
c.buildAutoImageChain()
|
|
c.resetChainCooldowns(c.autoImageChain.Load())
|
|
return nil
|
|
}
|
|
|
|
func (c *Core) resetChainCooldowns(ch *scheduler.Chain) {
|
|
if ch == nil {
|
|
return
|
|
}
|
|
for _, tn := range ch.Tiers {
|
|
for _, sl := range tn.Slots {
|
|
if p := c.registry.ProviderForSlot(sl.Model, sl.Source); p != nil {
|
|
p.ResetModelCooldown(sl.Model)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ResetHealth clears the scheduling backoff state of every provider.
|
|
func (c *Core) ResetHealth() {
|
|
for _, p := range c.registry.Providers() {
|
|
p.ResetHealth()
|
|
}
|
|
}
|
|
|
|
func (c *Core) ProviderForModel(model string) *provider.Provider {
|
|
return c.registry.ProviderForModel(model)
|
|
}
|
|
|
|
func (c *Core) ProviderForSlot(model, source string) *provider.Provider {
|
|
return c.registry.ProviderForSlot(model, source)
|
|
}
|
|
|
|
func (c *Core) Config() *config.Config { return c.cfg }
|
|
|
|
// mergedSources = base YAML sources + runtime sources (runtime wins by name).
|
|
func (c *Core) mergedSources() []config.Source {
|
|
byName := map[string]config.Source{}
|
|
order := []string{}
|
|
for _, s := range c.cfg.Sources {
|
|
byName[s.Name] = s
|
|
order = append(order, s.Name)
|
|
}
|
|
for _, s := range c.store.List() {
|
|
if _, ok := byName[s.Name]; !ok {
|
|
order = append(order, s.Name)
|
|
}
|
|
byName[s.Name] = s
|
|
}
|
|
out := make([]config.Source, 0, len(order))
|
|
seen := map[string]bool{}
|
|
for _, n := range order {
|
|
if !seen[n] {
|
|
seen[n] = true
|
|
s := c.resolveSourceKey(byName[n])
|
|
if s.Timeout == 0 {
|
|
s.Timeout = config.DefaultSourceTimeout
|
|
}
|
|
if s.QueueTimeout == 0 {
|
|
s.QueueTimeout = config.DefaultSourceQueueTimeout
|
|
}
|
|
if s.MaxConcurrent == 0 {
|
|
s.MaxConcurrent = config.DefaultSourceConcurrency
|
|
}
|
|
out = append(out, s)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// resolveSourceKey applies api_key_env and decrypts enc:v1: ciphertext.
|
|
// Config values are already unsealed at startup (Config.NormalizeSecretsForRun),
|
|
// so the decrypt branch is the belt-and-braces path for a source that arrived
|
|
// already sealed through another route. A failure there leaves the ciphertext in
|
|
// place, which makes the upstream reject the key loudly rather than sending an
|
|
// empty Authorization header that might look like a config-less source.
|
|
func (c *Core) resolveSourceKey(s config.Source) config.Source {
|
|
if s.APIKeyEnv != "" {
|
|
if v := os.Getenv(s.APIKeyEnv); v != "" {
|
|
s.APIKey = v
|
|
}
|
|
return s
|
|
}
|
|
if box := c.store.SecretBox(); box != nil && strings.HasPrefix(s.APIKey, "enc:v1:") {
|
|
if v, err := box.Decrypt(s.APIKey); err == nil {
|
|
s.APIKey = v
|
|
} else {
|
|
log.Printf("[core] source %s: api_key decrypt failed: %v", s.Name, err)
|
|
}
|
|
}
|
|
return s
|
|
}
|
|
|
|
func (c *Core) rebuildRegistry() error {
|
|
srcs := c.mergedSources()
|
|
providers := make([]*provider.Provider, 0, len(srcs))
|
|
adapterConcurrency := map[string]int{}
|
|
for _, s := range srcs {
|
|
providers = append(providers, provider.New(s, c.vm))
|
|
if s.Adapter != "" {
|
|
adapterConcurrency[s.Adapter] += s.MaxConcurrent
|
|
}
|
|
}
|
|
c.vm.ConfigureConcurrency(adapterConcurrency)
|
|
if c.registry == nil {
|
|
c.registry = provider.NewRegistry(providers, c.cfg.DefaultModel)
|
|
} else {
|
|
c.registry.Replace(providers)
|
|
}
|
|
c.buildAutoChain()
|
|
c.buildAutoImageChain()
|
|
return nil
|
|
}
|
|
|
|
// buildAutoChain rebuilds the AUTO chain snapshot from config.yaml rules
|
|
// against the current providers. Slot model ids are normalized to the exact
|
|
// configured spelling (case-insensitive match), otherwise ModelFor would fall
|
|
// back to the source's best chat model and cooldown/quota bookkeeping would
|
|
// key on a name that never matches.
|
|
func (c *Core) buildAutoChain() {
|
|
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" {
|
|
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
|
|
if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
|
|
if exact := p.ModelIDFold(e.Model); exact != "" {
|
|
model = exact
|
|
}
|
|
if m := p.ModelByID(model); m != nil && m.Kind == "image" {
|
|
continue // image-kind models never join the chat AUTO chain
|
|
}
|
|
if source == "" {
|
|
source = p.Name()
|
|
}
|
|
}
|
|
sr = append(sr, scheduler.Rule{
|
|
Model: model,
|
|
Source: source,
|
|
Tier: e.Tier,
|
|
Quota: e.TokenQuota,
|
|
Period: e.Period,
|
|
Hours: e.Hours,
|
|
})
|
|
}
|
|
c.autoChain.Store(scheduler.BuildChain(sr, prov))
|
|
}
|
|
|
|
// buildAutoImageChain rebuilds the image-generation AUTO chain snapshot.
|
|
// Unlike buildAutoChain, only image-kind models participate: chat models in
|
|
// the rules are skipped so a stale chat slot can't receive image traffic.
|
|
func (c *Core) buildAutoImageChain() {
|
|
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" {
|
|
return nil
|
|
}
|
|
return p
|
|
}
|
|
rules := c.cfg.AutoImage
|
|
sr := make([]scheduler.Rule, 0, len(rules))
|
|
for _, e := range rules {
|
|
model, source := e.Model, e.Source
|
|
if p := c.registry.ProviderForSlot(e.Model, e.Source); p != nil {
|
|
if exact := p.ModelIDFold(e.Model); exact != "" {
|
|
model = exact
|
|
}
|
|
if m := p.ModelByID(model); m != nil && m.Kind != "image" {
|
|
continue // chat-kind models never join the image AUTO chain
|
|
}
|
|
if source == "" {
|
|
source = p.Name()
|
|
}
|
|
}
|
|
sr = append(sr, scheduler.Rule{
|
|
Model: model,
|
|
Source: source,
|
|
Tier: e.Tier,
|
|
Quota: e.TokenQuota,
|
|
Period: e.Period,
|
|
Hours: e.Hours,
|
|
})
|
|
}
|
|
c.autoImageChain.Store(scheduler.BuildChain(sr, prov))
|
|
}
|
|
|
|
// AutoSlotState is the UI-facing health snapshot of one AUTO chain slot.
|
|
type AutoSlotState struct {
|
|
Model string `json:"model"`
|
|
Source string `json:"source"`
|
|
Pref int64 `json:"pref"`
|
|
FailCount int64 `json:"fail_count"`
|
|
CooldownUntil int64 `json:"cooldown_until"`
|
|
Cooling bool `json:"cooling"`
|
|
// CooldownFrom is when the current cooldown window opened and ProbeAfter is
|
|
// its midpoint: from ProbeAfter on, a single probe request is allowed
|
|
// through so a recovered upstream does not have to sit out the rest of the
|
|
// window. Probing reports whether such a probe is in flight right now.
|
|
CooldownFrom int64 `json:"cooldown_from"`
|
|
ProbeAfter int64 `json:"probe_after"`
|
|
Probing bool `json:"probing"`
|
|
// Probeable marks a cooling slot that is past its midpoint, i.e. the next
|
|
// AUTO request may use it as a probe.
|
|
Probeable bool `json:"probeable"`
|
|
}
|
|
|
|
// AutoSlotStates returns per-slot health for every slot of the current chain.
|
|
func (c *Core) AutoSlotStates() []AutoSlotState {
|
|
ch := c.autoChain.Load()
|
|
if ch == nil {
|
|
return nil
|
|
}
|
|
now := time.Now().Unix()
|
|
var out []AutoSlotState
|
|
for _, tn := range ch.Tiers {
|
|
for _, sl := range tn.Slots {
|
|
pp, ok := sl.Prov.(*provider.Provider)
|
|
if !ok {
|
|
continue
|
|
}
|
|
pref, fail, until := pp.ModelHealthInfo(sl.Model)
|
|
from, probeAfter, probing := pp.ModelProbeInfo(sl.Model)
|
|
out = append(out, AutoSlotState{
|
|
Model: sl.Model,
|
|
Source: sl.Source,
|
|
Pref: pref,
|
|
FailCount: fail,
|
|
CooldownUntil: until,
|
|
Cooling: until > now,
|
|
CooldownFrom: from,
|
|
ProbeAfter: probeAfter,
|
|
Probing: probing,
|
|
Probeable: until > now && probeAfter > 0 && now >= probeAfter,
|
|
})
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Reload re-reads the runtime store and rebuilds sources.
|
|
func (c *Core) Reload() error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if err := c.store.Load(); err != nil {
|
|
return err
|
|
}
|
|
return c.rebuildRegistry()
|
|
}
|
|
|
|
// ---- adapter management (web UI) ----
|
|
|
|
func (c *Core) ListAdapters() []lua.APIAdapter {
|
|
return c.vm.ListAdapters()
|
|
}
|
|
|
|
func (c *Core) UploadAdapter(name, code string) error {
|
|
if name == "" {
|
|
return fmt.Errorf("adapter name required")
|
|
}
|
|
if err := os.MkdirAll(c.cfg.AdapterDir, 0755); err != nil {
|
|
return err
|
|
}
|
|
path := filepath.Join(c.cfg.AdapterDir, name+".lua")
|
|
if err := os.WriteFile(path, []byte(code), 0644); err != nil {
|
|
return err
|
|
}
|
|
return c.vm.LoadAdapter(path)
|
|
}
|
|
|
|
func (c *Core) RemoveAdapter(name string) error {
|
|
path := filepath.Join(c.cfg.AdapterDir, name+".lua")
|
|
_ = os.Remove(path)
|
|
c.vm.RemoveAdapter(name)
|
|
return nil
|
|
}
|
|
|
|
// ---- source management (web UI) ----
|
|
|
|
func (c *Core) AddSource(src config.Source) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if err := normalizeSource(&src); err != nil {
|
|
return err
|
|
}
|
|
if err := config.UpsertSourceInYAML(c.cfg.Path, src.Name, src, c.cfg.SecretBox()); err != nil {
|
|
return err
|
|
}
|
|
// Update in-memory Sources so mergedSources() finds the entry.
|
|
replaced := false
|
|
for i, s := range c.cfg.Sources {
|
|
if s.Name == src.Name {
|
|
c.cfg.Sources[i] = src
|
|
replaced = true
|
|
break
|
|
}
|
|
}
|
|
if !replaced {
|
|
c.cfg.Sources = append(c.cfg.Sources, src)
|
|
}
|
|
return c.rebuildRegistry()
|
|
}
|
|
|
|
func (c *Core) RemoveSource(name string) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
for _, s := range c.cfg.Sources {
|
|
if s.Name == name {
|
|
if err := config.RemoveSourceFromYAML(c.cfg.Path, name); err != nil {
|
|
return err
|
|
}
|
|
c.cfg.Sources = c.removeCfgSource(name)
|
|
break
|
|
}
|
|
}
|
|
if _, err := c.store.Remove(name); err != nil {
|
|
return err
|
|
}
|
|
return c.rebuildRegistry()
|
|
}
|
|
|
|
func (c *Core) removeCfgSource(name string) []config.Source {
|
|
out := c.cfg.Sources[:0]
|
|
for _, s := range c.cfg.Sources {
|
|
if s.Name != name {
|
|
out = append(out, s)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (c *Core) Sources() []config.Source {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
return c.mergedSources()
|
|
}
|
|
|
|
// Templates returns the saved source templates.
|
|
func (c *Core) Templates() []config.SourceTemplate {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
return c.store.ListTemplates()
|
|
}
|
|
|
|
// SaveTemplate upserts a source template.
|
|
func (c *Core) SaveTemplate(t config.SourceTemplate) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if t.Name == "" {
|
|
return fmt.Errorf("template requires a name")
|
|
}
|
|
return c.store.UpsertTemplate(t)
|
|
}
|
|
|
|
// RemoveTemplate deletes a source template by name.
|
|
func (c *Core) RemoveTemplate(name string) error {
|
|
c.mu.Lock()
|
|
defer c.mu.Unlock()
|
|
if _, err := c.store.RemoveTemplate(name); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func normalizeSource(s *config.Source) error {
|
|
if s.Name == "" || s.BaseURL == "" {
|
|
return fmt.Errorf("source requires name and base_url")
|
|
}
|
|
if len(s.Models) == 0 {
|
|
return fmt.Errorf("source requires at least one model")
|
|
}
|
|
if s.Adapter == "" {
|
|
s.Adapter = "openai"
|
|
}
|
|
if s.MaxConcurrent == 0 {
|
|
s.MaxConcurrent = 8
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// Close releases resources.
|
|
func (c *Core) Close() {
|
|
if c.vm != nil {
|
|
c.vm.Stop()
|
|
}
|
|
}
|