Files
ModelRouter/internal/core/core.go
JianFeeeee ce2032435a fix(plugin): 插件 state 持久化 —— 重启不再丢账
## 问题(压测实测)
插件 state 活在 Lua VM 里,进程一死就没了。实测线上量级:
  重启前 {"requests":218241,"prompt_tokens":26188920,...}
  重启后 {"requests":0,"prompt_tokens":0,...}
对计费插件来说这不是舍入误差,是功能本身没生效 —— 它存在的意义就是那个
不断累加的数字,而一次 systemctl restart 就能把它抹掉。

## 设计
- prices(配置)与 state(累计历史)**分开存**在同一个文件里但不同字段。
  SetState 在内存里已经这么分,磁盘必须同意:合并会让改价看起来像清零,
  或让恢复历史时顺带复活过期价格。
- 原子写(临时文件 + rename):写一半崩掉时上一份仍可读,而不是留下一个
  解析失败的半截 JSON —— 那等于这次也丢。
- 损坏文件只警告不阻断启动。转发不能依赖插件的账本活着。
- 防抖后台刷:钩子路径只标记,真正的写在一个 goroutine 里合并进行。
  计费插件每请求都改 state,同步写会把一次 JSON 编码 + 文件写放到热路径上
  (实测钩子本身已经 14.6µs,写会盖过它)。
- Core.Close 必须先刷插件再停 VM:flush 要读 Lua 表,vm.Stop() 之后读的是
  已释放的内存。

## ★ 实现中踩的四个坑(都由测试或崩溃直接暴露,不是推测)
1. **后台 goroutine 碰 Lua = use-after-free**。最初让 flush 线程去读 Lua 状态,
   vm.Stop() 后那是已释放内存 —— 表现为 golua 里的 SIGSEGV,不是干净报错。
   改成:钩子路径(VM 必然存活、已持 p.mu)取快照,后台只写文件。
2. **自死锁**:markDirtyLocked 被 invoke 调用,而 invoke 全程持 p.mu,
   再 Lock 一次就是死锁。lua 包测试直接挂到超时。
3. **luaToJSON 独占整个栈**(每条路径结尾都 SetTop(0))。连续调两次读两个
   字段时第二次访问的是不存在的槽位 —— 这个绑定不 panic,直接 SIGABRT。
   改为每次重建栈。中间还因为提前 return 没 Pop 而让栈逐次错位。
4. **快照顺序**:先快照后读返回值,会把钩子的返回值清掉,于是每个"有意见"的
   插件静默变成"没意见",而文档承诺的"返回 table 合并进 payload"就废了,
   且没有任何报错。

## 判据(7 项,全部变异验证过)
重启后总计保留 / prices 与 state 分离 / 纯 prices 更新也持久化 /
钩子返回值不被快照吃掉 / 损坏文件降级不阻断 / 500 次变更合并成个位数次写 /
Close 刷出尾部。

变异结果:
  关掉 mark          → TestStateSurvivesRestart + TestPricesAndStateAreSeparate 红
  Close 不等 flush   → TestCloseFlushesTail 红
  prices-only 不写盘 → TestPricesOnlyUpdatePersists 红
  还原快照顺序       → TestHookReturnValueSurvivesSnapshot 红
★ 第一次跑「关掉 mark」时判据没报错,原因是我的变异脚本写出未使用变量导致
  编译失败 —— go test 根本没跑测试,我却读成了"通过"。换成 _, _ = 后如期变红。

## 端到端
隔离实例发 12 次请求 → systemctl restart → requests 仍为 12,token 数不变。
371 个测试全绿。
2026-10-02 10:17:07 +08:00

857 lines
25 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
plugins *lua.Plugins
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)
}
// Plugins load AFTER the VM is up (a plugin is just another Lua state) but
// BEFORE anything can serve traffic, so a plugin's request_end hook sees
// the very first request. A plugin that fails to load is logged and skipped
// rather than failing startup: plugins are optional extensions, and failing
// to boot the gateway because a third-party .lua has a typo would be the
// wrong trade.
c.plugins = lua.NewPlugins(c.vm, cfg.PluginDir)
if err := c.plugins.SeedBundled(); err != nil {
log.Printf("[core] seed bundled plugins: %v", err)
}
if err := c.plugins.LoadDir(); err != nil {
log.Printf("[core] plugin dir: %v", err)
}
for _, p := range c.plugins.List() {
if msg, bad := p["error"]; bad {
log.Printf("[core] plugin %v failed to load: %v", p["name"], msg)
continue
}
log.Printf("[core] plugin %v %v loaded (hooks=%v ui=%v)",
p["name"], p["version"], p["hooks"], p["ui"])
}
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, before anything reads a
// credential.
//
// ORDER IS LOAD-BEARING. These steps each consume secrets and must run
// after the unseal:
// - seedKeys compares cfg.Keys[i].Key against the plaintext
// gateway_keys entries; with ciphertext keys the comparison never
// matched and every restart appended another duplicate admin key
// (production had four copies of the same admin key).
// - rebuildRegistry hands cfg.Sources[i].APIKey to the providers; with
// ciphertext it built every upstream client with "enc:v1:..." as its
// bearer token.
// Previously the unseal happened at the END of this function, and things
// only appeared to work because seedKeys' Save() unsealed memory as a side
// effect. Removing the redundant saves exposed the real ordering bug.
cfg.AttachSecretBox(c.store.SecretBox())
hadPlaintextSecrets, err := cfg.UnsealSecrets(c.store.SecretBox())
if err != nil {
return nil, fmt.Errorf("unseal secrets: %w", err)
}
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. Startup writes
// nothing when the file was already sealed (hadPlaintextSecrets is decided
// from the on-disk state, before the unseal above).
if err := cfg.SealIfNeeded(hadPlaintextSecrets); 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 }
// Plugins exposes the loaded plugin set. Never nil once New* has returned, so
// the gateway can fire stages unconditionally; an unconfigured plugin
// directory yields an empty registry whose Fire is a no-op.
func (c *Core) Plugins() *lua.Plugins { return c.plugins }
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.
// Quotas live on the model scope entries, so a new key's budget is whatever
// its scopes carry.
func (c *Core) CreateKey(name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
models = cleanScopes(models)
if err := ValidateScopeQuotas(models); err != nil {
return config.GWKey{}, err
}
key := make([]byte, 16)
if _, err := rand.Read(key); 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)
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. The scope
// entries carry their own quotas, so replacing the scope replaces the budgets.
func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, note string) (config.GWKey, error) {
c.mu.Lock()
defer c.mu.Unlock()
if models != nil {
if err := ValidateScopeQuotas(models); 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 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() {
// Plugin state must be flushed BEFORE the VM stops. The saver's final write
// reads each plugin's Lua tables; once vm.Stop() has closed those states the
// read finds nothing and the last interval of accumulation is lost — which
// is the exact failure this persistence was added to prevent.
if c.plugins != nil {
c.plugins.Close()
}
if c.vm != nil {
c.vm.Stop()
}
}
// ValidateScopeQuotas checks every scope entry's caps before they are stored.
// A typo in a period must be rejected at write time rather than silently
// becoming a never-resetting budget — the opposite of what was typed.
func ValidateScopeQuotas(entries []config.ModelScope) error {
for _, e := range entries {
if err := (config.KeyQuota{
TokenQuota: e.TokenQuota,
ReqQuota: e.ReqQuota,
Period: e.Period,
Hours: e.Hours,
}).Validate(); err != nil {
return fmt.Errorf("model %q: %w", e.Model, err)
}
}
return nil
}