mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 17:07:59 +00:00
feat: opencode zen adapter + first-run config generation, fix stats/stream bugs
- adapters/opencode.lua: opencode.ai zen free pool adapter — sends the opencode client User-Agent (zen fingerprints clients by UA; non-official clients hit FreeUsageLimitError); pairs with api_key: public - config: no config file ships in the repo; first run generates a default config at the -config path with a random admin key, loopback listen and a keyless zen source (config.EnsureDefault); remove config.example.yaml - lua: seed bundled adapters from the embedded FS instead of a hardcoded name list - ui: widen model kind select (chat was clipped to 'cha') - phase 5 bugfixes: stats ms/s bucket mixing, cleanScopes nil, ctx.Err guards, direct-path ModelAvailable, empty stream body failure, bestImageModel rewrite, transform failure recording, Core.mu, timer, effective model for tool-calls
This commit is contained in:
@ -11,6 +11,7 @@ import (
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
@ -20,8 +21,12 @@ import (
|
||||
"llmsproxy/internal/scheduler"
|
||||
)
|
||||
|
||||
// Core owns the running configuration and adapters.
|
||||
// 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
|
||||
@ -192,6 +197,8 @@ func (c *Core) PublicBaseURL() string { return c.cfg.PublicBaseURL }
|
||||
|
||||
// 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
|
||||
@ -199,6 +206,8 @@ func (c *Core) ListKeys() []config.GWKey {
|
||||
|
||||
// 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
|
||||
@ -209,6 +218,8 @@ func (c *Core) FindKey(key string) (config.GWKey, bool) {
|
||||
|
||||
// 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) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
models = cleanScopes(models)
|
||||
key := make([]byte, 16)
|
||||
if _, err := rand.Read(key); err != nil {
|
||||
@ -234,6 +245,8 @@ func (c *Core) CreateKey(name, role string, models []config.ModelScope, note str
|
||||
|
||||
// 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) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
for i, k := range c.cfg.Keys {
|
||||
if k.Key == key {
|
||||
if name != "" {
|
||||
@ -242,7 +255,11 @@ func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, not
|
||||
if role == "admin" || role == "user" {
|
||||
c.cfg.Keys[i].Role = role
|
||||
}
|
||||
c.cfg.Keys[i].Models = cleanScopes(models)
|
||||
// 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
|
||||
@ -255,6 +272,8 @@ func (c *Core) UpdateKey(key, name, role string, models []config.ModelScope, not
|
||||
|
||||
// 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:]...)
|
||||
@ -268,13 +287,20 @@ func (c *Core) DeleteKey(key string) (bool, error) {
|
||||
|
||||
// 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
|
||||
}
|
||||
|
||||
// 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 {
|
||||
clean := make([]config.ModelScope, 0, len(entries))
|
||||
var clean []config.ModelScope
|
||||
for _, e := range entries {
|
||||
if e.Model == "" {
|
||||
continue
|
||||
@ -292,6 +318,8 @@ func cleanScopes(entries []config.ModelScope) []config.ModelScope {
|
||||
// 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
|
||||
@ -396,7 +424,10 @@ func (c *Core) rebuildRegistry() error {
|
||||
}
|
||||
|
||||
// buildAutoChain rebuilds the AUTO chain snapshot from config.yaml rules
|
||||
// against the current providers.
|
||||
// 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)
|
||||
@ -411,20 +442,26 @@ func (c *Core) buildAutoChain() {
|
||||
rules := c.cfg.Auto
|
||||
sr := make([]scheduler.Rule, 0, len(rules))
|
||||
for _, e := range rules {
|
||||
r := scheduler.Rule{
|
||||
Model: e.Model,
|
||||
Source: e.Source,
|
||||
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,
|
||||
}
|
||||
if r.Source == "" {
|
||||
if p := c.registry.ProviderForSlot(e.Model, ""); p != nil {
|
||||
r.Source = p.Name()
|
||||
}
|
||||
}
|
||||
sr = append(sr, r)
|
||||
})
|
||||
}
|
||||
c.autoChain.Store(scheduler.BuildChain(sr, prov))
|
||||
}
|
||||
@ -469,6 +506,8 @@ func (c *Core) AutoSlotStates() []AutoSlotState {
|
||||
|
||||
// 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
|
||||
}
|
||||
@ -505,6 +544,8 @@ func (c *Core) RemoveAdapter(name string) error {
|
||||
// ---- 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
|
||||
}
|
||||
@ -527,6 +568,8 @@ func (c *Core) AddSource(src config.Source) error {
|
||||
}
|
||||
|
||||
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 {
|
||||
@ -552,7 +595,11 @@ func (c *Core) removeCfgSource(name string) []config.Source {
|
||||
return out
|
||||
}
|
||||
|
||||
func (c *Core) Sources() []config.Source { return c.mergedSources() }
|
||||
func (c *Core) Sources() []config.Source {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
return c.mergedSources()
|
||||
}
|
||||
|
||||
func normalizeSource(s *config.Source) error {
|
||||
if s.Name == "" || s.BaseURL == "" {
|
||||
|
||||
135
internal/core/core_test.go
Normal file
135
internal/core/core_test.go
Normal file
@ -0,0 +1,135 @@
|
||||
package core
|
||||
|
||||
import (
|
||||
"path/filepath"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"llmsproxy/internal/config"
|
||||
)
|
||||
|
||||
// newTestConfig returns an in-memory config backed by temp files so
|
||||
// NewFromConfig (store load, key/auto seeding, registry build) is exercised
|
||||
// like production.
|
||||
func newTestConfig(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-test"},
|
||||
Sources: []config.Source{
|
||||
{Name: "s1", BaseURL: "http://127.0.0.1:1/v1", Adapter: "openai", Models: []config.Model{{ID: "gpt-4o", Priority: 10}}},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func newTestCore(t *testing.T, cfg *config.Config) *Core {
|
||||
t.Helper()
|
||||
c, err := NewFromConfig(cfg)
|
||||
if err != nil {
|
||||
t.Fatalf("core: %v", err)
|
||||
}
|
||||
t.Cleanup(c.Close)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestCleanScopesNilForEmpty(t *testing.T) {
|
||||
if got := cleanScopes(nil); got != nil {
|
||||
t.Fatalf("nil in must yield nil out, got %#v", got)
|
||||
}
|
||||
if got := cleanScopes([]config.ModelScope{{Model: ""}}); got != nil {
|
||||
t.Fatalf("all-empty entries must yield nil, got %#v", got)
|
||||
}
|
||||
got := cleanScopes([]config.ModelScope{{Model: "m", Source: "undefined"}})
|
||||
if len(got) != 1 || got[0].Source != "" || got[0].Model != "m" {
|
||||
t.Fatalf("placeholder normalization wrong: %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCreateKeyWithoutScopeUnrestricted: a user key created without a model
|
||||
// scope must be unrestricted in-process (nil), identical to what config.yaml
|
||||
// omitempty produces after a restart — no 403 lock-out, no silent widening.
|
||||
func TestCreateKeyWithoutScopeUnrestricted(t *testing.T) {
|
||||
c := newTestCore(t, newTestConfig(t))
|
||||
k, err := c.CreateKey("svc", "user", nil, "")
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
rec, ok := c.FindKey(k.Key)
|
||||
if !ok {
|
||||
t.Fatal("key not found")
|
||||
}
|
||||
if rec.Models != nil {
|
||||
t.Fatalf("user key without scope must be unrestricted, got %#v", rec.Models)
|
||||
}
|
||||
}
|
||||
|
||||
// TestUpdateKeyPreservesScopeWhenOmitted: PUT /api/keys without a models field
|
||||
// must keep the existing scope; only an explicit [] clears it.
|
||||
func TestUpdateKeyPreservesScopeWhenOmitted(t *testing.T) {
|
||||
c := newTestCore(t, newTestConfig(t))
|
||||
k, err := c.CreateKey("svc", "user", []config.ModelScope{{Model: "gpt-4o"}}, "")
|
||||
if err != nil {
|
||||
t.Fatalf("create: %v", err)
|
||||
}
|
||||
upd, err := c.UpdateKey(k.Key, "renamed", "", nil, "")
|
||||
if err != nil {
|
||||
t.Fatalf("update: %v", err)
|
||||
}
|
||||
if len(upd.Models) != 1 || upd.Models[0].Model != "gpt-4o" {
|
||||
t.Fatalf("scope must survive an update without models, got %#v", upd.Models)
|
||||
}
|
||||
cleared, err := c.UpdateKey(k.Key, "", "", []config.ModelScope{}, "")
|
||||
if err != nil {
|
||||
t.Fatalf("clear: %v", err)
|
||||
}
|
||||
if cleared.Models != nil {
|
||||
t.Fatalf("explicit empty list must clear the scope, got %#v", cleared.Models)
|
||||
}
|
||||
}
|
||||
|
||||
// TestSaveAutoRulesNormalizesModelCase: a rule whose model name differs only
|
||||
// in case must be normalized to the exact configured spelling so cooldown,
|
||||
// quota windows and the upstream model id all agree.
|
||||
func TestSaveAutoRulesNormalizesModelCase(t *testing.T) {
|
||||
c := newTestCore(t, newTestConfig(t))
|
||||
if err := c.SaveAutoRules([]config.ModelScope{{Model: "GPT-4O", Source: "s1", Tier: 1}}); err != nil {
|
||||
t.Fatalf("save: %v", err)
|
||||
}
|
||||
ch := c.AutoChain()
|
||||
if ch == nil || len(ch.Tiers) == 0 || len(ch.Tiers[0].Slots) == 0 {
|
||||
t.Fatalf("chain empty: %#v", ch)
|
||||
}
|
||||
if got := ch.Tiers[0].Slots[0].Model; got != "gpt-4o" {
|
||||
t.Fatalf("slot model must be normalized to the config spelling, got %q", got)
|
||||
}
|
||||
if got := ch.Tiers[0].Slots[0].Source; got != "s1" {
|
||||
t.Fatalf("slot source wrong: %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestConcurrentKeyMutation exercises the c.cfg.Keys lock under -race: key
|
||||
// creation/lookup/deletion from many goroutines at once.
|
||||
func TestConcurrentKeyMutation(t *testing.T) {
|
||||
c := newTestCore(t, newTestConfig(t))
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 16; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
k, err := c.CreateKey("k", "user", nil, "")
|
||||
if err == nil {
|
||||
if rec, ok := c.FindKey(k.Key); ok && rec.Key != k.Key {
|
||||
t.Errorf("wrong key returned")
|
||||
}
|
||||
c.UpdateKey(k.Key, "renamed", "", nil, "")
|
||||
c.DeleteKey(k.Key)
|
||||
}
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
Reference in New Issue
Block a user